Merge pull request #285 from KumarLabJax/task/add-mobilenet-unet-model

Add implementation for mobilnet unet model to jabs-vision

Author: bergsalex

.github/workflows/_run-tests-action.yml

@@ -37,25 +37,18 @@ jobs:
       # Run tests - discovers all test directories automatically
       - name: Test with pytest
         run: |
-          # Collect all test directories
-          TEST_DIRS=""
-          
-          # Root tests directory
+          set -euo pipefail
+
+          # Root tests
           if [ -d "tests" ]; then
-            TEST_DIRS="tests"
+            echo "Running root tests"
+            uv run pytest tests
           fi
-          
-          # Package-level test directories
-          for pkg_tests in packages/*/tests; do
-            if [ -d "$pkg_tests" ]; then
-              TEST_DIRS="$TEST_DIRS $pkg_tests"
+
+          # Package tests: run each with that package as a root on sys.path
+          for pkg in packages/*; do
+            if [ -d "$pkg/tests" ]; then
+              echo "Running tests in $pkg/tests"
+              uv run pytest "$pkg/tests"
             fi
           done
-          
-          if [ -n "$TEST_DIRS" ]; then
-            echo "Running tests in: $TEST_DIRS"
-            uv run pytest $TEST_DIRS
-          else
-            echo "No test directories found"
-            exit 1
-          fi

docs/api-reference/jabs-vision/backbones.md

@@ -0,0 +1,5 @@
+---
+title: JABS Vision Backbones
+---
+
+::: jabs.vision.backbones

docs/api-reference/jabs-vision/index.md

@@ -0,0 +1,30 @@
+# JABS Vision (`jabs-vision`)
+
+This package handles raw video processing and deep learning training and inference.
+
+## Overview
+
+`jabs-vision` is responsible for converting raw video frames into pose estimation data.
+It houses the heavy machine learning frameworks and GPU-accelerated code.
+
+## Responsibilities
+
+- **Pose Estimation Inference**: Running deep learning models (PyTorch) on video frames
+  to detect keypoints.
+- **Static Object Detection**
+- **Segmentation Masking**
+- **Identity Matching**: Tracking individual animals across frames using vectorized
+  features.
+- **Video Processing**: Handling frame extraction and normalization for input to vision
+  models.
+
+## Key Components
+
+- `jabs.vision.inference`: Wrappers for pose estimation model execution.
+- `jabs.vision.tracking`: Identity matching and re-identification logic.
+
+## Dependencies
+
+- `torch` / `torchvision`
+- `opencv-python-headless`
+- `jabs-io`, `jabs-core`

docs/api-reference/jabs-vision/modules.md

@@ -0,0 +1,7 @@
+---
+title: JABS Vision Modules
+---
+
+::: jabs.vision.modules
+    options:
+        show_submodules: true

packages/jabs-vision/src/jabs/vision/backbones/__init__.py

@@ -0,0 +1,8 @@
+"""Backbone modules for jabs-vision models."""
+
+from .timm import TimmBackbone, TimmBackboneConfig
+
+__all__ = [
+    "TimmBackbone",
+    "TimmBackboneConfig",
+]

packages/jabs-vision/src/jabs/vision/backbones/timm.py

@@ -0,0 +1,87 @@
+"""Timm backbone wrapper for feature extraction."""
+
+from dataclasses import dataclass, field
+
+import torch.nn as nn
+from torch import Tensor
+
+try:
+    import timm
+except ImportError:
+    timm = None  # type: ignore[assignment]
+
+
+@dataclass
+class TimmBackboneConfig:
+    """Configuration for TimmBackbone.
+
+    Attributes:
+        name: Name of the timm model to use.
+        pretrained: Whether to load pretrained weights.
+        out_indices: Which feature stages to return (0=stem, 1-4=stages).
+    """
+
+    name: str = "mobilenetv3_large_100"
+    pretrained: bool = True
+    out_indices: tuple[int, ...] = field(default_factory=lambda: (0, 1, 2, 3, 4))
+
+
+class TimmBackbone(nn.Module):
+    """Wrapper around timm models for multi-scale feature extraction.
+
+    This module creates a feature extractor from any timm model that supports
+    `features_only=True`, returning feature maps at multiple scales.
+
+    Example:
+        ``` py
+        cfg = TimmBackboneConfig(name="mobilenetv3_small_100", pretrained=True)
+        backbone = TimmBackbone(cfg)
+        x = torch.randn(1, 3, 256, 256)
+        features = backbone(x)  # List of feature tensors
+        print([f.shape for f in features])
+        ```
+    """
+
+    def __init__(self, cfg: TimmBackboneConfig) -> None:
+        super().__init__()
+
+        if timm is None:
+            raise ImportError(
+                "timm is required for TimmBackbone. "
+                "Install it with: pip install 'jabs-vision[timm]' or pip install timm"
+            )
+
+        self.cfg = cfg
+
+        # Create feature extractor
+        self.model = timm.create_model(
+            cfg.name,
+            pretrained=cfg.pretrained,
+            features_only=True,
+            out_indices=cfg.out_indices,
+        )
+
+        # Get feature info for channels and strides
+        self._channels = [info["num_chs"] for info in self.model.feature_info]
+        self._strides = [info["reduction"] for info in self.model.feature_info]
+
+    @property
+    def channels(self) -> list[int]:
+        """Number of channels at each feature level."""
+        return self._channels
+
+    @property
+    def strides(self) -> list[int]:
+        """Spatial reduction (stride) at each feature level."""
+        return self._strides
+
+    def forward(self, x: Tensor) -> list[Tensor]:
+        """Extract multi-scale features.
+
+        Args:
+            x: Input tensor of shape (B, C, H, W).
+
+        Returns:
+            List of feature tensors, one per output index.
+        """
+        return self.model(x)

packages/jabs-vision/src/jabs/vision/core/__init__.py

@@ -0,0 +1,14 @@
+"""Core components for jabs-vision."""
+
+from .interfaces import (
+    BaseVisionModel,
+    OutputKeys,
+)
+from .registry import MODEL_REGISTRY, ModelRegistry
+
+__all__ = [
+    "MODEL_REGISTRY",
+    "BaseVisionModel",
+    "ModelRegistry",
+    "OutputKeys",
+]

packages/jabs-vision/src/jabs/vision/core/interfaces.py

@@ -0,0 +1,40 @@
+"""Core interfaces and protocols for jabs-vision models."""
+
+from abc import ABC, abstractmethod
+
+import torch
+from torch import Tensor
+
+# =============================================================================
+# Abstract Base Classes (for implementation guidance)
+# =============================================================================
+
+
+class BaseVisionModel(ABC, torch.nn.Module):
+    """Abstract base class for vision models.
+
+    Provides common functionality and enforces the interface contract.
+    """
+
+    @abstractmethod
+    def forward(self, x: Tensor) -> dict[str, Tensor]:
+        """Forward pass returning named outputs."""
+        pass
+
+
+# =============================================================================
+# Output Key Conventions
+# =============================================================================
+
+
+class OutputKeys:
+    """Standard keys for model output dictionaries.
+
+    Using these constants prevents typos and documents the contract.
+    """
+
+    # Keypoint outputs
+    HEATMAPS = "heatmaps"
+    COORDS = "coords"
+    KEYPOINT_CONFIDENCE = "keypoint_confidence"
+    CONFIDENCE_MAPS = "confidence_maps"  # Learned confidence (not derived from heatmaps)

packages/jabs-vision/src/jabs/vision/core/registry.py

@@ -0,0 +1,82 @@
+"""Model registry for jabs-vision.
+
+Provides a centralized registry for model classes with decorator-based registration.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TypeVar
+
+T = TypeVar("T")
+
+
+class ModelRegistry:
+    """Registry for model classes.
+
+    Usage:
+        @MODEL_REGISTRY.register("my_model")
+        class MyModel:
+            ...
+
+        # Later:
+        model_cls = MODEL_REGISTRY.get("my_model")
+    """
+
+    def __init__(self) -> None:
+        self._registry: dict[str, type] = {}
+
+    def register(self, name: str) -> Callable[[T], T]:
+        """Decorator to register a model class.
+
+        Args:
+            name: Unique name for the model.
+
+        Returns:
+            Decorator function that registers the class.
+
+        Raises:
+            ValueError: If a model with this name is already registered.
+        """
+
+        def _register(cls: T) -> T:
+            if name in self._registry:
+                raise ValueError(f"Model '{name}' is already registered.")
+            self._registry[name] = cls  # type: ignore[assignment]
+            return cls
+
+        return _register
+
+    def get(self, name: str) -> type:
+        """Retrieve a model class from the registry.
+
+        Args:
+            name: Name of the registered model.
+
+        Returns:
+            The registered model class.
+
+        Raises:
+            KeyError: If the model name is not found.
+        """
+        if name not in self._registry:
+            raise KeyError(
+                f"Model '{name}' not found in registry. "
+                f"Registered models: {list(self._registry.keys())}"
+            )
+        return self._registry[name]
+
+    def list_models(self) -> list[str]:
+        """List all registered model names.
+
+        Returns:
+            Sorted list of registered model names.
+        """
+        return sorted(self._registry.keys())
+
+    def __contains__(self, name: str) -> bool:
+        """Check if a model name is registered."""
+        return name in self._registry
+
+
+MODEL_REGISTRY = ModelRegistry()

packages/jabs-vision/src/jabs/vision/models/__init__.py

@@ -0,0 +1 @@
+"""Model definitions for JabsVision."""