diff --git a/deeplabcut/pose_estimation_pytorch/modelzoo/inference.py b/deeplabcut/pose_estimation_pytorch/modelzoo/inference.py index 3582d20d9..766abe5d9 100644 --- a/deeplabcut/pose_estimation_pytorch/modelzoo/inference.py +++ b/deeplabcut/pose_estimation_pytorch/modelzoo/inference.py @@ -9,6 +9,7 @@ # Licensed under GNU Lesser General Public License v3.0 # import json +import logging from pathlib import Path import numpy as np @@ -30,6 +31,8 @@ ) from deeplabcut.utils.make_labeled_video import create_video +logger = logging.getLogger(__name__) + class NumpyEncoder(json.JSONEncoder): """Special json encoder for numpy types.""" @@ -170,6 +173,8 @@ def _video_inference_superanimal( pose_runner=pose_runner, detector_runner=detector_runner, ) + if not predictions: + raise RuntimeError(f"No pose predictions were made for video {video_path}. Were no individuals detected?") bbox_keys_in_predictions = {"bboxes", "bbox_scores"} bboxes_list = [ diff --git a/deeplabcut/pose_estimation_pytorch/modelzoo/utils.py b/deeplabcut/pose_estimation_pytorch/modelzoo/utils.py index 991115610..de9bfd05b 100644 --- a/deeplabcut/pose_estimation_pytorch/modelzoo/utils.py +++ b/deeplabcut/pose_estimation_pytorch/modelzoo/utils.py @@ -26,6 +26,9 @@ # COCO category ID for the "person" class. COCO_PERSON_CATEGORY_ID = 1 +MODEL_FILENAME_MAPPING = { + "superanimal_humanbody_rtmpose_x": "rtmpose-x_simcc-body7.pt", +} def get_model_configs_folder_path() -> Path: @@ -126,10 +129,15 @@ def download_super_animal_snapshot(dataset: str, model_name: str) -> Path: model_filename = f"{model_name}.pt" model_path = snapshot_dir / model_filename + source_filename = MODEL_FILENAME_MAPPING.get(model_name, model_filename) + if model_filename == source_filename: + rename_mapping = None + else: + rename_mapping = {source_filename: model_filename} download_huggingface_model( model_name, target_dir=str(snapshot_dir), - rename_mapping={model_filename: model_filename}, + rename_mapping=rename_mapping, ) if not model_path.exists(): raise RuntimeError(f"Failed to download {model_name} to {model_path}")