improved docs

Author: n-poulsen

dlclive/benchmark.py

@@ -549,7 +549,6 @@ def benchmark_videos(
         )
 
     while True:
-
         ret, frame = cap.read()
         if not ret:
             break
@@ -656,8 +655,6 @@ def save_poses_to_files(video_path, save_dir, bodyparts, poses, timestamp):
             writer.writerow(row)
 
 
-
-
 import argparse
 import os
 

dlclive/benchmark_pytorch.py

@@ -214,7 +214,6 @@ def analyze_video(
         )
 
     while True:
-
         ret, frame = cap.read()
         if not ret:
             break

dlclive/benchmark_tf.py

@@ -305,7 +305,6 @@ def benchmark(
 
     iterator = range(n_frames) if (print_rate) or (display) else tqdm(range(n_frames))
     for i in iterator:
-
         ret, frame = cap.read()
 
         if not ret:
@@ -321,7 +320,6 @@ def benchmark(
         inf_times[i] = time.time() - start_pose
 
         if save_video:
-
             if colors is None:
                 all_colors = getattr(cc, cmap)
                 colors = [
@@ -399,15 +397,13 @@ def benchmark(
         vwriter.release()
 
     if save_poses:
-
         cfg_path = os.path.normpath(f"{model_path}/pose_cfg.yaml")
         ruamel_file = ruamel.yaml.YAML()
         dlc_cfg = ruamel_file.load(open(cfg_path, "r"))
         bodyparts = dlc_cfg["all_joints_names"]
         poses = np.array(poses)
 
         if use_pandas:
-
             poses = poses.reshape((poses.shape[0], poses.shape[1] * poses.shape[2]))
             pdindex = pd.MultiIndex.from_product(
                 [bodyparts, ["x", "y", "likelihood"]], names=["bodyparts", "coords"]
@@ -426,7 +422,6 @@ def benchmark(
             pose_df.to_hdf(out_dlc_file, key="df_with_missing", mode="w")
 
         else:
-
             out_vid_base = os.path.basename(video_path)
             out_dlc_file = os.path.normpath(
                 f"{out_dir}/{os.path.splitext(out_vid_base)[0]}_DLCLIVE_POSES.npy"
@@ -614,14 +609,12 @@ def benchmark_videos(
     # loop over videos
 
     for v in video_path:
-
         # initialize full inference times
 
         inf_times = []
         im_size_out = []
 
         for i in range(len(resize)):
-
             print(f"\nRun {i+1} / {len(resize)}\n")
 
             this_inf_times, this_im_size, TFGPUinference, meta = benchmark(

dlclive/core/inferenceutils.py

@@ -841,7 +841,6 @@ def assemble(self, chunk_size=1, n_processes=None):
         # work nicely with the GUI or interactive sessions.
         # In that case, we fall back to the serial assembly.
         if chunk_size == 0 or multiprocessing.get_start_method() == "spawn":
-
             for i, data_dict in enumerate(tqdm(self)):
                 assemblies, unique = self._assemble(data_dict, i)
                 if assemblies:

dlclive/dlclive.py

@@ -28,31 +28,41 @@ class DLCLive:
     -----------
 
     model_path: Path
-        Full path to exported model file
+        Full path to exported model (created when `deeplabcut.export_model(...)` was
+        called). For PyTorch models, this is a single model file. For TensorFlow models,
+        this is a directory containing the model snapshots.
 
     model_type: string, optional
-        which model to use: 'pytorch' or 'onnx' for exported snapshot
+        Which model to use. For the PyTorch engine, options are [`pytorch`]. For the
+        TensorFlow engine, options are [`base`, `tensorrt`, `lite`].
+
+    precision: string, optional
+        Precision of model weights, for model_type "pytorch" and "tensorrt". Options
+        are, for different model_types:
+            "pytorch": {"FP32", "FP16"}
+            "tensorrt": {"FP32", "FP16", "INT8"}
 
     tf_config:
+        TensorFlow only. Optional ConfigProto for the TensorFlow session.
 
+    single_animal: bool, default=True
+        PyTorch only.
 
-    precision: string, optional
-        precision of model weights, for model_type='onnx' or 'pytorch'. Can be 'FP32'
-        (default) or 'FP16'
+    device: str, optional, default=None
+        PyTorch only.
+
+    top_down_config: dict, optional, default=None
+
+    top_down_dynamic: dict, optional, default=None
 
     cropping: list of int
-        cropping parameters in pixel number: [x1, x2, y1, y2] #A: Maybe this is the
-        dynamic cropping of each frame to speed of processing, so instead of analyzing
-        the whole frame, it analyzes only the part of the frame where the animal is
-
-    dynamic: triple containing (state, detectiontreshold, margin) #A: margin adds some
-        space so the 'bbox' isn't too narrow around the animal'. First key points are
-        predicted, then dynamic cropping is performed to 'single out' the animal, and
-        then pose is estimated, we think.
+        Cropping parameters in pixel number: [x1, x2, y1, y2]
+
+    dynamic: triple containing (state, detectiontreshold, margin)
         If the state is true, then dynamic cropping will be performed. That means that
         if an object is detected (i.e. any body part > detectiontreshold), then object
         boundaries are computed according to the smallest/largest x position and
-        smallest/largest y position of all body parts. This  window is expanded by the
+        smallest/largest y position of all body parts. This window is expanded by the
         margin and from then on only the posture within this crop is analyzed (until the
         object is lost, i.e. <detectiontreshold). The current position is utilized for
         updating the crop window for the next frame (this is why the margin is important
@@ -63,8 +73,7 @@ class DLCLive:
         For example, resize=0.5 will downsize both the height and width of the image by
         a factor of 2.
 
-    processor: dlc pose processor object, optional #A: this is possibly the 'predictor'
-        - or is it what enables use on jetson boards?
+    processor: dlc pose processor object, optional
         User-defined processor object. Must contain two methods: process and save.
         The 'process' method takes in a pose, performs some processing, and returns
         processed pose.
@@ -80,12 +89,19 @@ class DLCLive:
         boolean flag to convert frames from BGR to RGB color scheme
 
     display: bool, optional
-        Display frames with DeepLabCut labels?
+        Open a display to show predicted pose in frames with DeepLabCut labels.
         This is useful for testing model accuracy and cropping parameters, but it is
         very slow.
 
+    pcutoff: float, default=0.5
+        Only used when display=True. The score threshold for displaying a bodypart in
+        the display.
+
+    display_radius: int, default=3
+        Only used when display=True. Radius for keypoint display in pixels, default=3
+
     display_cmap: str, optional
-        String indicating the Matplotlib colormap to use.
+        Only used when display=True. String indicating the Matplotlib colormap to use.
     """
 
     PARAMETERS = (
@@ -103,33 +119,36 @@ def __init__(
         self,
         model_path: str | Path,
         model_type: str = "base",
-        # tf_config: Any = None,
         precision: str = "FP32",
-        # single_animal: bool = True,
-        # device: str | None = None,
+        tf_config: Any = None,
+        single_animal: bool = True,
+        device: str | None = None,
+        top_down_config: dict | None = None,
+        top_down_dynamic: dict | None = None,
         cropping: list[int] | None = None,
         dynamic: tuple[bool, float, float] = (False, 0.5, 10),
         resize: float | None = None,
         convert2rgb: bool = True,
         processor: Processor | None = None,
         display: bool | Display = False,
         pcutoff: float = 0.5,
-        # bbox_cutoff: float = 0.6,
-        # max_detections: int = 1,
         display_radius: int = 3,
         display_cmap: str = "bmy",
-        **kwargs,
     ):
         self.path = Path(model_path)
         self.runner: BaseRunner = factory.build_runner(
             model_type,
             model_path,
-            **kwargs,
+            precision=precision,
+            tf_config=tf_config,
+            single_animal=single_animal,
+            device=device,
+            dynamic=top_down_dynamic,
+            top_down_config=top_down_config,
         )
         self.is_initialized = False
 
         self.model_type = model_type
-        self.precision = precision
         self.cropping = cropping
         self.dynamic = dynamic
         self.dynamic_cropping = None

dlclive/factory.py

@@ -1,33 +1,56 @@
-
 """Factory to build runners for DeepLabCut-Live inference"""
+from __future__ import annotations
+
 from pathlib import Path
+from typing import Literal
 
 from dlclive.core.runner import BaseRunner
 
 
 def build_runner(
-    model_type: str,
+    model_type: Literal["pytorch", "tensorflow", "base", "tensorrt", "lite"],
     model_path: str | Path,
     **kwargs,
 ) -> BaseRunner:
     """
 
     Parameters
     ----------
-    model_type
-    model_path
-    kwargs
+    model_type: str, optional
+        Which model to use. For the PyTorch engine, options are [`pytorch`]. For the
+        TensorFlow engine, options are [`base`, `tensorrt`, `lite`].
+    model_path: str, Path
+        Full path to exported model (created when `deeplabcut.export_model(...)` was
+        called). For PyTorch models, this is a single model file. For TensorFlow models,
+        this is a directory containing the model snapshots.
+
+    kwargs: dict, optional
+        PyTorch Engine Kwargs:
+
+        TensorFlow Engine Kwargs:
 
     Returns
     -------
 
     """
     if model_type.lower() == "pytorch":
         from dlclive.pose_estimation_pytorch.runner import PyTorchRunner
-        return PyTorchRunner(model_path, **kwargs)
+
+        valid = {"device", "precision", "single_animal", "dynamic", "top_down_config"}
+        return PyTorchRunner(model_path, **filter_keys(valid, kwargs))
 
     elif model_type.lower() in ("tensorflow", "base", "tensorrt", "lite"):
         from dlclive.pose_estimation_tensorflow.runner import TensorFlowRunner
-        return TensorFlowRunner(model_path, model_type, **kwargs)
+
+        if model_type.lower() == "tensorflow":
+            model_type = "base"
+
+        valid = {"tf_config", "precision"}
+        return TensorFlowRunner(model_path, model_type, **filter_keys(valid, kwargs))
 
     raise ValueError(f"Unknown model type: {model_type}")
+
+
+def filter_keys(valid: set[str], kwargs: dict) -> dict:
+    """Filters the keys in kwargs, only keeping those in valid."""
+    return {k: v for k, v in kwargs.items() if k in valid}

dlclive/graph.py

@@ -106,12 +106,13 @@ def get_output_tensors(graph):
 
 
 def get_input_tensor(graph):
-
     input_tensor = str(graph.get_operations()[0].name) + ":0"
     return input_tensor
 
 
-def extract_graph(graph, tf_config=None) -> tuple[tf.Session, tf.Tensor, list[tf.Tensor]]:
+def extract_graph(
+    graph, tf_config=None
+) -> tuple[tf.Session, tf.Tensor, list[tf.Tensor]]:
     """
     Initializes a tensorflow session with the specified graph and extracts the model's inputs and outputs
 

dlclive/live_inference.py

@@ -197,7 +197,6 @@ def analyze_live_video(
     ]
 
     if save_video:
-
         # Define output video path
         output_video_path = os.path.join(
             save_dir, f"{experiment_name}_DLCLIVE_LABELLED_{timestamp}.mp4"
@@ -217,7 +216,6 @@ def analyze_live_video(
         )
 
     while True:
-
         ret, frame = cap.read()
         if not ret:
             break

dlclive/pose_estimation_pytorch/data/image.py

@@ -109,7 +109,7 @@ def top_down_crop(
 
     # crop the pixels we care about
     image_crop = np.zeros((crop_h, crop_w, c), dtype=image.dtype)
-    image_crop[pad_top:pad_top + h, pad_left:pad_left + w] = image[y1:y2, x1:x2]
+    image_crop[pad_top : pad_top + h, pad_left : pad_left + w] = image[y1:y2, x1:x2]
 
     # resize the cropped image
     image = cv2.resize(image_crop, (out_w, out_h), interpolation=cv2.INTER_LINEAR)

dlclive/pose_estimation_pytorch/dynamic_cropping.py

@@ -39,9 +39,9 @@ class DynamicCropper:
             The margin used to expand an individuals bounding box before cropping it.
 
     Examples:
-        >>> import deeplabcut.pose_estimation_pytorch.models as models
+        >>> import torch.nn as nn
         >>>
-        >>> model: models.PoseModel
+        >>> model: nn.Module  # pose estimation model
         >>> frames: torch.Tensor  # shape (num_frames, 3, H, W)
         >>>
         >>> dynamic = DynamicCropper(threshold=0.6, margin=25)
@@ -57,6 +57,7 @@ class DynamicCropper:
         >>>     predictions.append(pose)
         >>>
     """
+
     threshold: float
     margin: int
     _crop: tuple[int, int, int, int] | None = field(default=None, repr=False)
@@ -424,16 +425,18 @@ def _prepare_bounding_box(
 
         input_ratio = w / h
         if input_ratio > self._td_ratio:  # h/w < h0/w0 => h' = w * h0/w0
-            h = w /  self._td_ratio
+            h = w / self._td_ratio
         elif input_ratio < self._td_ratio:  # w/h < w0/h0 => w' = h * w0/h0
-            w = h *  self._td_ratio
+            w = h * self._td_ratio
 
         x1, y1 = int(round(cx - (w / 2))), int(round(cy - (h / 2)))
         w, h = max(int(w), self.min_bbox_size[0]), max(int(h), self.min_bbox_size[1])
         return x1, y1, w, h
 
     def _crop_bounding_box(
-        self, image: torch.Tensor, bbox: tuple[int, int, int, int],
+        self,
+        image: torch.Tensor,
+        bbox: tuple[int, int, int, int],
     ) -> torch.Tensor:
         """Applies a top-down crop to an image given a bounding box.
 
@@ -487,7 +490,7 @@ def _extract_best_patch(self, pose: torch.Tensor) -> torch.Tensor:
         # set the crop to the one used for the best patch
         self._crop = self._patches[best_patch]
 
-        return pose[best_patch:best_patch + 1]
+        return pose[best_patch : best_patch + 1]
 
     def generate_patches(self) -> list[tuple[int, int, int, int]]:
         """Generates patch coordinates for splitting an image.