Add general detector + global moth model to example pipeline

processing_services/README.md

@@ -43,6 +43,7 @@ If your goal is to run an ML backend locally, simply copy the `example` app and
         - `ZeroShotObjectDetectorPipeline`
         - `ZeroShotObjectDetectorWithRandomSpeciesClassifierPipeline`
         - `ZeroShotObjectDetectorWithConstantClassifierPipeline`
+        - `ZeroShotObjectDetectorWithGlobalMothClassifierPipeline` (see below)
 
 4. Add `NewPipeline` to `processing_services/example/api/api.py`
 
@@ -76,6 +77,15 @@ PipelineChoice = typing.Literal[
     "new-pipeline-slug",
 ]
 ```
+### Global moth classifier pipeline
+
+`ZeroShotObjectDetectorWithGlobalMothClassifierPipeline` finds insects with the zero-shot detector and classifies each crop with a ResNet50 trained on about 29,000 moth species (`api/global_moth_classifier.py`). It is a working template for adding a real model with downloaded weights.
+
+- The weights (about 330 MB) and the category map are downloaded on first use into `processing_services/example/cache/` and reused afterwards. The first request is slow on CPU while the models load.
+- Set the `MODEL_BASE_URL` environment variable to download from a mirror, or to read from a local directory laid out as `moths/classification/<file>`. It defaults to the public AMI model store.
+- The classifier returns the raw logits for every class (in the order of the category map published in `/info`) and no probabilities. The receiving system derives probabilities from the logits.
+- Run its tests without any downloads: `docker compose -f processing_services/docker-compose.yml run --rm ml_backend_example python -m unittest api.test_global_moth_classifier api.test_utils`
+
 ## Demo
 
 ## `minimal` Pipelines and Output Images

processing_services/docker-compose.yml

@@ -16,6 +16,8 @@ services:
       context: ./example
     volumes:
       - ./example/:/app:z
+    environment:
+      - MODEL_BASE_URL  # Optional: mirror or local directory for model weights
     ports:
       - "2003:2000"
     extra_hosts:

processing_services/example/.gitignore

@@ -0,0 +1,4 @@
+# Cache directories for models and dependencies
+cache/
+huggingface_cache/
+pytorch_cache/

processing_services/example/api/algorithms.py

@@ -3,8 +3,6 @@
 import math
 import random
 
-import torch
-
 from .schemas import (
     AlgorithmCategoryMapResponse,
     AlgorithmConfigResponse,
@@ -14,25 +12,14 @@
     Detection,
     SourceImage,
 )
+from .utils import get_best_device
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
 SAVED_MODELS = {}
 
 
-def get_best_device() -> str:
-    """
-    Returns the best available device for running the model.
-
-    MPS is not supported by the current algoritms.
-    """
-    if torch.cuda.is_available():
-        return f"cuda:{torch.cuda.current_device()}"
-    else:
-        return "cpu"
-
-
 class Algorithm:
     algorithm_config_response: AlgorithmConfigResponse
 

processing_services/example/api/api.py

@@ -11,6 +11,7 @@
     ZeroShotHFClassifierPipeline,
     ZeroShotObjectDetectorPipeline,
     ZeroShotObjectDetectorWithConstantClassifierPipeline,
+    ZeroShotObjectDetectorWithGlobalMothClassifierPipeline,
     ZeroShotObjectDetectorWithRandomSpeciesClassifierPipeline,
 )
 from .schemas import (
@@ -41,6 +42,7 @@
     ZeroShotObjectDetectorPipeline,
     ZeroShotObjectDetectorWithConstantClassifierPipeline,
     ZeroShotObjectDetectorWithRandomSpeciesClassifierPipeline,
+    ZeroShotObjectDetectorWithGlobalMothClassifierPipeline,
 ]
 pipeline_choices: dict[str, type[Pipeline]] = {pipeline.config.slug: pipeline for pipeline in pipelines}
 algorithm_choices: dict[str, AlgorithmConfigResponse] = {
@@ -62,7 +64,7 @@ async def info() -> ProcessingServiceInfoResponse:
     info = ProcessingServiceInfoResponse(
         name="Custom ML Backend",
         description=("A template for running custom models locally."),
-        pipelines=[pipeline.config for pipeline in pipelines],
+        pipelines=[pipeline.get_config() for pipeline in pipelines],
         # algorithms=list(algorithm_choices.values()),
     )
     return info

processing_services/example/api/base.py

@@ -0,0 +1,233 @@
+"""
+Simplified base classes for inference models without database dependencies.
+Adapted from trapdata.ml.models.base but streamlined for processing service use.
+"""
+
+import json
+import logging
+from typing import Any
+
+import torch
+import torchvision.transforms
+
+from .utils import get_best_device, get_or_download_file
+
+logger = logging.getLogger(__name__)
+
+
+# Standard normalization transforms
+imagenet_normalization = torchvision.transforms.Normalize(
+    mean=[0.485, 0.456, 0.406],  # RGB
+    std=[0.229, 0.224, 0.225],  # RGB
+)
+
+tensorflow_normalization = torchvision.transforms.Normalize(
+    mean=[0.5, 0.5, 0.5],  # RGB
+    std=[0.5, 0.5, 0.5],  # RGB
+)
+
+
+class SimplifiedInferenceBase:
+    """
+    Simplified base class for inference models without database or queue dependencies.
+    """
+
+    name: str = "Unknown Inference Model"
+    description: str = ""
+    weights_path: str | None = None
+    labels_path: str | None = None
+    category_map: dict[int, str] = {}
+    num_classes: int | None = None
+    default_taxon_rank: str = "SPECIES"
+    normalization = tensorflow_normalization
+    batch_size: int = 4
+    device: str | None = None
+
+    def __init__(self, **kwargs):
+        # Override any class attributes with provided kwargs
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+        logger.info(f"Initializing simplified inference class {self.name}")
+
+        self.device = self.device or get_best_device()
+        self.category_map = self.get_labels(self.labels_path)
+        self.num_classes = self.num_classes or len(self.category_map)
+        self.weights = self.get_weights(self.weights_path)
+        self.transforms = self.get_transforms()
+
+        logger.info(f"Loading model for {self.name} with {len(self.category_map or [])} categories")
+        self.model = self.get_model()
+
+    @classmethod
+    def get_key(cls) -> str:
+        """Generate a unique key for this algorithm."""
+        if hasattr(cls, "key") and cls.key:
+            return cls.key
+        else:
+            return cls.name.lower().replace(" ", "-").replace("/", "-")
+
+    def get_weights(self, weights_path: str | None) -> str | None:
+        """Download and cache model weights."""
+        if weights_path:
+            logger.info(f"⬇️  Downloading model weights from: {weights_path}")
+            weights_file = str(get_or_download_file(weights_path, tempdir_prefix="models"))
+            logger.info(f"✅ Model weights downloaded to: {weights_file}")
+            return weights_file
+        else:
+            logger.warning(f"No weights specified for model {self.name}")
+            return None
+
+    def get_labels(self, labels_path: str | None) -> dict[int, str]:
+        """Download and load category labels."""
+        if not labels_path:
+            return {}
+
+        logger.info(f"⬇️  Downloading category labels from: {labels_path}")
+        local_path = get_or_download_file(labels_path, tempdir_prefix="models")
+        logger.info(f"📝 Loading category labels from: {local_path}")
+
+        with open(local_path) as f:
+            labels = json.load(f)
+
+        # Convert label->index mapping to index->label mapping
+        index_to_label = {index: label for label, index in labels.items()}
+        if len(index_to_label) != len(labels) or sorted(index_to_label) != list(range(len(labels))):
+            raise ValueError(
+                f"Category labels in {local_path} must map each label to a unique index from 0 to {len(labels) - 1}"
+            )
+        logger.info(f"✅ Loaded {len(index_to_label)} category labels")
+        return index_to_label
+
+    def get_transforms(self) -> torchvision.transforms.Compose:
+        """Get image preprocessing transforms."""
+        return torchvision.transforms.Compose(
+            [
+                torchvision.transforms.Resize((224, 224)),
+                torchvision.transforms.ToTensor(),
+                self.normalization,
+            ]
+        )
+
+    def get_model(self) -> torch.nn.Module:
+        """
+        Load and return the PyTorch model.
+        Must be implemented by subclasses.
+        """
+        raise NotImplementedError("Subclasses must implement get_model()")
+
+    def predict_batch(self, batch: torch.Tensor) -> torch.Tensor:
+        """
+        Run inference on a batch of images.
+        Must be implemented by subclasses.
+        """
+        raise NotImplementedError("Subclasses must implement predict_batch()")
+
+    def post_process_batch(self, logits: torch.Tensor) -> Any:
+        """
+        Post-process model outputs.
+        Must be implemented by subclasses.
+        """
+        raise NotImplementedError("Subclasses must implement post_process_batch()")
+
+
+class ResNet50Base(SimplifiedInferenceBase):
+    """
+    Base class for ResNet50-based models.
+    """
+
+    input_size: int = 224
+    normalization = imagenet_normalization
+
+    def get_transforms(self) -> torchvision.transforms.Compose:
+        """Get ResNet50-specific transforms."""
+        return torchvision.transforms.Compose(
+            [
+                torchvision.transforms.Resize((self.input_size, self.input_size)),
+                torchvision.transforms.ToTensor(),
+                self.normalization,
+            ]
+        )
+
+    def get_model(self) -> torch.nn.Module:
+        """Load ResNet50 model with custom classifier."""
+        import torchvision.models as models
+
+        logger.info("🏗️  Creating ResNet50 model architecture...")
+        # Create ResNet50 backbone
+        model = models.resnet50(weights=None)
+
+        # Replace final classifier layer
+        if self.num_classes is None:
+            raise ValueError("num_classes must be set before loading model")
+        logger.info(f"🔧 Setting up classifier layer for {self.num_classes} classes...")
+        model.fc = torch.nn.Linear(model.fc.in_features, self.num_classes)
+
+        # Load pretrained weights
+        if self.weights:
+            logger.info(f"📂 Loading pretrained weights from: {self.weights}")
+            checkpoint = torch.load(self.weights, map_location=self.device)
+
+            # Handle different checkpoint formats
+            if "model_state_dict" in checkpoint:
+                logger.info("📥 Loading state dict from 'model_state_dict' key...")
+                model.load_state_dict(checkpoint["model_state_dict"])
+            elif "state_dict" in checkpoint:
+                logger.info("📥 Loading state dict from 'state_dict' key...")
+                model.load_state_dict(checkpoint["state_dict"])
+            else:
+                logger.info("📥 Loading state dict directly...")
+                model.load_state_dict(checkpoint)
+            logger.info("✅ Model weights loaded successfully!")
+        else:
+            logger.warning("⚠️  No pretrained weights provided - using random initialization")
+
+        logger.info(f"📱 Moving model to device: {self.device}")
+        model = model.to(self.device)
+        model.eval()
+        logger.info("✅ Model ready for inference!")
+        return model
+
+    def predict_batch(self, batch: torch.Tensor) -> torch.Tensor:
+        """Run inference on batch."""
+        with torch.no_grad():
+            batch = batch.to(self.device)
+            outputs = self.model(batch)
+            return outputs
+
+    def post_process_batch(self, logits: torch.Tensor) -> list:
+        """Convert logits to predictions."""
+        probabilities = torch.softmax(logits, dim=1)
+        return [
+            {"scores": probs.cpu().numpy().tolist(), "logits": row.cpu().numpy().tolist()}
+            for probs, row in zip(probabilities, logits, strict=True)
+        ]
+
+
+class TimmResNet50Base(ResNet50Base):
+    """
+    Base class for timm ResNet50-based models.
+    """
+
+    def get_model(self) -> torch.nn.Module:
+        """Load timm ResNet50 model."""
+        import timm
+
+        # Create timm ResNet50 model
+        model = timm.create_model("resnet50", pretrained=False, num_classes=self.num_classes)
+
+        # Load pretrained weights
+        if self.weights:
+            checkpoint = torch.load(self.weights, map_location=self.device)
+
+            # Handle different checkpoint formats
+            if "model_state_dict" in checkpoint:
+                model.load_state_dict(checkpoint["model_state_dict"])
+            elif "state_dict" in checkpoint:
+                model.load_state_dict(checkpoint["state_dict"])
+            else:
+                model.load_state_dict(checkpoint)
+
+        model = model.to(self.device)
+        model.eval()
+        return model