diff --git a/_toc.yml b/_toc.yml index 5d2a303f45..6d70d0299c 100644 --- a/_toc.yml +++ b/_toc.yml @@ -77,7 +77,6 @@ parts: - file: docs/recipes/OpenVINO - file: docs/recipes/flip_and_rotate - file: docs/recipes/pose_cfg_file_breakdown - - file: docs/recipes/fmpose3d - file: docs/recipes/publishing_notebooks_into_the_DLC_main_cookbook - caption: Hardware Tips diff --git a/deeplabcut/modelzoo/fmpose_3d/__init__.py b/deeplabcut/modelzoo/fmpose_3d/__init__.py deleted file mode 100644 index fc1ada573b..0000000000 --- a/deeplabcut/modelzoo/fmpose_3d/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -""" -DeepLabCut2.0-3.0 Toolbox (deeplabcut.org) -© A. & M. Mathis Labs -https://github.com/DeepLabCut/DeepLabCut -Please see AUTHORS for contributors. -https://github.com/DeepLabCut/DeepLabCut/blob/main/AUTHORS -Licensed under GNU Lesser General Public License v3.0 -""" diff --git a/deeplabcut/modelzoo/fmpose_3d/fmpose3d.py b/deeplabcut/modelzoo/fmpose_3d/fmpose3d.py deleted file mode 100644 index a7ce458ef4..0000000000 --- a/deeplabcut/modelzoo/fmpose_3d/fmpose3d.py +++ /dev/null @@ -1,56 +0,0 @@ -""" -DeepLabCut2.0-3.0 Toolbox (deeplabcut.org) -© A. & M. Mathis Labs -https://github.com/DeepLabCut/DeepLabCut -Please see AUTHORS for contributors. -https://github.com/DeepLabCut/DeepLabCut/blob/main/AUTHORS -Licensed under GNU Lesser General Public License v3.0 -""" - -from fmpose3d import ( - FMPose3DInference, - FMPose3DConfig, - SupportedModel, -) - - -def get_fmpose3d_inference_api( - model_type: SupportedModel = "fmpose3d_humans", - snapshot_path: str | None = None, - device: str | None = None, - config_kwargs: dict = {}, - ) -> FMPose3DInference: - """ - Get a FMPose3DInference API for a given model type and snapshot path. - - Args: - model_type: one of the supported model types: "fmpose3d_humans", "fmpose3d_animals", - snapshot_path: The path to the snapshot file. If None, FMPose3D will download the default snapshot. - device: The device to use. If None, the device will be inferred from the environment. - config_kwargs: Additional keyword arguments to pass to the FMPose3DConfig. - Returns: - FMPose3DInference: An FMPose3DInference API runner. - - Example Usages - ```python - # Initialize the API (downloads the default weights automatically from huggingface) - fmpose = get_fmpose3d_inference_api( - model_type="fmpose3d_animals", - device="cuda:0", - ) - - # Run inference on an image - predictions_3d = fmpose.predict(source="path/to/image.jpg") # or (H, W, 3) numpy array - - # Lift 2d predictions to 3d - keypoints_2d = np.random.rand(num_frames, num_joints, 2) - predictions_3d = fmpose.pose_3d(keypoints_2d=keypoints_2d) - ``` - """ - model_config = FMPose3DConfig(model_type=model_type, **config_kwargs) - fmpose3d_api = FMPose3DInference( - model_config, - model_weights_path=snapshot_path, - device=device - ) - return fmpose3d_api \ No newline at end of file diff --git a/deeplabcut/modelzoo/models_to_framework.json b/deeplabcut/modelzoo/models_to_framework.json index 1f2174a5fe..760af11c83 100644 --- a/deeplabcut/modelzoo/models_to_framework.json +++ b/deeplabcut/modelzoo/models_to_framework.json @@ -3,5 +3,7 @@ "hrnet_w32": "pytorch", "resnet_50": "pytorch", "rtmpose_s": "pytorch", - "rtmpose_x": "pytorch" + "rtmpose_x": "pytorch", + "fmpose3d_humans": "pytorch", + "fmpose3d_animals": "pytorch", } diff --git a/deeplabcut/modelzoo/video_inference.py b/deeplabcut/modelzoo/video_inference.py index 90e082880c..51b435cd88 100644 --- a/deeplabcut/modelzoo/video_inference.py +++ b/deeplabcut/modelzoo/video_inference.py @@ -11,7 +11,9 @@ from __future__ import annotations import json +import logging import os +import warnings from pathlib import Path from typing import Optional, Union @@ -34,6 +36,8 @@ video_to_frames, ) +logger = logging.getLogger(__name__) + def get_checkpoint_epoch(checkpoint_path): """ @@ -80,6 +84,7 @@ def video_inference_superanimal( customized_model_config: Optional[str] = None, plot_bboxes: bool = True, create_labeled_video: bool = True, + fmpose_return_3d: bool = False, ): """ This function performs inference on videos using a pretrained SuperAnimal model. @@ -177,6 +182,13 @@ def video_inference_superanimal( create_labeled_video (bool): Specifies if a labeled video needs to be created, True by default. + fmpose_return_3d (bool): + Only used when ``model_name`` starts with ``"fmpose3d"``. + If True, include in-memory 3D poses in the return payload + (per video: ``{"df_2d": ..., "df_3d": ...}``). + If False (default), keep the legacy return payload with only + the 2D DataFrame per video. + Raises: NotImplementedError: If the model is not found in the modelzoo. @@ -316,8 +328,10 @@ def video_inference_superanimal( """ if scale_list is None: scale_list = [] - - print(f"Running video inference on {videos} with {superanimal_name}_{model_name}") + if not model_name.startswith("fmpose3d"): + print( + f"Running video inference on {videos} with {superanimal_name}_{model_name}" + ) dlc_root_path = get_deeplabcut_path() modelzoo_path = os.path.join(dlc_root_path, "modelzoo") available_architectures = json.load( @@ -352,6 +366,42 @@ def video_inference_superanimal( create_labeled_video=create_labeled_video, ) elif framework == "pytorch": + if model_name.startswith("fmpose3d"): + logger.info("Running video inference on %s using %s", videos, model_name) + + recommended_superanimal_name = { + "fmpose3d_animals": "quadruped", + "fmpose3d_humans": "human", + }.get(model_name) + + provided_superanimal_name = superanimal_name or "" + if superanimal_name != recommended_superanimal_name: + warnings.warn( + "For FMPose3D models, model selection is driven by 'model_name'. But for API " + "consistency, it is recommended to set 'superanimal_name' to the corresponding value." + f"Provided superanimal_name={provided_superanimal_name!r} differs from the " + f"recommended value for {model_name!r}: " + f"{recommended_superanimal_name!r}.", + stacklevel=2, + ) + + from deeplabcut.pose_estimation_pytorch.modelzoo.fmpose_3d.inference import ( + _video_inference_fmpose3d, + ) + + return _video_inference_fmpose3d( + video_paths=videos, + model_name=model_name, + max_individuals=max_individuals, + pcutoff=pcutoff, + batch_size=batch_size, + dest_folder=dest_folder, + device=device, + create_labeled_video=create_labeled_video, + cropping=cropping, + include_3d_in_return=fmpose_return_3d, + ) + torchvision_detector_name = None if superanimal_name != "superanimal_humanbody" and detector_name is None: raise ValueError( @@ -483,9 +533,9 @@ def video_inference_superanimal( if superanimal_name != "superanimal_humanbody": detector_snapshot_prefix = f"snapshot-{detector_name}" - config["detector"]["runner"][ - "snapshot_prefix" - ] = detector_snapshot_prefix + config["detector"]["runner"]["snapshot_prefix"] = ( + detector_snapshot_prefix + ) # the model config's parameters need to be updated for adaptation training model_config_path = model_folder / "pytorch_config.yaml" diff --git a/deeplabcut/modelzoo/fmpose_3d/README.md b/deeplabcut/pose_estimation_pytorch/modelzoo/fmpose_3d/README.md similarity index 95% rename from deeplabcut/modelzoo/fmpose_3d/README.md rename to deeplabcut/pose_estimation_pytorch/modelzoo/fmpose_3d/README.md index 9a6ce4053a..f72d35f9af 100644 --- a/deeplabcut/modelzoo/fmpose_3d/README.md +++ b/deeplabcut/pose_estimation_pytorch/modelzoo/fmpose_3d/README.md @@ -9,4 +9,4 @@ Model weights are hosted on HuggingFace Hub and are downloaded automatically when no local path is provided. The library is installable via `pip install fmpose3d` and requires Python >= 3.8. -For a full overview and documentation on the API, see https://github.com/AdaptiveMotorControlLab/FMPose3D. +For a full overview and documentation on the API, see https://github.com/AdaptiveMotorControlLab/FMPose3D. diff --git a/deeplabcut/pose_estimation_pytorch/modelzoo/fmpose_3d/__init__.py b/deeplabcut/pose_estimation_pytorch/modelzoo/fmpose_3d/__init__.py new file mode 100644 index 0000000000..1ea902fbca --- /dev/null +++ b/deeplabcut/pose_estimation_pytorch/modelzoo/fmpose_3d/__init__.py @@ -0,0 +1,15 @@ +""" +DeepLabCut2.0-3.0 Toolbox (deeplabcut.org) +© A. & M. Mathis Labs +https://github.com/DeepLabCut/DeepLabCut +Please see AUTHORS for contributors. +https://github.com/DeepLabCut/DeepLabCut/blob/main/AUTHORS +Licensed under GNU Lesser General Public License v3.0 +""" + +# NOTE: this module may contain items that need refactoring during +# the keypoint migration. + +# kpt_refactor - Needs attention when refactoring keypoints +# i_o - This module writes keypoints to disk +# pandas - This module relies on pandas (might be moved to polars) diff --git a/deeplabcut/pose_estimation_pytorch/modelzoo/fmpose_3d/fmpose3d.py b/deeplabcut/pose_estimation_pytorch/modelzoo/fmpose_3d/fmpose3d.py new file mode 100644 index 0000000000..41f835afb9 --- /dev/null +++ b/deeplabcut/pose_estimation_pytorch/modelzoo/fmpose_3d/fmpose3d.py @@ -0,0 +1,137 @@ +""" +DeepLabCut2.0-3.0 Toolbox (deeplabcut.org) +© A. & M. Mathis Labs +https://github.com/DeepLabCut/DeepLabCut +Please see AUTHORS for contributors. +https://github.com/DeepLabCut/DeepLabCut/blob/main/AUTHORS +Licensed under GNU Lesser General Public License v3.0 +""" + +from dataclasses import dataclass + +from fmpose3d import ( + FMPose3DConfig, + FMPose3DInference, + SupportedModel, +) + + +@dataclass(frozen=True) +class FMPose3DModelMetadata: + """Metadata for an FMPose3D model variant.""" + + superanimal_name: str + bodyparts: tuple[str, ...] + + @property + def num_bodyparts(self) -> int: + return len(self.bodyparts) + + def build_model_cfg(self, max_individuals: int) -> dict: + """Build a DLC-compatible model_cfg dict for create_df_from_prediction.""" + return { + "metadata": { + "bodyparts": list(self.bodyparts), + "unique_bodyparts": [], + "individuals": [f"individual{i + 1}" for i in range(max_individuals)], + }, + } + + +FMPOSE3D_MODEL_METADATA: dict[str, FMPose3DModelMetadata] = { + "fmpose3d_humans": FMPose3DModelMetadata( + superanimal_name="superanimal_humanbody", + bodyparts=( + "pelvis", + "right_hip", + "right_knee", + "right_ankle", + "left_hip", + "left_knee", + "left_ankle", + "spine", + "thorax", + "neck", + "head", + "left_shoulder", + "left_elbow", + "left_wrist", + "right_shoulder", + "right_elbow", + "right_wrist", + ), + ), + "fmpose3d_animals": FMPose3DModelMetadata( + superanimal_name="superanimal_quadruped", + bodyparts=( + "left_eye", + "right_eye", + "nose", + "neck", + "root_of_tail", + "left_shoulder", + "left_elbow", + "left_front_paw", + "right_shoulder", + "right_elbow", + "right_front_paw", + "left_hip", + "left_knee", + "left_back_paw", + "right_hip", + "right_knee", + "right_back_paw", + "withers", + "throat", + "left_ear", + "right_ear", + "mouth", + "chin", + "left_hock", + "right_hock", + "tail_tip", + ), + ), +} + + +def get_fmpose3d_inference_api( + model_type: SupportedModel = "fmpose3d_humans", + snapshot_path: str | None = None, + device: str | None = None, + config_kwargs: dict = {}, +) -> FMPose3DInference: + """ + Get a FMPose3DInference API for a given model type and snapshot path. + + Args: + model_type: one of the supported model types: "fmpose3d_humans", "fmpose3d_animals", + snapshot_path: The path to the snapshot file. If None, FMPose3D will download the default snapshot. + device: The device to use. If None, the device will be inferred from the environment. + config_kwargs: Additional keyword arguments to pass to the FMPose3DConfig. + Returns: + FMPose3DInference: An FMPose3DInference API runner. + + Example Usages + ```python + # Initialize the API (downloads the default weights automatically from huggingface) + fmpose = get_fmpose3d_inference_api( + model_type="fmpose3d_animals", + device="cuda:0", + ) + + # Run inference on an image + predictions_3d = fmpose.predict(source="path/to/image.jpg") # or (H, W, 3) numpy array + + # Lift 2d predictions to 3d + keypoints_2d = np.random.rand(num_frames, num_joints, 2) + predictions_3d = fmpose.pose_3d(keypoints_2d=keypoints_2d) + ``` + """ + model_config = FMPose3DConfig(model_type=model_type, **config_kwargs) + fmpose3d_api = FMPose3DInference( + model_config, + model_weights_path=snapshot_path, + device=device, + ) + return fmpose3d_api diff --git a/deeplabcut/pose_estimation_pytorch/modelzoo/fmpose_3d/inference.py b/deeplabcut/pose_estimation_pytorch/modelzoo/fmpose_3d/inference.py new file mode 100644 index 0000000000..eda263d76d --- /dev/null +++ b/deeplabcut/pose_estimation_pytorch/modelzoo/fmpose_3d/inference.py @@ -0,0 +1,262 @@ +# +# DeepLabCut Toolbox (deeplabcut.org) +# © A. & M.W. Mathis Labs +# https://github.com/DeepLabCut/DeepLabCut +# +# Please see AUTHORS for contributors. +# https://github.com/DeepLabCut/DeepLabCut/blob/main/AUTHORS +# +# Licensed under GNU Lesser General Public License v3.0 +# +import json +import logging +from pathlib import Path + +import numpy as np + +from deeplabcut.modelzoo.utils import get_superanimal_colormaps +from deeplabcut.pose_estimation_pytorch.apis.videos import ( + VideoIterator, + create_df_from_prediction, +) +from deeplabcut.pose_estimation_pytorch.modelzoo.fmpose_3d.fmpose3d import ( + FMPOSE3D_MODEL_METADATA, + get_fmpose3d_inference_api, +) +from deeplabcut.utils import auxiliaryfunctions_3d +from deeplabcut.utils.make_labeled_video import create_video + +logger = logging.getLogger(__name__) + + +class NumpyEncoder(json.JSONEncoder): + """Special json encoder for numpy types.""" + + def default(self, obj): + if isinstance(obj, np.ndarray): + return obj.tolist() # Convert ndarray to list + return json.JSONEncoder.default(self, obj) + + +def _pose2d_to_dlc_predictions( + pose_2d, + max_individuals: int, + num_bodyparts: int, +) -> list[dict[str, np.ndarray]]: + """Convert FMPose3D 2D output to DLC per-frame prediction format.""" + all_kpts = np.asarray(pose_2d.keypoints) + all_scores = np.asarray(pose_2d.scores) + if all_kpts.ndim != 4 or all_scores.ndim != 3: + raise ValueError( + "Expected pose_2d keypoints/scores shaped as (num_persons, num_frames, num_bodyparts, {2 or score})." + ) + + num_frames = all_kpts.shape[1] + num_persons = all_kpts.shape[0] + per_frame: list[dict[str, np.ndarray]] = [] + for frame_idx in range(num_frames): + n_det = min(num_persons, max_individuals) + bodyparts_array = np.zeros((max_individuals, num_bodyparts, 3)) + bodyparts_array[:n_det, :, :2] = all_kpts[:n_det, frame_idx, :num_bodyparts, :2] + bodyparts_array[:n_det, :, 2] = all_scores[:n_det, frame_idx, :num_bodyparts] + per_frame.append({"bodyparts": bodyparts_array}) + return per_frame + + +# NOTE: i_o; pandas; kpt_refactor; this function may need to change in the future, to improve dataframe +# i/o migration to validated keypoint schemas (parquet) +def _poses3d_to_dataframe(poses_3d: list[np.ndarray], df_2d, scorer_3d: str): + """Create and fill a 3D dataframe using the shared auxiliary helper.""" + df_3d, scorer_3d, bodyparts = auxiliaryfunctions_3d.create_empty_df( + df_2d, scorer_3d, "3d" + ) + n_frames = len(poses_3d) + n_bodyparts = len(bodyparts) + arr = np.full((n_frames, n_bodyparts, 3), np.nan, dtype=float) + + for frame_idx, pose in enumerate(poses_3d): + pose_np = np.asarray(pose) + if pose_np.ndim == 3: + if pose_np.shape[0] == 0: + continue + pose_np = pose_np[0] + if pose_np.ndim != 2 or pose_np.shape[-1] != 3: + continue + + n = min(n_bodyparts, pose_np.shape[0]) + arr[frame_idx, :n] = pose_np[:n] + + xyz_cols = [(scorer_3d, bp, coord) for bp in bodyparts for coord in ("x", "y", "z")] + df_3d.loc[:, xyz_cols] = arr.reshape(n_frames, -1) + + return df_3d + + +def _video_inference_fmpose3d( + video_paths: str | Path | list[str | Path], + model_name: str, + max_individuals: int = 1, + pcutoff: float = 0.1, + batch_size: int = 1, + dest_folder: str | Path | None = None, + device: str | None = None, + create_labeled_video: bool = True, + cropping: list[int] | None = None, + include_3d_in_return: bool = False, +) -> dict: + """Perform FMPose3D video inference with a lightweight DLC loop.""" + import torch + from tqdm import tqdm + + if max_individuals != 1: + logger.warning( + "FMPose3D 3D lifting currently supports only one individual. " + "Clamping max_individuals=%s to 1 for this pipeline.", + max_individuals, + ) + max_individuals = 1 + + if device is None or device == "auto": + device = "cuda:0" if torch.cuda.is_available() else "cpu" + + if isinstance(video_paths, (str, Path)): + video_paths = [video_paths] + + if model_name not in FMPOSE3D_MODEL_METADATA: + raise ValueError( + f"Unsupported FMPose3D model '{model_name}'. " + "Use one of: " + ", ".join(sorted(FMPOSE3D_MODEL_METADATA.keys())) + ) + metadata = FMPOSE3D_MODEL_METADATA[model_name] + model_cfg = metadata.build_model_cfg(max_individuals) + bodyparts = list(metadata.bodyparts) + num_bodyparts = metadata.num_bodyparts + superanimal_name = metadata.superanimal_name + + api = get_fmpose3d_inference_api(model_type=model_name, device=device) + + dest_folder = ( + Path(video_paths[0]).parent if dest_folder is None else Path(dest_folder) + ) + dest_folder.mkdir(parents=True, exist_ok=True) + + if create_labeled_video: + superanimal_colormaps = get_superanimal_colormaps() + colormap = superanimal_colormaps[superanimal_name] + + dlc_scorer = f"DLC_{model_name}" + results = {} + + for video_path in video_paths: + print(f"Processing video {video_path} with {model_name}") + video = VideoIterator(video_path, cropping=cropping) + vid_w, vid_h = video.dimensions + + predictions_2d: list[dict[str, np.ndarray]] = [] + all_poses_3d: list[np.ndarray] = [] + warned_multi_person_2d = False + + def _process_batch(frames: list[np.ndarray]) -> None: + nonlocal warned_multi_person_2d + pose_2d = api.prepare_2d(source=np.stack(frames)) + num_detected = int(np.asarray(pose_2d.keypoints).shape[0]) + if num_detected > 1 and not warned_multi_person_2d: + logger.warning( + "Multiple 2D detections (%s) were found, but FMPose3D 3D lifting uses only the first individual.", + num_detected, + ) + warned_multi_person_2d = True + predictions_2d.extend( + _pose2d_to_dlc_predictions( + pose_2d, + max_individuals=max_individuals, + num_bodyparts=num_bodyparts, + ) + ) + try: + pose_3d = api.pose_3d( + keypoints_2d=pose_2d.keypoints, + image_size=pose_2d.image_size, + ) + all_poses_3d.extend(np.asarray(pose_3d.poses_3d)) + except ValueError as e: + logger.info( + "Skipping 3D lifting for batch due to invalid 2D result: %s", e + ) + all_poses_3d.extend([np.zeros((0, num_bodyparts, 3)) for _ in frames]) + + batch: list[np.ndarray] = [] + for frame in tqdm(video, desc="FMPose3D inference"): + batch.append(frame) + if len(batch) == batch_size: + _process_batch(batch) + batch.clear() + if batch: + _process_batch(batch) + + output_prefix = f"{Path(video_path).stem}_{dlc_scorer}" + output_h5 = dest_folder / f"{output_prefix}.h5" + + print(f"Saving 2D results to {dest_folder}") + df = create_df_from_prediction( + predictions=predictions_2d, + dlc_scorer=dlc_scorer, + multi_animal=True, + model_cfg=model_cfg, + output_path=dest_folder, + output_prefix=output_prefix, + ) + scorer_3d = f"{dlc_scorer}_3d" + df_3d = _poses3d_to_dataframe(all_poses_3d, df, scorer_3d) + output_3d_h5 = dest_folder / f"{output_prefix}_3d.h5" + df_3d.to_hdf(output_3d_h5, key="df_with_missing", mode="w", format="table") + print(f"3D dataframe saved to {output_3d_h5}") + + if include_3d_in_return: + results[video_path] = { + "df_2d": df, + "df_3d": df_3d, + } + else: + results[video_path] = df + + output_json = dest_folder / f"{output_prefix}.json" + with open(output_json, "w") as f: + json.dump(predictions_2d, f, cls=NumpyEncoder) + + poses_3d_serialisable = [ + pose.tolist() if isinstance(pose, np.ndarray) else pose + for pose in all_poses_3d + ] + output_3d_json = dest_folder / f"{output_prefix}_3d.json" + with open(output_3d_json, "w") as f: + json.dump( + { + "model": model_name, + "bodyparts": bodyparts, + "poses_3d": poses_3d_serialisable, + }, + f, + ) + print(f"3D predictions saved to {output_3d_json}") + + if create_labeled_video: + bbox = cropping + if cropping is None: + bbox = (0, vid_w, 0, vid_h) + output_video = dest_folder / f"{output_prefix}_labeled.mp4" + create_video( + video_path, + output_h5, + pcutoff=pcutoff, + fps=video.fps, + bbox=bbox, + cmap=colormap, + output_path=output_video, + plot_bboxes=False, + bboxes_list=[], + bboxes_pcutoff=0.0, + ) + print(f"Video with predictions was saved as {output_video}") + + return results diff --git a/docs/ModelZoo.md b/docs/ModelZoo.md index fe5f3496b4..4ae7c8dbaa 100644 --- a/docs/ModelZoo.md +++ b/docs/ModelZoo.md @@ -93,6 +93,33 @@ deeplabcut.video_inference_superanimal([video_path], video_adapt = False) ``` +### Practical example: FMPose3D monocular 3D inference + +For FMPose3D models, use `model_name="fmpose3d_animals"` or +`model_name="fmpose3d_humans"`. Model selection is still determined by +`model_name`, but to stay aligned with SuperAnimal naming conventions use: +`superanimal_name="superanimal_quadruped"` for `fmpose3d_animals`, and +`superanimal_name="superanimal_humanbody"` for `fmpose3d_humans`. + +Like the 2D superanimal models, this inference branch writes +intermediate 2D predictions to `