deeplabcut.pose_estimation_pytorch.runners.inference

Classes:

Name	Description
CTDInferenceRunner	Runner for pose estimation inference.
DetectorInferenceRunner	Runner for object detection inference.
InferenceRunner	Base class for inference runners.
PoseInferenceRunner	Runner for pose estimation inference.

Functions:

Name	Description
build_inference_runner	Build a runner object according to a pytorch configuration file.

CTDInferenceRunner

Bases: PoseInferenceRunner

Runner for pose estimation inference.

Parameters:

Name	Type	Description	Default
model	PoseModel	The CTD model to run inference with.	required
bu_runner	PoseInferenceRunner | None	A runner for the BU model to run inference with. If no BU runner is given, conditions must be given in the context for the data. Otherwise an error will be raised during inference.	None
ctd_tracking	bool | CTDTrackingConfig	Whether to track using the CTD model. If	False

Methods:

Name	Description
inference	Run CTD model inference on the given dataset.
predict	Makes predictions from a model input and output.

Source code in deeplabcut/pose_estimation_pytorch/runners/inference.py

class CTDInferenceRunner(PoseInferenceRunner):
    """Runner for pose estimation inference.

    Args:
        model: The CTD model to run inference with.
        bu_runner: A runner for the BU model to run inference with. If no BU runner is
            given, conditions must be given in the context for the data. Otherwise an
            error will be raised during inference.
        ctd_tracking: Whether to track using the CTD model. If
    """

    def __init__(
        self,
        model: PoseModel,
        bu_runner: PoseInferenceRunner | None = None,
        ctd_tracking: bool | ctd.CTDTrackingConfig = False,
        **kwargs,
    ):
        super().__init__(model, **kwargs)
        self.bu_runner = bu_runner
        if bu_runner is not None:
            self.bu_runner.model.eval()

        self.tracking = None
        if isinstance(ctd_tracking, ctd.CTDTrackingConfig):
            self.tracking = ctd_tracking
        elif ctd_tracking:  # generate default config
            self.tracking = ctd.CTDTrackingConfig()

        if self.tracking and self.batch_size != 1:
            print("CTD tracking can only be used with batch size 1. Updating it.")
            self.batch_size = 1

        self._image_loader = LoadImage()

        # Stored poses and IDX -> ID map for CTD tracking
        self._bu_age = -1
        self._missing_idvs = False
        self._prev_pose = None
        self._idx_to_id = None
        self._ctd_track_ages = None  # the age of each CTD tracklet

    @_inference_mode_decorator
    def inference(
        self,
        images: (Iterable[str | Path | np.ndarray] | Iterable[tuple[str | Path | np.ndarray, dict[str, Any]]]),
        shelf_writer: shelving.ShelfWriter | None = None,
    ) -> list[dict[str, np.ndarray]]:
        """Run CTD model inference on the given dataset.

        Args:
            images: the images to run inference on, optionally with context
            shelf_writer: by default, data are saved in a list and returned at the end
                of inference. Passing a shelf manager writes data to disk on-the-fly
                using a "shelf" (a pickle-based, persistent, database-like object by
                default, resulting in constant memory footprint). The returned list is
                then empty.

        Returns:
            a dict containing head predictions for each image
            [
                {
                    "bodypart": {"poses": np.array},
                    "unique_bodypart": {"poses": np.array},
                }
            ]
        """
        if self.tracking:
            return self._ctd_tracking_inference(images, shelf_writer)

        results = []
        for data in images:
            data = self.add_conditions(data)
            self._prepare_inputs(data)
            self._process_full_batches()
            results += self._extract_results(shelf_writer)

        # Process the last batch even if not full
        if self._inputs_waiting_for_processing():
            self._process_batch()
            results += self._extract_results(shelf_writer)

        return results

    def predict(self, inputs: torch.Tensor, **kwargs) -> list[dict[str, dict[str, np.ndarray]]]:
        """Makes predictions from a model input and output.

        Args:
            the inputs to the model, of shape (batch_size, ...)

        Returns:
            predictions for each of the 'batch_size' inputs, made by each head, e.g.
            [
                {
                    "bodypart": {"poses": np.ndarray},
                    "unique_bodypart": {"poses": np.ndarray},
                }
            ]
        """
        cond_kpts = kwargs.get("cond_kpts", None)
        if cond_kpts is not None and cond_kpts.shape[0] == 0:
            # No conditions, so just return an empty prediction list
            return []

        # Normal prediction path
        if self.inference_cfg.autocast.enabled:
            with torch.autocast(device_type=str(self.device)):
                outputs = self.model(inputs.to(self.device), **kwargs)
                raw_predictions = self.model.get_predictions(outputs)
        else:
            outputs = self.model(inputs.to(self.device), **kwargs)
            raw_predictions = self.model.get_predictions(outputs)

        predictions = [
            {
                head: {pred_name: pred[b].cpu().numpy() for pred_name, pred in head_outputs.items()}
                for head, head_outputs in raw_predictions.items()
            }
            for b in range(len(inputs))
        ]

        return predictions

    def add_conditions(
        self,
        data: str | Path | np.ndarray | tuple[str | Path | np.ndarray, dict],
    ) -> tuple[np.ndarray, dict]:
        if isinstance(data, (str, Path, np.ndarray)):
            inputs, context = data, {}
        else:
            inputs, context = data

        # Load the image once - then given as a numpy array to CTD
        image, _ = self._image_loader(inputs, context)

        # If the conditional keypoints are in the context, return the context
        if "cond_kpts" in context:
            return image, context

        # Run the pre-processor
        if self.bu_runner.preprocessor is not None:
            inputs, context = self.bu_runner.preprocessor(image, context)
        else:
            inputs = torch.as_tensor(image)

        # Get and post-process the predictions
        with _directml_runtime_error_hint():
            predictions = self.bu_runner.predict(inputs)
        if self.bu_runner.postprocessor is not None:
            predictions, context = self.bu_runner.postprocessor(predictions, context)

        # Extract the conditions
        conds = predictions["bodyparts"][..., :3]
        pred_mask = ~np.all(np.any((conds <= 0) | np.isnan(conds), axis=2), axis=1)
        if np.sum(pred_mask) > 0:
            conds = conds[pred_mask]
        else:
            conds = np.zeros((0, conds.shape[1], 3))

        return image, {"cond_kpts": conds}

    def _ctd_tracking_inference(
        self,
        images: (Iterable[str | Path | np.ndarray] | Iterable[tuple[str | Path | np.ndarray, dict[str, Any]]]),
        shelf_writer: shelving.ShelfWriter | None = None,
    ) -> list[dict[str, np.ndarray]]:
        results = []
        for data in images:
            inputs, context = self._prepare_ctd_inputs(data)
            model_kwargs = context.pop("model_kwargs", {})
            with _directml_runtime_error_hint():
                predictions = self.predict(inputs, **model_kwargs)
            if self.postprocessor is not None:
                # Pop the "cond_kpts" from the context so there's no re-scoring
                # This is required when tracking with CTD, otherwise scores go to 0
                if self._prev_pose is not None:
                    context.pop("cond_kpts")

                predictions, _ = self.postprocessor(predictions, context)

            # Set the predictions as context for the next frame
            self._ctd_tracking_postprocess(predictions, context["image_size"])

            if shelf_writer is not None:
                shelf_writer.add_prediction(
                    bodyparts=predictions["bodyparts"],
                    unique_bodyparts=predictions.get("unique_bodyparts"),
                    identity_scores=predictions.get("identity_scores"),
                    features=predictions.get("features"),
                )
            else:
                results.append(predictions)

        return results

    def _prepare_ctd_inputs(self, data) -> tuple[torch.Tensor, dict[str, Any]]:
        # If there's no valid poses, use the BU model to get conditions
        self._bu_age += 1
        if (
            self._prev_pose is None
            or (self._missing_idvs and self.tracking.bu_on_lost_idv and self._bu_age >= self.tracking.bu_max_frequency)
            or (self.tracking.bu_min_frequency is not None and self._bu_age >= self.tracking.bu_min_frequency)
        ):
            self._bu_age = 0
            inputs, context = self.add_conditions(data)

            if self._prev_pose is not None:
                context["cond_kpts"] = self._merge_conditions(context["cond_kpts"])

        else:
            if isinstance(data, (str, Path, np.ndarray)):
                inputs, context = data, {}
            else:
                inputs, context = data

            context["cond_kpts"] = self._prev_pose

        if self.preprocessor is None:
            return torch.as_tensor(inputs), context

        inputs, context = self.preprocessor(inputs, context)
        return inputs, context

    def _ctd_tracking_postprocess(
        self,
        predictions: dict[str, np.ndarray],
        image_size: tuple[int, int],
    ) -> None:
        """Post-processes predictions.

        In-place changes to the predictions dict.
        """
        # reorder the previous poses so the indices match the track IDs
        if self._idx_to_id is not None:
            predictions["bodyparts"] = predictions["bodyparts"][self._idx_to_id]

        # mask all keypoints below the CTD tracking threshold
        prev_pose = predictions["bodyparts"][..., :3].copy()
        prev_pose[prev_pose[..., 2] <= self.tracking.threshold_ctd] = np.nan

        # mask all keypoints outside the image
        w, h = image_size
        prev_pose[prev_pose[..., 0] < 0] = np.nan
        prev_pose[prev_pose[..., 1] < 0] = np.nan
        prev_pose[prev_pose[..., 0] >= w] = np.nan
        prev_pose[prev_pose[..., 1] >= h] = np.nan

        # apply NMS on the conditions, keeping older tracks
        order = None
        if self._ctd_track_ages is not None:
            ordering = self._ctd_track_ages.copy()

            # sort by track age, then score
            vis = np.sum(np.all(~np.isnan(prev_pose), axis=-1), axis=-1) > 1
            scores = np.nanmean(prev_pose[vis, :, 2], axis=-1)
            ordering[vis] += scores

            # only keep non-zero scores
            order = ordering.argsort()[::-1]
            order = order[ordering[order] > 0]

        nms_mask = nms.nms_oks(
            prev_pose,
            oks_threshold=self.tracking.threshold_nms,
            oks_sigmas=0.1,
            oks_margin=1.0,
            score_threshold=self.tracking.threshold_ctd,
            order=order,
        )

        # Set the previous pose and ID ordering
        if np.any(nms_mask):
            self._prev_pose = prev_pose[nms_mask]

            # get the IDs of the kept poses
            found_idx_to_id = np.where(nms_mask)[0]
            missing_ids = np.where(~nms_mask)[0]
            self._idx_to_id = np.concatenate([found_idx_to_id, missing_ids])

            # add 1 to the age of kept tracks
            if self._ctd_track_ages is None:
                self._ctd_track_ages = np.zeros(len(self._idx_to_id))
            self._ctd_track_ages[nms_mask] += 1
            self._ctd_track_ages[~nms_mask] = 0

            # check if there are any missing individuals
            self._missing_idvs = len(self._prev_pose) != len(self._idx_to_id)
        else:
            self._prev_pose = None
            self._idx_to_id = None
            self._ctd_track_ages = None

    def _merge_conditions(self, bu_cond: np.ndarray) -> np.ndarray:
        """Merges conditions made by a BU model with existing conditions from CTD
        tracking.
        """
        # prepare the BU conditions for matching
        bu_cond = bu_cond.copy()[:, :, :3]
        # mask low-quality keypoints
        bu_cond[bu_cond[..., 2] < self.tracking.threshold_ctd] = np.nan

        # remove non-visible individuals
        kpt_vis = np.all(~np.isnan(bu_cond), axis=-1)
        idv_vis = np.sum(kpt_vis, axis=-1) > 1  # need at least 2 kpts for OKS

        # if no valid BU predictions are left, return the CTD conditions
        if np.sum(idv_vis) == 0:
            return self._prev_pose

        # match BU conditions to CTD poses from the highest score to the lowest
        bu_cond = bu_cond[idv_vis]
        new_conditions = []
        for bu_pose in bu_cond:
            best_oks = 0
            for ctd_pose in self._prev_pose:
                best_oks = max(
                    best_oks,
                    calc_object_keypoint_similarity(bu_pose, ctd_pose, sigma=0.1),
                )

            if best_oks < self.tracking.threshold_bu_add:
                new_conditions.append((best_oks, bu_pose))

        # add the conditions with the lowest OKS score
        new_conditions = [c[1] for c in sorted(new_conditions, key=lambda x: x[0])]

        # if there are no new conditions,
        if len(new_conditions) == 0:
            return self._prev_pose

        new_conditions = np.stack(new_conditions, axis=0)
        cond_pose = np.concatenate([self._prev_pose, new_conditions], axis=0)
        return cond_pose[: len(self._idx_to_id)]

inference

inference(
    images: Iterable[str | Path | ndarray] | Iterable[tuple[str | Path | ndarray, dict[str, Any]]],
    shelf_writer: ShelfWriter | None = None,
) -> list[dict[str, np.ndarray]]

Run CTD model inference on the given dataset.

Parameters:

Name	Type	Description	Default
images	Iterable[str | Path | ndarray] | Iterable[tuple[str | Path | ndarray, dict[str, Any]]]	the images to run inference on, optionally with context	required
shelf_writer	ShelfWriter | None	by default, data are saved in a list and returned at the end of inference. Passing a shelf manager writes data to disk on-the-fly using a "shelf" (a pickle-based, persistent, database-like object by default, resulting in constant memory footprint). The returned list is then empty.	None

Returns:

Type	Description
list[dict[str, ndarray]]	a dict containing head predictions for each image [ { "bodypart": {"poses": np.array}, "unique_bodypart": {"poses": np.array}, } ]

Source code in deeplabcut/pose_estimation_pytorch/runners/inference.py

@_inference_mode_decorator
def inference(
    self,
    images: (Iterable[str | Path | np.ndarray] | Iterable[tuple[str | Path | np.ndarray, dict[str, Any]]]),
    shelf_writer: shelving.ShelfWriter | None = None,
) -> list[dict[str, np.ndarray]]:
    """Run CTD model inference on the given dataset.

    Args:
        images: the images to run inference on, optionally with context
        shelf_writer: by default, data are saved in a list and returned at the end
            of inference. Passing a shelf manager writes data to disk on-the-fly
            using a "shelf" (a pickle-based, persistent, database-like object by
            default, resulting in constant memory footprint). The returned list is
            then empty.

    Returns:
        a dict containing head predictions for each image
        [
            {
                "bodypart": {"poses": np.array},
                "unique_bodypart": {"poses": np.array},
            }
        ]
    """
    if self.tracking:
        return self._ctd_tracking_inference(images, shelf_writer)

    results = []
    for data in images:
        data = self.add_conditions(data)
        self._prepare_inputs(data)
        self._process_full_batches()
        results += self._extract_results(shelf_writer)

    # Process the last batch even if not full
    if self._inputs_waiting_for_processing():
        self._process_batch()
        results += self._extract_results(shelf_writer)

    return results

predict

predict(inputs: Tensor, **kwargs) -> list[dict[str, dict[str, np.ndarray]]]

Makes predictions from a model input and output.

Returns:

Type	Description
list[dict[str, dict[str, ndarray]]]	predictions for each of the 'batch_size' inputs, made by each head, e.g. [ { "bodypart": {"poses": np.ndarray}, "unique_bodypart": {"poses": np.ndarray}, } ]

Source code in deeplabcut/pose_estimation_pytorch/runners/inference.py

def predict(self, inputs: torch.Tensor, **kwargs) -> list[dict[str, dict[str, np.ndarray]]]:
    """Makes predictions from a model input and output.

    Args:
        the inputs to the model, of shape (batch_size, ...)

    Returns:
        predictions for each of the 'batch_size' inputs, made by each head, e.g.
        [
            {
                "bodypart": {"poses": np.ndarray},
                "unique_bodypart": {"poses": np.ndarray},
            }
        ]
    """
    cond_kpts = kwargs.get("cond_kpts", None)
    if cond_kpts is not None and cond_kpts.shape[0] == 0:
        # No conditions, so just return an empty prediction list
        return []

    # Normal prediction path
    if self.inference_cfg.autocast.enabled:
        with torch.autocast(device_type=str(self.device)):
            outputs = self.model(inputs.to(self.device), **kwargs)
            raw_predictions = self.model.get_predictions(outputs)
    else:
        outputs = self.model(inputs.to(self.device), **kwargs)
        raw_predictions = self.model.get_predictions(outputs)

    predictions = [
        {
            head: {pred_name: pred[b].cpu().numpy() for pred_name, pred in head_outputs.items()}
            for head, head_outputs in raw_predictions.items()
        }
        for b in range(len(inputs))
    ]

    return predictions

DetectorInferenceRunner

Bases: InferenceRunner[BaseDetector]

Runner for object detection inference.

Methods:

Name	Description
__init__	Args:
predict	Makes predictions from a model input and output.

Source code in deeplabcut/pose_estimation_pytorch/runners/inference.py

class DetectorInferenceRunner(InferenceRunner[BaseDetector]):
    """Runner for object detection inference."""

    def __init__(self, model: BaseDetector, **kwargs):
        """
        Args:
            model: The detector to use for inference.
            **kwargs: Inference runner kwargs.
        """
        super().__init__(model, **kwargs)

    def predict(self, inputs: torch.Tensor, **kwargs) -> list[dict[str, dict[str, np.ndarray]]]:
        """Makes predictions from a model input and output.

        Args:
            the inputs to the model, of shape (batch_size, ...)

        Returns:
            predictions for each of the 'batch_size' inputs, made by each head, e.g.
            [
                {
                    "bodypart": {"poses": np.ndarray},
                    "unique_bodypart": "poses": np.ndarray},
                }
            ]
        """
        if self.inference_cfg.autocast.enabled:
            with torch.autocast(device_type=str(self.device)):
                _, raw_predictions = self.model(inputs.to(self.device))
        else:
            _, raw_predictions = self.model(inputs.to(self.device))
        predictions = [
            {
                "detection": {
                    "bboxes": item["boxes"].cpu().numpy().reshape(-1, 4),
                    "scores": item["scores"].cpu().numpy().reshape(-1),
                }
            }
            for item in raw_predictions
        ]
        return predictions

__init__

__init__(model: BaseDetector, **kwargs)

Parameters:

Name	Type	Description	Default
model	BaseDetector	The detector to use for inference.	required
**kwargs		Inference runner kwargs.	{}

Source code in deeplabcut/pose_estimation_pytorch/runners/inference.py

def __init__(self, model: BaseDetector, **kwargs):
    """
    Args:
        model: The detector to use for inference.
        **kwargs: Inference runner kwargs.
    """
    super().__init__(model, **kwargs)

predict

predict(inputs: Tensor, **kwargs) -> list[dict[str, dict[str, np.ndarray]]]

Makes predictions from a model input and output.

Returns:

Type	Description
list[dict[str, dict[str, ndarray]]]	predictions for each of the 'batch_size' inputs, made by each head, e.g. [ { "bodypart": {"poses": np.ndarray}, "unique_bodypart": "poses": np.ndarray}, } ]

Source code in deeplabcut/pose_estimation_pytorch/runners/inference.py

def predict(self, inputs: torch.Tensor, **kwargs) -> list[dict[str, dict[str, np.ndarray]]]:
    """Makes predictions from a model input and output.

    Args:
        the inputs to the model, of shape (batch_size, ...)

    Returns:
        predictions for each of the 'batch_size' inputs, made by each head, e.g.
        [
            {
                "bodypart": {"poses": np.ndarray},
                "unique_bodypart": "poses": np.ndarray},
            }
        ]
    """
    if self.inference_cfg.autocast.enabled:
        with torch.autocast(device_type=str(self.device)):
            _, raw_predictions = self.model(inputs.to(self.device))
    else:
        _, raw_predictions = self.model(inputs.to(self.device))
    predictions = [
        {
            "detection": {
                "bboxes": item["boxes"].cpu().numpy().reshape(-1, 4),
                "scores": item["scores"].cpu().numpy().reshape(-1),
            }
        }
        for item in raw_predictions
    ]
    return predictions

InferenceRunner

Bases: Runner, Generic[ModelType]

Base class for inference runners.

A runner takes a model and runs actions on it, such as training or inference

Methods:

Name	Description
__del__	Cleanup method to ensure threads are stopped.
__init__	Args:
inference	Run model inference on the given dataset.
predict	Makes predictions from a model input and output.

Source code in deeplabcut/pose_estimation_pytorch/runners/inference.py

class InferenceRunner(Runner, Generic[ModelType], metaclass=ABCMeta):
    """Base class for inference runners.

    A runner takes a model and runs actions on it, such as training or inference
    """

    def __init__(
        self,
        model: ModelType,
        batch_size: int = 1,
        device: str = "cpu",
        snapshot_path: str | Path | None = None,
        preprocessor: Preprocessor | None = None,
        postprocessor: Postprocessor | None = None,
        load_weights_only: bool | None = None,
        inference_cfg: InferenceConfig | dict | None = None,
    ):
        """
        Args:
            model: The model to run actions on
            device: The device to use (e.g. {'cpu', 'cuda:0', 'mps'})
            snapshot_path: If defined, the path of a snapshot from which to load
                pretrained weights
            preprocessor: The preprocessor to use on images before inference
            postprocessor: The postprocessor to use on images after inference
            load_weights_only: Value for the torch.load() `weights_only` parameter.
                If False, the python pickle module is used implicitly, which is known to
                    be insecure. Only set to False if you're loading data that you trust
                    (e.g. snapshots that you created). For more information, see:
                        https://pytorch.org/docs/stable/generated/torch.load.html
                If None, the default value is used:
                    `deeplabcut.pose_estimation_pytorch.get_load_weights_only()`
            inference_cfg: Configuration for the inference runner
        """
        super().__init__(model=model, device=device, snapshot_path=snapshot_path)
        if not isinstance(batch_size, int) or batch_size <= 0:
            raise ValueError(f"batch_size must be a positive integer; is {batch_size}")

        self.batch_size = batch_size
        self.preprocessor = preprocessor
        self.postprocessor = postprocessor

        if isinstance(inference_cfg, InferenceConfig):
            self.inference_cfg = inference_cfg
        elif isinstance(inference_cfg, dict):
            self.inference_cfg = InferenceConfig.from_dict(inference_cfg)
        elif inference_cfg is None:
            self.inference_cfg = InferenceConfig()
        else:
            raise ValueError(f"Invalid inference config: {inference_cfg}")

        if self.snapshot_path is not None and self.snapshot_path != "":
            self.load_snapshot(
                self.snapshot_path,
                self.device,
                self.model,
                weights_only=load_weights_only,
            )

        self.model.to(self.device)
        self.model.eval()

        if self.inference_cfg.compile.enabled:
            try:
                self.model = torch.compile(self.model, backend=self.inference_cfg.compile.backend)
            except Exception as e:
                warnings.warn(
                    f"torch.compile failed with backend='{self.inference_cfg.compile.backend}', "
                    f"falling back to eager mode. Error: {e}",
                    stacklevel=2,
                )

        self._batch_list: list[torch.Tensor] = []
        self._model_kwargs: dict[str, np.ndarray | torch.Tensor] = {}

        self._contexts: list[dict] = []
        self._image_batch_sizes: list[int] = []
        self._predictions: list = []

        # Async-specific attributes
        if self.inference_cfg.multithreading.enabled:
            self._input_queue = Queue(maxsize=self.inference_cfg.multithreading.queue_length)
            self._preprocessing_thread = None
            self._stop_event = threading.Event()
            self._exception = None

    @abstractmethod
    def predict(self, inputs: torch.Tensor, **kwargs) -> list[dict[str, dict[str, np.ndarray]]]:
        """Makes predictions from a model input and output.

        Args:
            the inputs to the model, of shape (batch_size, ...)

        Returns:
            the predictions for each of the 'batch_size' inputs
        """

    @_inference_mode_decorator
    def inference(
        self,
        images: (Iterable[str | Path | np.ndarray] | Iterable[tuple[str | Path | np.ndarray, dict[str, Any]]]),
        shelf_writer: shelving.ShelfWriter | None = None,
    ) -> list[dict[str, np.ndarray]]:
        """Run model inference on the given dataset.

        TODO: Add an option to also return head outputs (such as heatmaps)? Can be
         super useful for debugging

        Args:
            images: the images to run inference on, optionally with context
            shelf_writer: by default, data are saved in a list and returned at the end
                of inference. Passing a shelf manager writes data to disk on-the-fly
                using a "shelf" (a pickle-based, persistent, database-like object by
                default, resulting in constant memory footprint). The returned list is
                then empty.

        Returns:
            a dict containing head predictions for each image
            [
                {
                    "bodypart": {"poses": np.array},
                    "unique_bodypart": {"poses": np.array},
                }
            ]
        """
        if self.inference_cfg.multithreading.enabled:
            return self._async_inference(images, shelf_writer)
        else:
            return self._sequential_inference(images, shelf_writer)

    def _sequential_inference(
        self,
        images: (Iterable[str | Path | np.ndarray] | Iterable[tuple[str | Path | np.ndarray, dict[str, Any]]]),
        shelf_writer: shelving.ShelfWriter | None = None,
    ) -> list[dict[str, np.ndarray]]:
        """Original sequential inference implementation."""
        results = []
        for data in images:
            self._prepare_inputs(data)
            self._process_full_batches()
            results += self._extract_results(shelf_writer)

        # Process the last batch even if not full
        if self._inputs_waiting_for_processing():
            self._process_batch()
            results += self._extract_results(shelf_writer)

        return results

    def _async_inference(
        self,
        images: (Iterable[str | Path | np.ndarray] | Iterable[tuple[str | Path | np.ndarray, dict[str, Any]]]),
        shelf_writer: shelving.ShelfWriter | None = None,
    ) -> list[dict[str, np.ndarray]]:
        """Async inference with pipeline parallelism."""
        # Reset state
        self._stop_event.clear()
        self._exception = None
        self._batch_list = []
        self._model_kwargs = {}
        self._contexts = []
        self._image_batch_sizes = []
        self._predictions = []

        # Start preprocessing thread
        self._preprocessing_thread = threading.Thread(target=self._preprocessing_worker, args=(images,))
        self._preprocessing_thread.start()

        results = []

        try:
            while True:
                # Get next batch from queue
                item = self._safe_get()

                # None means either producer finished or stop_event triggered
                if item is None:
                    break

                batch, model_kwargs = item

                # Run model inference
                with _directml_runtime_error_hint():
                    predictions = self.predict(batch, **model_kwargs)
                self._predictions.extend(predictions)

                # Extract and return results
                batch_results = self._extract_results(shelf_writer)
                results.extend(batch_results)

                # propagate any exception from the producer immediately
                if self._exception is not None:
                    raise self._exception

        except BaseException as e:  # catches KeyboardInterrupt, SystemExit, etc.
            # tell producer to quit
            self._stop_event.set()
            raise e
        finally:
            # Wait for preprocessing thread to finish
            if self._preprocessing_thread is not None:
                self._preprocessing_thread.join(timeout=self.inference_cfg.multithreading.timeout)

            # Check for exceptions in preprocessing thread
            if self._exception is not None:
                raise self._exception

        return results

    def _prepare_inputs(
        self,
        data: str | Path | np.ndarray | tuple[str | Path | np.ndarray, dict],
    ) -> None:
        """Prepares inputs for an image and adds them to the data ready to be
        processed.
        """
        if isinstance(data, (str, Path, np.ndarray)):
            inputs, context = data, {}
        else:
            inputs, context = data

        if self.preprocessor is not None:
            inputs, context = self.preprocessor(inputs, context)
        else:
            inputs = torch.as_tensor(inputs)

        # add new model_kwargs from the inputs
        model_kwargs = context.pop("model_kwargs", {})
        for k, v in model_kwargs.items():
            curr_v = self._model_kwargs.get(k)
            if curr_v is None or len(curr_v) == 0:
                curr_v = v
            elif len(v) == 0:
                continue
            elif isinstance(curr_v, np.ndarray):
                curr_v = np.concatenate([curr_v, v], axis=0)
            elif isinstance(curr_v, torch.Tensor):
                curr_v = torch.cat([curr_v, v], dim=0)
            else:
                raise ValueError(f"model_kwargs {k} must be a numpy array or torch tensor - found '{type(v)}'.")
            self._model_kwargs[k] = curr_v

        self._contexts.append(context)
        self._image_batch_sizes.append(len(inputs))

        # skip when there are no inputs for an image
        if len(inputs) == 0:
            return

        # extend the list with individual image tensors (slice along first dim)
        self._batch_list.extend(list(inputs))

    def _process_full_batches(self) -> None:
        """Processes prepared inputs in batches of the desired batch size."""
        while len(self._batch_list) >= self.batch_size:
            self._process_batch()

    def _extract_results(self, shelf_writer: shelving.ShelfWriter) -> list:
        """Obtains results that were obtained from processing a batch."""
        results = []
        while len(self._image_batch_sizes) > 0 and len(self._predictions) >= self._image_batch_sizes[0]:
            num_predictions = self._image_batch_sizes[0]
            image_predictions = self._predictions[:num_predictions]
            context = self._contexts[0]
            if self.postprocessor is not None:
                # TODO: Should we return context?
                # TODO: typing update - the post-processor can remove a dict level
                image_predictions, _ = self.postprocessor(image_predictions, context)

            if shelf_writer is not None:
                shelf_writer.add_prediction(
                    bodyparts=image_predictions["bodyparts"],
                    unique_bodyparts=image_predictions.get("unique_bodyparts"),
                    identity_scores=image_predictions.get("identity_scores"),
                    features=image_predictions.get("features"),
                )
            else:
                results.append(image_predictions)

            self._contexts = self._contexts[1:]
            self._image_batch_sizes = self._image_batch_sizes[1:]
            self._predictions = self._predictions[num_predictions:]

        return results

    def _process_batch(self) -> None:
        """Processes a batch.

        There must be inputs waiting to be processed before this is called, otherwise
        this method will raise an error.
        """
        batch = torch.stack(self._batch_list[: self.batch_size], dim=0)
        model_kwargs = {mk: v[: self.batch_size] for mk, v in self._model_kwargs.items()}

        with _directml_runtime_error_hint():
            self._predictions += self.predict(batch, **model_kwargs)

        # remove processed inputs
        if len(self._batch_list) <= self.batch_size:
            self._batch_list = []
            self._model_kwargs = {}
        else:
            self._batch_list = self._batch_list[self.batch_size :]
            self._model_kwargs = {mk: v[self.batch_size :] for mk, v in self._model_kwargs.items()}

    def _inputs_waiting_for_processing(self) -> bool:
        """Returns: Whether there are inputs which have not yet been processed"""
        return len(self._batch_list) > 0

    def _safe_put(self, item: Any) -> bool:
        """Put item in the queue, retrying until successful or stop_event is set."""
        while not self._stop_event.is_set():
            try:
                self._input_queue.put(item, timeout=1.0)
                return True
            except Full:
                continue
        return False

    def _safe_get(self) -> Any:
        """Get the next item from the queue safely, retrying until successful or
        stop_event is set.

        Returns:
            The item from the queue, or None if the producer is dead or stop_signal is raised and queue empty.
        """
        while True:
            try:
                item = self._input_queue.get(timeout=1.0)
                return item
            except Empty:
                # check if producer is still running
                if (
                    self._stop_event.is_set()
                    or self._preprocessing_thread is None
                    or not self._preprocessing_thread.is_alive()
                ):
                    return None
                continue

    def _preprocessing_worker(self, images: Iterable) -> None:
        """Background worker that prepares inputs and puts them in the input queue."""
        try:
            for data in images:
                if self._stop_event.is_set():
                    break

                # Prepare inputs using the parent class method
                self._prepare_inputs(data)

                # Process full batches and put them in the queue
                while len(self._batch_list) >= self.batch_size:
                    batch = torch.stack(self._batch_list[: self.batch_size], dim=0)
                    model_kwargs = {mk: v[: self.batch_size] for mk, v in self._model_kwargs.items()}

                    self._safe_put((batch, model_kwargs))

                    # Remove processed inputs from batch
                    if len(self._batch_list) <= self.batch_size:
                        self._batch_list, self._model_kwargs = [], {}
                    else:
                        self._batch_list = self._batch_list[self.batch_size :]
                        self._model_kwargs = {mk: v[self.batch_size :] for mk, v in self._model_kwargs.items()}

            # Process any remaining inputs
            if len(self._batch_list) > 0:
                batch = torch.stack(self._batch_list, dim=0)
                self._safe_put((batch, self._model_kwargs))

        except BaseException as e:  # catches KeyboardInterrupt, SystemExit, etc.
            self._exception = e
            self._stop_event.set()
        finally:
            # Signal that preprocessing is done
            self._safe_put(None)

    def __del__(self):
        """Cleanup method to ensure threads are stopped."""
        if hasattr(self, "_stop_event"):
            self._stop_event.set()
        if hasattr(self, "_preprocessing_thread") and self._preprocessing_thread is not None:
            self._preprocessing_thread.join(timeout=1.0)

__del__

__del__()

Cleanup method to ensure threads are stopped.

Source code in deeplabcut/pose_estimation_pytorch/runners/inference.py

def __del__(self):
    """Cleanup method to ensure threads are stopped."""
    if hasattr(self, "_stop_event"):
        self._stop_event.set()
    if hasattr(self, "_preprocessing_thread") and self._preprocessing_thread is not None:
        self._preprocessing_thread.join(timeout=1.0)

__init__

__init__(
    model: ModelType,
    batch_size: int = 1,
    device: str = "cpu",
    snapshot_path: str | Path | None = None,
    preprocessor: Preprocessor | None = None,
    postprocessor: Postprocessor | None = None,
    load_weights_only: bool | None = None,
    inference_cfg: InferenceConfig | dict | None = None,
)

Parameters:

Name	Type	Description	Default
model	ModelType	The model to run actions on	required
device	str	The device to use (e.g. {'cpu', 'cuda:0', 'mps'})	'cpu'
snapshot_path	str | Path | None	If defined, the path of a snapshot from which to load pretrained weights	None
preprocessor	Preprocessor | None	The preprocessor to use on images before inference	None
postprocessor	Postprocessor | None	The postprocessor to use on images after inference	None
load_weights_only	bool | None	Value for the torch.load() weights_only parameter. If False, the python pickle module is used implicitly, which is known to be insecure. Only set to False if you're loading data that you trust (e.g. snapshots that you created). For more information, see: https://pytorch.org/docs/stable/generated/torch.load.html If None, the default value is used: deeplabcut.pose_estimation_pytorch.get_load_weights_only()	None
inference_cfg	InferenceConfig | dict | None	Configuration for the inference runner	None

Source code in deeplabcut/pose_estimation_pytorch/runners/inference.py

def __init__(
    self,
    model: ModelType,
    batch_size: int = 1,
    device: str = "cpu",
    snapshot_path: str | Path | None = None,
    preprocessor: Preprocessor | None = None,
    postprocessor: Postprocessor | None = None,
    load_weights_only: bool | None = None,
    inference_cfg: InferenceConfig | dict | None = None,
):
    """
    Args:
        model: The model to run actions on
        device: The device to use (e.g. {'cpu', 'cuda:0', 'mps'})
        snapshot_path: If defined, the path of a snapshot from which to load
            pretrained weights
        preprocessor: The preprocessor to use on images before inference
        postprocessor: The postprocessor to use on images after inference
        load_weights_only: Value for the torch.load() `weights_only` parameter.
            If False, the python pickle module is used implicitly, which is known to
                be insecure. Only set to False if you're loading data that you trust
                (e.g. snapshots that you created). For more information, see:
                    https://pytorch.org/docs/stable/generated/torch.load.html
            If None, the default value is used:
                `deeplabcut.pose_estimation_pytorch.get_load_weights_only()`
        inference_cfg: Configuration for the inference runner
    """
    super().__init__(model=model, device=device, snapshot_path=snapshot_path)
    if not isinstance(batch_size, int) or batch_size <= 0:
        raise ValueError(f"batch_size must be a positive integer; is {batch_size}")

    self.batch_size = batch_size
    self.preprocessor = preprocessor
    self.postprocessor = postprocessor

    if isinstance(inference_cfg, InferenceConfig):
        self.inference_cfg = inference_cfg
    elif isinstance(inference_cfg, dict):
        self.inference_cfg = InferenceConfig.from_dict(inference_cfg)
    elif inference_cfg is None:
        self.inference_cfg = InferenceConfig()
    else:
        raise ValueError(f"Invalid inference config: {inference_cfg}")

    if self.snapshot_path is not None and self.snapshot_path != "":
        self.load_snapshot(
            self.snapshot_path,
            self.device,
            self.model,
            weights_only=load_weights_only,
        )

    self.model.to(self.device)
    self.model.eval()

    if self.inference_cfg.compile.enabled:
        try:
            self.model = torch.compile(self.model, backend=self.inference_cfg.compile.backend)
        except Exception as e:
            warnings.warn(
                f"torch.compile failed with backend='{self.inference_cfg.compile.backend}', "
                f"falling back to eager mode. Error: {e}",
                stacklevel=2,
            )

    self._batch_list: list[torch.Tensor] = []
    self._model_kwargs: dict[str, np.ndarray | torch.Tensor] = {}

    self._contexts: list[dict] = []
    self._image_batch_sizes: list[int] = []
    self._predictions: list = []

    # Async-specific attributes
    if self.inference_cfg.multithreading.enabled:
        self._input_queue = Queue(maxsize=self.inference_cfg.multithreading.queue_length)
        self._preprocessing_thread = None
        self._stop_event = threading.Event()
        self._exception = None

inference

inference(
    images: Iterable[str | Path | ndarray] | Iterable[tuple[str | Path | ndarray, dict[str, Any]]],
    shelf_writer: ShelfWriter | None = None,
) -> list[dict[str, np.ndarray]]

Run model inference on the given dataset.

Add an option to also return head outputs (such as heatmaps)? Can be

super useful for debugging

Parameters:

Name	Type	Description	Default
images	Iterable[str | Path | ndarray] | Iterable[tuple[str | Path | ndarray, dict[str, Any]]]	the images to run inference on, optionally with context	required
shelf_writer	ShelfWriter | None	by default, data are saved in a list and returned at the end of inference. Passing a shelf manager writes data to disk on-the-fly using a "shelf" (a pickle-based, persistent, database-like object by default, resulting in constant memory footprint). The returned list is then empty.	None

Returns:

Type	Description
list[dict[str, ndarray]]	a dict containing head predictions for each image [ { "bodypart": {"poses": np.array}, "unique_bodypart": {"poses": np.array}, } ]

Source code in deeplabcut/pose_estimation_pytorch/runners/inference.py

@_inference_mode_decorator
def inference(
    self,
    images: (Iterable[str | Path | np.ndarray] | Iterable[tuple[str | Path | np.ndarray, dict[str, Any]]]),
    shelf_writer: shelving.ShelfWriter | None = None,
) -> list[dict[str, np.ndarray]]:
    """Run model inference on the given dataset.

    TODO: Add an option to also return head outputs (such as heatmaps)? Can be
     super useful for debugging

    Args:
        images: the images to run inference on, optionally with context
        shelf_writer: by default, data are saved in a list and returned at the end
            of inference. Passing a shelf manager writes data to disk on-the-fly
            using a "shelf" (a pickle-based, persistent, database-like object by
            default, resulting in constant memory footprint). The returned list is
            then empty.

    Returns:
        a dict containing head predictions for each image
        [
            {
                "bodypart": {"poses": np.array},
                "unique_bodypart": {"poses": np.array},
            }
        ]
    """
    if self.inference_cfg.multithreading.enabled:
        return self._async_inference(images, shelf_writer)
    else:
        return self._sequential_inference(images, shelf_writer)

predict abstractmethod

predict(inputs: Tensor, **kwargs) -> list[dict[str, dict[str, np.ndarray]]]

Makes predictions from a model input and output.

Returns:

Type	Description
list[dict[str, dict[str, ndarray]]]	the predictions for each of the 'batch_size' inputs

Source code in deeplabcut/pose_estimation_pytorch/runners/inference.py

@abstractmethod
def predict(self, inputs: torch.Tensor, **kwargs) -> list[dict[str, dict[str, np.ndarray]]]:
    """Makes predictions from a model input and output.

    Args:
        the inputs to the model, of shape (batch_size, ...)

    Returns:
        the predictions for each of the 'batch_size' inputs
    """

PoseInferenceRunner

Bases: InferenceRunner[PoseModel]

Runner for pose estimation inference.

Methods:

Name	Description
predict	Makes predictions from a model input and output.

Source code in deeplabcut/pose_estimation_pytorch/runners/inference.py

class PoseInferenceRunner(InferenceRunner[PoseModel]):
    """Runner for pose estimation inference."""

    def __init__(
        self,
        model: PoseModel,
        dynamic: DynamicCropper | None = None,
        **kwargs,
    ):
        super().__init__(model, **kwargs)
        self.dynamic = dynamic
        if dynamic is not None and self.batch_size != 1:
            raise ValueError("Dynamic cropping can only be used with batch size 1. Please set your batch size to 1.")

    def predict(self, inputs: torch.Tensor, **kwargs) -> list[dict[str, dict[str, np.ndarray]]]:
        """Makes predictions from a model input and output.

        Args:
            the inputs to the model, of shape (batch_size, ...)

        Returns:
            predictions for each of the 'batch_size' inputs, made by each head, e.g.
            [
                {
                    "bodypart": {"poses": np.ndarray},
                    "unique_bodypart": {"poses": np.ndarray},
                }
            ]
        """
        batch_size = len(inputs)
        if self.dynamic is not None:
            # dynamic cropping can use patches
            inputs = self.dynamic.crop(inputs)
        if self.inference_cfg.autocast.enabled:
            with torch.autocast(device_type=str(self.device)):
                outputs = self.model(inputs.to(self.device), **kwargs)
                raw_predictions = self.model.get_predictions(outputs)
        else:
            outputs = self.model(inputs.to(self.device), **kwargs)
            raw_predictions = self.model.get_predictions(outputs)

        if self.dynamic is not None:
            raw_predictions["bodypart"]["poses"] = self.dynamic.update(raw_predictions["bodypart"]["poses"])

        predictions = [
            {
                head: {pred_name: pred[b].cpu().numpy() for pred_name, pred in head_outputs.items()}
                for head, head_outputs in raw_predictions.items()
            }
            for b in range(batch_size)
        ]
        return predictions

predict

predict(inputs: Tensor, **kwargs) -> list[dict[str, dict[str, np.ndarray]]]

Makes predictions from a model input and output.

Returns:

Type	Description
list[dict[str, dict[str, ndarray]]]	predictions for each of the 'batch_size' inputs, made by each head, e.g. [ { "bodypart": {"poses": np.ndarray}, "unique_bodypart": {"poses": np.ndarray}, } ]

Source code in deeplabcut/pose_estimation_pytorch/runners/inference.py

def predict(self, inputs: torch.Tensor, **kwargs) -> list[dict[str, dict[str, np.ndarray]]]:
    """Makes predictions from a model input and output.

    Args:
        the inputs to the model, of shape (batch_size, ...)

    Returns:
        predictions for each of the 'batch_size' inputs, made by each head, e.g.
        [
            {
                "bodypart": {"poses": np.ndarray},
                "unique_bodypart": {"poses": np.ndarray},
            }
        ]
    """
    batch_size = len(inputs)
    if self.dynamic is not None:
        # dynamic cropping can use patches
        inputs = self.dynamic.crop(inputs)
    if self.inference_cfg.autocast.enabled:
        with torch.autocast(device_type=str(self.device)):
            outputs = self.model(inputs.to(self.device), **kwargs)
            raw_predictions = self.model.get_predictions(outputs)
    else:
        outputs = self.model(inputs.to(self.device), **kwargs)
        raw_predictions = self.model.get_predictions(outputs)

    if self.dynamic is not None:
        raw_predictions["bodypart"]["poses"] = self.dynamic.update(raw_predictions["bodypart"]["poses"])

    predictions = [
        {
            head: {pred_name: pred[b].cpu().numpy() for pred_name, pred in head_outputs.items()}
            for head, head_outputs in raw_predictions.items()
        }
        for b in range(batch_size)
    ]
    return predictions

build_inference_runner

build_inference_runner(
    task: Task,
    model: Module,
    device: str,
    snapshot_path: str | Path | None = None,
    batch_size: int = 1,
    preprocessor: Preprocessor | None = None,
    postprocessor: Postprocessor | None = None,
    dynamic: DynamicCropper | None = None,
    load_weights_only: bool | None = None,
    inference_cfg: InferenceConfig | dict | None = None,
    **kwargs
) -> InferenceRunner

Build a runner object according to a pytorch configuration file.

Parameters:

Name	Type	Description	Default
task	Task	the inference task to run	required
model	Module	the model to run	required
device	str	the device to use (e.g. {'cpu', 'cuda:0', 'mps'})	required
snapshot_path	str | Path | None	the snapshot from which to load the weights	None
batch_size	int	the batch size to use to run inference	1
preprocessor	Preprocessor | None	the preprocessor to use on images before inference	None
postprocessor	Postprocessor | None	the postprocessor to use on images after inference	None
dynamic	DynamicCropper | None	The DynamicCropper used for video inference, or None if dynamic cropping should not be used. Only for bottom-up pose estimation models. Should only be used when creating inference runners for video pose estimation with batch size 1.	None
load_weights_only	bool | None	Value for the torch.load() weights_only parameter. If False, the python pickle module is used implicitly, which is known to be insecure. Only set to False if you're loading data that you trust (e.g. snapshots that you created). For more information, see: https://pytorch.org/docs/stable/generated/torch.load.html If None, the default value is used: deeplabcut.pose_estimation_pytorch.get_load_weights_only()	None
inference_cfg	InferenceConfig | dict | None	Configuration for the inference runner	None
**kwargs		Other arguments for the InferenceRunner.	{}

Returns:

Type	Description
InferenceRunner	The inference runner.

Source code in deeplabcut/pose_estimation_pytorch/runners/inference.py

def build_inference_runner(
    task: Task,
    model: nn.Module,
    device: str,
    snapshot_path: str | Path | None = None,
    batch_size: int = 1,
    preprocessor: Preprocessor | None = None,
    postprocessor: Postprocessor | None = None,
    dynamic: DynamicCropper | None = None,
    load_weights_only: bool | None = None,
    inference_cfg: InferenceConfig | dict | None = None,
    **kwargs,
) -> InferenceRunner:
    """Build a runner object according to a pytorch configuration file.

    Args:
        task: the inference task to run
        model: the model to run
        device: the device to use (e.g. {'cpu', 'cuda:0', 'mps'})
        snapshot_path: the snapshot from which to load the weights
        batch_size: the batch size to use to run inference
        preprocessor: the preprocessor to use on images before inference
        postprocessor: the postprocessor to use on images after inference
        dynamic: The DynamicCropper used for video inference, or None if dynamic
            cropping should not be used. Only for bottom-up pose estimation models.
            Should only be used when creating inference runners for video pose
            estimation with batch size 1.
        load_weights_only: Value for the torch.load() `weights_only` parameter.
            If False, the python pickle module is used implicitly, which is known to
            be insecure. Only set to False if you're loading data that you trust (e.g.
            snapshots that you created). For more information, see:
                https://pytorch.org/docs/stable/generated/torch.load.html
            If None, the default value is used:
                `deeplabcut.pose_estimation_pytorch.get_load_weights_only()`
        inference_cfg: Configuration for the inference runner
        **kwargs: Other arguments for the InferenceRunner.

    Returns:
        The inference runner.
    """
    kwargs = dict(
        model=model,
        device=device,
        snapshot_path=snapshot_path,
        batch_size=batch_size,
        preprocessor=preprocessor,
        postprocessor=postprocessor,
        load_weights_only=load_weights_only,
        inference_cfg=inference_cfg,
        **kwargs,
    )

    if task == Task.DETECT:
        if dynamic is not None:
            raise ValueError(
                "The DynamicCropper can only be used for pose estimation; not object "
                "detection. Please turn off dynamic cropping."
            )
        return DetectorInferenceRunner(**kwargs)

    if task != Task.BOTTOM_UP:
        if dynamic is not None and not isinstance(dynamic, TopDownDynamicCropper):
            print(
                "Turning off dynamic cropping. It should only be used for bottom-up "
                f"pose estimation models, but you are using a {task} model. To use "
                f"dynamic cropping with {task}, use a TopDownDynamicCropper."
            )
            dynamic = None

    if task == Task.COND_TOP_DOWN:
        return CTDInferenceRunner(**kwargs)

    return PoseInferenceRunner(dynamic=dynamic, **kwargs)