Skip to content

Fix deeplabcut.analyze_images() with CTD model #2990

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MMathisLab merged 6 commits into from
Sep 18, 2025
Merged

Fix deeplabcut.analyze_images() with CTD model #2990

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
604ab5b
Add ctd_conditions argument to analyze_images() method
maximpavliv May 21, 2025
db0eda0
Add torch_kwargs to analyze_images in compat.py
maximpavliv May 21, 2025
2d4dd61
Fix create_labeled_images() with CTD predictions crash
maximpavliv May 21, 2025
1f102b4
Black formatting
maximpavliv May 21, 2025
5f5c4eb
Merge branch 'main' into maxim/fix_ctd_analyze_images
maximpavliv Sep 12, 2025
9131ba9
Merge branch 'main' into maxim/fix_ctd_analyze_images
maximpavliv Sep 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions deeplabcut/compat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1091,6 +1091,7 @@ def create_tracking_dataset(
)
elif engine == Engine.PYTORCH:
from deeplabcut.pose_estimation_pytorch.apis import create_tracking_dataset

return create_tracking_dataset(
config,
videos,
Expand Down Expand Up @@ -1126,6 +1127,7 @@ def analyze_images(
pcutoff: float | None = None,
bbox_pcutoff: float | None = None,
plot_skeleton: bool = False,
**torch_kwargs,
) -> dict[str, dict[str, np.ndarray | np.ndarray]]:
"""Analyzes images with a DeepLabCut model and stores the output in an H5 file.

Expand Down Expand Up @@ -1213,6 +1215,9 @@ def analyze_images(
If a skeleton is defined in the project's config.yaml, whether
to plot the skeleton connecting the predicted bodyparts on the images.

torch_kwargs:
Any extra parameters to pass to the PyTorch API, such as ``ctd_conditions``

Returns
-------
A dictionary mapping image paths (as strings) to model predictions.
Expand Down Expand Up @@ -1275,6 +1280,7 @@ def analyze_images(
pcutoff=pcutoff,
bbox_pcutoff=bbox_pcutoff,
plot_skeleton=plot_skeleton,
**torch_kwargs,
)

raise NotImplementedError(f"This function is not implemented for {engine}")
Expand Down
41 changes: 41 additions & 0 deletions deeplabcut/pose_estimation_pytorch/apis/analyze_images.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import deeplabcut.pose_estimation_pytorch.modelzoo as modelzoo
from deeplabcut.core.engine import Engine
from deeplabcut.modelzoo.utils import get_superanimal_colormaps
from deeplabcut.pose_estimation_pytorch.apis.ctd import get_condition_provider
from deeplabcut.pose_estimation_pytorch.apis.utils import (
get_detector_inference_runner,
build_predictions_dataframe,
Expand All @@ -38,6 +39,7 @@
parse_snapshot_index_for_analysis,
get_filtered_coco_detector_inference_runner,
)
from deeplabcut.pose_estimation_pytorch.data.ctd import CondFromModel
from deeplabcut.pose_estimation_pytorch.modelzoo.utils import update_config
from deeplabcut.pose_estimation_pytorch.task import Task
from deeplabcut.pose_estimation_pytorch.utils import resolve_device
Expand Down Expand Up @@ -207,6 +209,7 @@ def superanimal_analyze_images(
device=device,
progress_bar=progress_bar,
filtered_detector_config=filtered_detector_config,
# TODO: when COND_TOP_DOWN SuperAnimal models will be released - create & pass a conditions provider

@MMathisLab MMathisLab May 23, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe already add this; we can release the mouse model now cc @AlexEMG @n-poulsen ?

@MMathisLab MMathisLab Jun 13, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc again @AlexEMG and @n-poulsen ...

Copilot AI Jun 26, 2025

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[nitpick] Consider either resolving this TODO by implementing the condition provider logic for COND_TOP_DOWN SuperAnimal models or referencing an issue for future tracking to avoid stale comments.

Suggested change
# TODO: when COND_TOP_DOWN SuperAnimal models will be released - create & pass a conditions provider
# TODO: when COND_TOP_DOWN SuperAnimal models will be released - create & pass a conditions provider
# Tracking issue: https://github.com/DeepLabCut/DeepLabCut/issues/<issue_number>

Copilot uses AI. Check for mistakes.

@MMathisLab MMathisLab Jun 26, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc @maximpavliv @yeshaokai

)

skeleton_bodyparts = config.get("skeleton", [])
Expand Down Expand Up @@ -249,6 +252,7 @@ def analyze_images(
pcutoff: float | None = None,
bbox_pcutoff: float | None = None,
plot_skeleton: bool = True,
ctd_conditions: dict | CondFromModel | None = None,
) -> dict[str, dict]:
"""Runs analysis on images using a pose model.

Expand Down Expand Up @@ -280,6 +284,14 @@ def analyze_images(
None or in (0, 1). If None, it is read from the project configuration file.
plot_skeleton: If a skeleton is defined in the model configuration file, whether
to plot the skeleton connecting the predicted bodyparts on the images.
ctd_conditions: Only for CTD models. If None, the configuration for the
condition provider will be loaded from the pytorch_config file (under the
"data": "conditions"). If the ctd_conditions is given as a dict, creates a
CondFromModel from the dict. Otherwise, a CondFromModel can be given
directly. Example configuration:
```
ctd_conditions = {"shuffle": 17, "snapshot": "snapshot-best-190.pt"}
```

Returns:
A dictionary mapping each image filename to the different types of predictions
Expand Down Expand Up @@ -311,6 +323,22 @@ def analyze_images(
detector_snapshot_index, train_folder, Task.DETECT
)[0]

# Load the BU model for the conditions provider
cond_provider = None
if pose_task == Task.COND_TOP_DOWN:
if ctd_conditions is None:
cond_provider = get_condition_provider(
condition_cfg=model_cfg["data"]["conditions"],
config=config,
)
elif isinstance(ctd_conditions, dict):
cond_provider = get_condition_provider(
condition_cfg=ctd_conditions,
config=config,
)
else:
cond_provider = ctd_conditions

predictions = analyze_image_folder(
model_cfg=model_cfg,
images=images,
Expand All @@ -320,6 +348,7 @@ def analyze_images(
device=device,
max_individuals=max_individuals,
progress_bar=progress_bar,
cond_provider=cond_provider,
)

if len(predictions) == 0:
Expand Down Expand Up @@ -409,6 +438,7 @@ def analyze_image_folder(
max_individuals: int | None = None,
progress_bar: bool = True,
filtered_detector_config: dict | None = None,
cond_provider: CondFromModel | None = None,
) -> dict[str, dict[str, np.ndarray | np.ndarray]]:
"""Runs pose inference on a folder of images and returns the predictions

Expand All @@ -428,6 +458,7 @@ def analyze_image_folder(
progress_bar: Whether to display a progress bar when running inference.
filtered_detector_config: If using a filtered torchvision detector instead of a saved detector snapshot,
specify the filtered detector configuration
cond_provider: If using a CTD model - this parameter is needed to provide the conditions

Returns:
A dictionary mapping each image filename to the different types of predictions
Expand Down Expand Up @@ -456,11 +487,18 @@ def analyze_image_folder(
if device is None:
device = resolve_device(model_cfg)

if pose_task == Task.COND_TOP_DOWN and cond_provider is None:
raise ValueError(
"A conditions provider must be specified for image analysis when using cond-top-down models"
f" Please specify the `cond_provider` parameter."
)

pose_runner = get_pose_inference_runner(
model_config=model_cfg,
snapshot_path=snapshot_path,
device=device,
max_individuals=max_individuals,
cond_provider=cond_provider,
)

image_suffixes = ".png", ".jpg", ".jpeg"
Expand Down Expand Up @@ -523,6 +561,7 @@ def plot_images_coco(
detector_path: str | Path | None = None,
device: str | None = None,
max_individuals: int | None = None,
cond_provider: CondFromModel | None = None,
) -> list[dict]:
"""
Runs pose inference on a folder of images from a COCO dataset, and plots all
Expand All @@ -538,6 +577,7 @@ def plot_images_coco(
if a top-down model was used.
device: The device on which to run image inference
max_individuals: The maximum number of individuals to detect in an image.
cond_provider: If using a CTD model - this parameter is needed to provide the conditions

Returns:
A list of dictionaries containing predictions made on each image.
Expand Down Expand Up @@ -583,6 +623,7 @@ def plot_images_coco(
device=device,
max_individuals=max_individuals,
progress_bar=True,
cond_provider=cond_provider,
)

os.makedirs(out_path, exist_ok=True)
Expand Down
6 changes: 3 additions & 3 deletions deeplabcut/pose_estimation_pytorch/apis/visualization.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,7 @@ def create_labeled_images(
if color_by_individual:
kwargs["c"] = cmap(idx / total_idv)
else:
c = np.linspace(0, 1, total_bodyparts)[:len(pose)][mask]
c = np.linspace(0, 1, total_bodyparts)[: len(pose)][mask]
kwargs["c"] = c
kwargs["cmap"] = cmap

Expand All @@ -139,14 +139,14 @@ def create_labeled_images(
kwargs["c"] = cmap(1)
else:
c = np.linspace(0, 1, total_bodyparts)
kwargs["c"] = c[-len(unique_pred):][mask]
kwargs["c"] = c[-len(unique_pred) :][mask]
kwargs["cmap"] = cmap

xy = xy[mask]
ax.scatter(xy[:, 0], xy[:, 1], **kwargs)

# plot bounding boxes
if "bboxes" in image_predictions:
if "bboxes" in image_predictions and "bbox_scores" in image_predictions:
bboxes = image_predictions["bboxes"]
bbox_scores = image_predictions["bbox_scores"]
for idx, (bbox, score) in enumerate(zip(bboxes, bbox_scores)):
Expand Down
Loading