Refactor benchmarking: model_type in benchmark_videos()

Author: maximpavliv

benchmarking/run_dlclive_benchmark.py

@@ -12,6 +12,7 @@
 import glob
 
 from dlclive import benchmark_videos, download_benchmarking_data
+from dlclive.engine import Engine
 
 datafolder = os.path.join(
     pathlib.Path(__file__).parent.absolute(), "Data-DLC-live-benchmark"
@@ -36,8 +37,22 @@
 if not os.path.isdir(out_dir):
     os.mkdir(out_dir)
 
-for m in dog_models:
-    benchmark_videos(m, dog_video, output=out_dir, n_frames=n_frames, pixels=pixels)
-
-for m in mouse_models:
-    benchmark_videos(m, mouse_video, output=out_dir, n_frames=n_frames, pixels=pixels)
+for model_path in dog_models:
+    benchmark_videos(
+        model_path=model_path,
+        model_type="base" if Engine.from_model_path(model_path) == Engine.TENSORFLOW else "pytorch",
+        video_path=dog_video,
+        output=out_dir,
+        n_frames=n_frames,
+        pixels=pixels
+    )
+
+for model_path in mouse_models:
+    benchmark_videos(
+        model_path=model_path,
+        model_type="base" if Engine.from_model_path(model_path) == Engine.TENSORFLOW else "pytorch",
+        video_path=mouse_video,
+        output=out_dir,
+        n_frames=n_frames,
+        pixels=pixels
+    )

dlclive/benchmark.py

@@ -66,6 +66,7 @@ def show_progress(count, block_size, total_size):
 
 def benchmark_videos(
     model_path,
+    model_type,
     video_path,
     output=None,
     n_frames=1000,
@@ -94,6 +95,9 @@ def benchmark_videos(
     ----------
     model_path : str
         path to exported DeepLabCut model
+    model_type: string, optional
+        Which model to use. For the PyTorch engine, options are [`pytorch`]. For the
+        TensorFlow engine, options are [`base`, `tensorrt`, `lite`].
     video_path : str or list
         path to video file or list of paths to video files
     output : str
@@ -169,7 +173,7 @@ def benchmark_videos(
 
             this_inf_times, this_im_size, meta = benchmark(
                 model_path=model_path,
-                model_type="base",
+                model_type=model_type,
                 video_path=video,
                 tf_config=tf_config,
                 resize=resize[i],
@@ -375,6 +379,7 @@ def benchmark(
         Path to the DeepLabCut model.
     model_type : str
         Which model to use. For the PyTorch engine, options are [`pytorch`]. For the
+        TensorFlow engine, options are [`base`, `tensorrt`, `lite`].
     video_path : str
         Path to the video file to be analyzed.
         TensorFlow engine, options are [`base`, `tensorrt`, `lite`].

dlclive/check_install/check_install.py

@@ -16,7 +16,8 @@
 
 from dlclibrary.dlcmodelzoo.modelzoo_download import download_huggingface_model
 
-from dlclive.benchmark_tf import benchmark_videos
+from dlclive.benchmark import benchmark_videos
+from dlclive.engine import Engine
 
 MODEL_NAME = "superanimal_quadruped"
 SNAPSHOT_NAME = "snapshot-700000.pb"
@@ -77,7 +78,12 @@ def main():
     # run benchmark videos
     print("\n Running inference...\n")
     benchmark_videos(
-        str(model_dir), video_file, display=display, resize=0.5, pcutoff=0.25
+        model_path=str(model_dir),
+        model_type="base" if Engine.from_model_path(model_dir) == Engine.TENSORFLOW else "pytorch",
+        video_path=video_file,
+        display=display,
+        resize=0.5,
+        pcutoff=0.25
     )
 
     # deleting temporary files

dlclive/engine.py

@@ -1,4 +1,5 @@
 from enum import Enum
+from pathlib import Path
 
 class Engine(Enum):
     TENSORFLOW = "tensorflow"
@@ -12,3 +13,21 @@ def from_model_type(cls, model_type: str) -> "Engine":
             return cls.TENSORFLOW
         else:
             raise ValueError(f"Unknown model type: {model_type}")
+
+    @classmethod
+    def from_model_path(cls, model_path: str | Path) -> "Engine":
+        path = Path(model_path)
+
+        if not path.exists():
+            raise FileNotFoundError(f"Model path does not exist: {model_path}")
+
+        if path.is_dir():
+            has_cfg = (path / "pose_cfg.yaml").is_file()
+            has_pb = any(p.suffix == ".pb" for p in path.glob("*.pb"))
+            if has_cfg and has_pb:
+                return cls.TENSORFLOW
+        elif path.is_file():
+            if path.suffix == ".pt":
+                return cls.PYTORCH
+
+        raise ValueError(f"Could not determine engine from model path: {model_path}")