diff --git a/README.md b/README.md index 2da193d..ee32252 100644 --- a/README.md +++ b/README.md @@ -36,6 +36,12 @@ conda activate fmpose_3d pip install fmpose3d ``` +For the animal pipeline, install the optional DeepLabCut dependency: + +```bash +pip install "fmpose3d[animals]" +``` + ## Demos ### Testing on in-the-wild images (humans) diff --git a/animals/demo/vis_animals.py b/animals/demo/vis_animals.py index 459c2bc..5772060 100644 --- a/animals/demo/vis_animals.py +++ b/animals/demo/vis_animals.py @@ -46,7 +46,15 @@ from fmpose3d.models import get_model CFM = get_model(args.model_type) -from deeplabcut.pose_estimation_pytorch.apis import superanimal_analyze_images +try: + from deeplabcut.pose_estimation_pytorch.apis import ( # pyright: ignore[reportMissingImports] + superanimal_analyze_images, + ) +except ImportError: + raise ImportError( + "DeepLabCut is required for the animal demo. " + "Install it with: pip install \"fmpose3d[animals]\"" + ) from None superanimal_name = "superanimal_quadruped" model_name = "hrnet_w32" diff --git a/fmpose3d/inference_api/README.md b/fmpose3d/inference_api/README.md index 87f6e58..5a15900 100644 --- a/fmpose3d/inference_api/README.md +++ b/fmpose3d/inference_api/README.md @@ -48,6 +48,12 @@ result = api.predict("dog.jpg") print(result.poses_3d.shape) # (1, 26, 3) ``` +Before using the animal pipeline, install the optional DeepLabCut dependency: + +```bash +pip install "fmpose3d[animals]" +``` + ## API Documentation @@ -221,6 +227,9 @@ Default 2D estimator for the human pipeline. Wraps HRNet + YOLO with a COCO → 2D estimator for the animal pipeline. Uses DeepLabCut SuperAnimal and maps quadruped80K keypoints to the 26-joint Animal3D layout. +If DeepLabCut is not installed, calling this estimator raises a clear `ImportError` +with the recommended install command: `pip install "fmpose3d[animals]"`. + - `setup_runtime()` — No-op (DLC loads lazily). - `predict(frames: ndarray)` → `(keypoints, scores, valid_frames_mask)` — Returns Animal3D-format 2D keypoints plus a frame-level validity mask. diff --git a/fmpose3d/inference_api/fmpose3d.py b/fmpose3d/inference_api/fmpose3d.py index 77d88ef..07a3417 100644 --- a/fmpose3d/inference_api/fmpose3d.py +++ b/fmpose3d/inference_api/fmpose3d.py @@ -189,6 +189,20 @@ def _compute_valid_frames_mask( } +def _require_superanimal_analyze_images() -> Callable[..., object]: + """Return DeepLabCut's SuperAnimal API or raise a clear ImportError.""" + try: + from deeplabcut.pose_estimation_pytorch.apis import ( # pyright: ignore[reportMissingImports] + superanimal_analyze_images, + ) + except ImportError: + raise ImportError( + "DeepLabCut is required for the animal 2D estimator. " + "Install it with: pip install \"fmpose3d[animals]\"" + ) from None + return superanimal_analyze_images + + class SuperAnimalEstimator: """2D pose estimator for animals: DeepLabCut SuperAnimal. @@ -236,9 +250,7 @@ def predict( """ import cv2 import tempfile - from deeplabcut.pose_estimation_pytorch.apis import ( - superanimal_analyze_images, - ) + superanimal_analyze_images = _require_superanimal_analyze_images() cfg = self.cfg num_frames = frames.shape[0] diff --git a/pyproject.toml b/pyproject.toml index 3e7c8b3..b40c122 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,6 @@ dependencies = [ "scikit-image>=0.19.0", "filterpy>=1.4.5", "pandas>=1.0.1", - "deeplabcut==3.0.0rc13", "huggingface_hub>=0.20.0", ] @@ -45,6 +44,7 @@ dependencies = [ dev = ["pytest>=7.0.0", "black>=22.0.0", "flake8>=4.0.0", "isort>=5.10.0"] wandb = ["wandb>=0.12.0"] viz = ["matplotlib>=3.5.0", "opencv-python>=4.5.0"] +animals = ["deeplabcut>=3.0.0rc13"] [tool.setuptools] include-package-data = true diff --git a/tests/fmpose3d_api/conftest.py b/tests/fmpose3d_api/conftest.py index 9f824f1..872509a 100644 --- a/tests/fmpose3d_api/conftest.py +++ b/tests/fmpose3d_api/conftest.py @@ -23,6 +23,7 @@ import os import socket +from importlib.util import find_spec import pytest @@ -72,12 +73,7 @@ def weights_ready(filename: str) -> bool: HUMAN_WEIGHTS_READY: bool = weights_ready(HUMAN_WEIGHTS_FILENAME) ANIMAL_WEIGHTS_READY: bool = weights_ready(ANIMAL_WEIGHTS_FILENAME) -try: - import deeplabcut # noqa: F401 - - DLC_AVAILABLE: bool = True -except ImportError: - DLC_AVAILABLE = False +DLC_AVAILABLE: bool = find_spec("deeplabcut") is not None # --------------------------------------------------------------------------- # Reusable skip markers diff --git a/tests/fmpose3d_api/test_fmpose3d.py b/tests/fmpose3d_api/test_fmpose3d.py index 1877baa..1e0b29e 100644 --- a/tests/fmpose3d_api/test_fmpose3d.py +++ b/tests/fmpose3d_api/test_fmpose3d.py @@ -741,6 +741,18 @@ def test_pose3d_result(self): class TestSuperAnimalPrediction: + def test_predict_raises_clear_error_without_deeplabcut(self): + """Missing DLC should raise a clear installation hint.""" + estimator = SuperAnimalEstimator() + frames = np.random.randint(0, 255, (1, 64, 64, 3), dtype=np.uint8) + + with patch( + "fmpose3d.inference_api.fmpose3d.importlib.util.find_spec", + return_value=None, + ): + with pytest.raises(ImportError, match=r"fmpose3d\[animals\]"): + estimator.predict(frames) + def test_predict_returns_zeros_when_no_bodyparts(self): """When DLC detects nothing, keypoints are zero-filled.""" pytest.importorskip("deeplabcut")