diff --git a/processing_services/README.md b/processing_services/README.md index 48bb254e7..fd4bba14c 100644 --- a/processing_services/README.md +++ b/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/`. 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 diff --git a/processing_services/docker-compose.yml b/processing_services/docker-compose.yml index 91a21c100..8c9dc3ef4 100644 --- a/processing_services/docker-compose.yml +++ b/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: diff --git a/processing_services/example/.gitignore b/processing_services/example/.gitignore new file mode 100644 index 000000000..95312cdaf --- /dev/null +++ b/processing_services/example/.gitignore @@ -0,0 +1,4 @@ +# Cache directories for models and dependencies +cache/ +huggingface_cache/ +pytorch_cache/ diff --git a/processing_services/example/api/algorithms.py b/processing_services/example/api/algorithms.py index 8a80038dd..7d1b93980 100644 --- a/processing_services/example/api/algorithms.py +++ b/processing_services/example/api/algorithms.py @@ -3,8 +3,6 @@ import math import random -import torch - from .schemas import ( AlgorithmCategoryMapResponse, AlgorithmConfigResponse, @@ -14,6 +12,7 @@ Detection, SourceImage, ) +from .utils import get_best_device logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG) @@ -21,18 +20,6 @@ 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 diff --git a/processing_services/example/api/api.py b/processing_services/example/api/api.py index 79ce5d83c..0bfc42040 100644 --- a/processing_services/example/api/api.py +++ b/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 diff --git a/processing_services/example/api/base.py b/processing_services/example/api/base.py new file mode 100644 index 000000000..8bf6718b2 --- /dev/null +++ b/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 diff --git a/processing_services/example/api/global_moth_classifier.py b/processing_services/example/api/global_moth_classifier.py new file mode 100644 index 000000000..4619d2d80 --- /dev/null +++ b/processing_services/example/api/global_moth_classifier.py @@ -0,0 +1,215 @@ +""" +Global moth species classifier for the example processing service. + +Adapted from the classifier of the same name in the AMI data companion (trapdata), without its database or queue +dependencies. It classifies cropped detections into ~29k species and returns the raw logits for every class; the +caller derives probabilities from them. +""" + +import datetime +import logging +import os +import time + +import requests +import torch +import torchvision.transforms + +from .algorithms import SAVED_MODELS, Algorithm +from .base import TimmResNet50Base, imagenet_normalization +from .schemas import ( + AlgorithmCategoryMapResponse, + AlgorithmConfigResponse, + AlgorithmReference, + ClassificationResponse, + Detection, +) + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + +# Public object store holding the model weights and category maps. Override with the MODEL_BASE_URL environment +# variable to use a mirror or a local directory laid out the same way. +DEFAULT_MODEL_BASE_URL = ( + "https://object-arbutus.alliancecan.ca/swift/v1/AUTH_3c987b8fc90743469d42899b1fdb48eb/ami-models/" +) +MODEL_BASE_URL_ENV_VAR = "MODEL_BASE_URL" + +WEIGHTS_RELATIVE_PATH = "moths/classification/global_resnet50_20240828_b06d3b3a.pth" +LABELS_RELATIVE_PATH = "moths/classification/global_category_map_with_names_20240828.json" + +# Decimal places kept for each logit. The model has ~29k classes, so every digit multiplies the response size. +LOGIT_DECIMALS = 4 + +# Category maps by labels path, so repeated `/info` calls do not re-read and re-parse the labels file. +_CATEGORY_MAPS: dict[str, dict[int, str]] = {} + + +def get_model_url(relative_path: str) -> str: + """Join a model file path onto MODEL_BASE_URL (or the default public store).""" + base_url = os.environ.get(MODEL_BASE_URL_ENV_VAR) or DEFAULT_MODEL_BASE_URL + return f"{base_url.rstrip('/')}/{relative_path}" + + +class GlobalMothClassifier(Algorithm, TimmResNet50Base): + """ + Classifies cropped detections into one of the species in the global category map. + + `Algorithm` provides the interface the pipelines use, and `TimmResNet50Base` loads the weights and labels. The + weights are only downloaded when `compile()` runs, so creating the object and describing it in `/info` is cheap. + The loaded model is shared between instances, so later requests do not reload it. + """ + + name = "Global Species Classifier - Aug 2024" + description = ( + "Trained on August 28th, 2024 for 29,176 species. " + "https://wandb.ai/moth-ai/global-moth-classifier/runs/h0cuqrbc/overview" + ) + + input_size = 128 + normalization = imagenet_normalization + default_taxon_rank = "SPECIES" + batch_size = 4 + + def __init__(self, **kwargs): + self._init_kwargs = { + "weights_path": get_model_url(WEIGHTS_RELATIVE_PATH), + "labels_path": get_model_url(LABELS_RELATIVE_PATH), + **kwargs, + } + self.weights_path = self._init_kwargs["weights_path"] + self.labels_path = self._init_kwargs["labels_path"] + self.model = None + self.transforms = None + self.category_map = {} + self.num_classes = None + self._algorithm_config_response: AlgorithmConfigResponse | None = None + + @property + def algorithm_config_response(self) -> AlgorithmConfigResponse: + """The algorithm description, including the category map once the labels can be loaded.""" + if self._algorithm_config_response is not None: + return self._algorithm_config_response + config = self.get_algorithm_config_response() + # Do not remember a description built without labels, so a later call can retry the download. + if self.category_map: + self._algorithm_config_response = config + return config + + def _saved_model_key(self) -> str: + return f"{type(self).__name__}:{self.weights_path}:{self.labels_path}:{self._init_kwargs.get('device')}" + + def compile(self): + """Download (if needed) and load the weights and labels. Reuses an already loaded model.""" + if self.model is not None: + return + + saved_key = self._saved_model_key() + if saved_key in SAVED_MODELS: + logger.info(f"Using saved model for {self.name}") + self.__dict__.update(SAVED_MODELS[saved_key]) + return + + logger.info(f"Compiling {self.name} from {self.weights_path}") + TimmResNet50Base.__init__(self, **self._init_kwargs) + SAVED_MODELS[saved_key] = { + name: getattr(self, name) + for name in ("model", "transforms", "category_map", "num_classes", "weights", "device") + } + logger.info(f"Compiled {self.name}: {self.num_classes} classes on {self.device}") + + def get_transforms(self) -> torchvision.transforms.Compose: + return torchvision.transforms.Compose( + [ + torchvision.transforms.Resize((self.input_size, self.input_size)), + torchvision.transforms.ToTensor(), + self.normalization, + ] + ) + + def post_process_batch(self, logits: torch.Tensor) -> list[dict]: + """Return the logits for every class, rounded to keep the response small. Probabilities are not computed.""" + return [{"logits": [round(value, LOGIT_DECIMALS) for value in row.tolist()]} for row in logits.cpu()] + + def run(self, detections: list[Detection]) -> list[Detection]: + """ + Classify detections from their cropped images. + + Detections without a cropped image are skipped with a warning and are not returned. + """ + if not detections: + return [] + if self.model is None or self.transforms is None: + raise RuntimeError("Model not compiled. Call compile() first.") + + logger.info(f"Running {self.name} on {len(detections)} detections") + classified_detections: list[Detection] = [] + + for i in range(0, len(detections), self.batch_size): + # Only detections that have an image are batched; predictions are paired with this list, not the input. + batched_detections: list[Detection] = [] + batch_images = [] + for detection in detections[i : i + self.batch_size]: + if detection._pil is None: + logger.warning(f"Detection {detection.id} has no cropped image, skipping it") + continue + batched_detections.append(detection) + batch_images.append(detection._pil.convert("RGB")) + + if not batched_detections: + continue + + batch_tensor = torch.stack([self.transforms(image) for image in batch_images]) + start_time = time.perf_counter() + predictions = self.post_process_batch(self.predict_batch(batch_tensor)) + inference_time = (time.perf_counter() - start_time) / len(batched_detections) + + for detection, prediction in zip(batched_detections, predictions, strict=True): + logits = prediction["logits"] + best_index = max(range(len(logits)), key=logits.__getitem__) + classification = ClassificationResponse( + classification=self.category_map[best_index], + logits=logits, + inference_time=inference_time, + timestamp=datetime.datetime.now(), + algorithm=AlgorithmReference(name=self.name, key=self.get_key()), + terminal=True, + ) + classified_detections.append(detection.model_copy(update={"classifications": [classification]})) + + logger.info(f"Classified {len(classified_detections)} of {len(detections)} detections") + return classified_detections + + def get_category_map(self) -> AlgorithmCategoryMapResponse: + if not self.category_map and self.labels_path in _CATEGORY_MAPS: + self.category_map = _CATEGORY_MAPS[self.labels_path] + if not self.category_map: + try: + self.category_map = self.get_labels(self.labels_path) + _CATEGORY_MAPS[self.labels_path] = self.category_map + except (requests.RequestException, OSError, ValueError) as error: + logger.warning(f"Could not load the category labels from {self.labels_path}: {error}") + + data = [ + {"index": index, "label": label, "taxon_rank": self.default_taxon_rank} + for index, label in sorted(self.category_map.items()) + ] + return AlgorithmCategoryMapResponse( + data=data, + labels=[row["label"] for row in data], + version="v1", + description=f"Global moth species classifier with {len(data)} species", + uri=self.labels_path, + ) + + def get_algorithm_config_response(self) -> AlgorithmConfigResponse: + return AlgorithmConfigResponse( + name=self.name, + key=self.get_key(), + task_type="classification", + description=self.description, + version=1, + version_name="v1", + category_map=self.get_category_map(), + uri=self.weights_path, + ) diff --git a/processing_services/example/api/pipelines.py b/processing_services/example/api/pipelines.py index 02b31d0d9..907cb8488 100644 --- a/processing_services/example/api/pipelines.py +++ b/processing_services/example/api/pipelines.py @@ -9,6 +9,7 @@ RandomSpeciesClassifier, ZeroShotObjectDetector, ) +from .global_moth_classifier import GlobalMothClassifier from .schemas import ( Detection, DetectionResponse, @@ -70,6 +71,14 @@ def __init__( self.batch_sizes = custom_batch_sizes or self.batch_sizes or [1] * len(self.stages) assert len(self.batch_sizes) == len(self.stages), "Number of batch sizes must match the number of stages." + @classmethod + def get_config(cls) -> PipelineConfigResponse: + """ + The pipeline description published in `/info`. + Override when part of it, such as a category map, is only available at runtime. + """ + return cls.config + def get_stages(self) -> list[Algorithm]: """ An optional function to initialize and return a list of algorithms/stages. @@ -346,3 +355,73 @@ def run(self) -> PipelineResultsResponse: logger.info(f"Successfully processed {len(detections_with_classifications)} detections.") return pipeline_response + + +class ZeroShotObjectDetectorWithGlobalMothClassifierPipeline(Pipeline): + """ + A pipeline that uses the HuggingFace zero shot object detector and the global moth classifier. + This provides high-quality moth species identification with 29,176+ species support. + """ + + batch_sizes = [1, 4] # Detector batch=1, Classifier batch=4 + config = PipelineConfigResponse( + name="Zero Shot Object Detector With Global Moth Classifier Pipeline", + slug="zero-shot-object-detector-with-global-moth-classifier-pipeline", + description=( + "HF zero shot object detector with global moth species classifier. " + "Supports 29,176+ moth species trained on global data." + ), + version=1, + algorithms=[], # Populated by get_config() and get_stages(), which need the downloaded category map + ) + + @classmethod + def get_config(cls) -> PipelineConfigResponse: + # The classifier's category map is read from its labels file, so it is not built at import time. + cls.config.algorithms = [ + ZeroShotObjectDetector().algorithm_config_response, + GlobalMothClassifier().algorithm_config_response, + ] + return cls.config + + def get_stages(self) -> list[Algorithm]: + zero_shot_object_detector = ZeroShotObjectDetector() + if "candidate_labels" in self.request_config: + zero_shot_object_detector.candidate_labels = self.request_config["candidate_labels"] + + global_moth_classifier = GlobalMothClassifier() + + self.config.algorithms = [ + zero_shot_object_detector.algorithm_config_response, + global_moth_classifier.algorithm_config_response, + ] + + return [zero_shot_object_detector, global_moth_classifier] + + def run(self) -> PipelineResultsResponse: + start_time = datetime.datetime.now() + detections: list[Detection] = [] + + if self.existing_detections: + logger.info("[1/2] Skipping the localizer, use existing detections...") + detections = self.existing_detections + else: + logger.info("[1/2] No existing detections, generating detections...") + detections = self._get_detections(self.stages[0], self.source_images, self.batch_sizes[0]) + + logger.info("[2/2] Running the global moth classifier...") + detections_with_classifications: list[Detection] = self._get_detections( + self.stages[1], detections, self.batch_sizes[1] + ) + + end_time = datetime.datetime.now() + elapsed_time = (end_time - start_time).total_seconds() + + pipeline_response: PipelineResultsResponse = self._get_pipeline_response( + detections_with_classifications, elapsed_time + ) + logger.info( + f"Successfully processed {len(detections_with_classifications)} detections with global moth classifier." + ) + + return pipeline_response diff --git a/processing_services/example/api/schemas.py b/processing_services/example/api/schemas.py index 05682c6f1..05966126b 100644 --- a/processing_services/example/api/schemas.py +++ b/processing_services/example/api/schemas.py @@ -89,11 +89,17 @@ class ClassificationResponse(pydantic.BaseModel): ) scores: list[float] = pydantic.Field( default_factory=list, - description="The calibrated probabilities for each class label, most commonly the softmax output.", + description=( + "The calibrated probabilities for each class label, most commonly the softmax output. " + "May be left empty when the model returns logits, from which probabilities can be derived." + ), ) logits: list[float] = pydantic.Field( default_factory=list, - description="The raw logits output by the model, before any calibration or normalization.", + description=( + "The raw logits output by the model, before any calibration or normalization, " + "one per class in the order of the algorithm's category map." + ), ) inference_time: float | None = None algorithm: AlgorithmReference @@ -127,7 +133,7 @@ class DetectionRequest(pydantic.BaseModel): class DetectionResponse(pydantic.BaseModel): # these fields are populated with values from a Detection, excluding source_image details source_image_id: str - bbox: BoundingBox + bbox: BoundingBox | None = None inference_time: float | None = None algorithm: AlgorithmReference timestamp: datetime.datetime @@ -213,6 +219,7 @@ class Config: "zero-shot-object-detector-pipeline", "zero-shot-object-detector-with-constant-classifier-pipeline", "zero-shot-object-detector-with-random-species-classifier-pipeline", + "zero-shot-object-detector-with-global-moth-classifier-pipeline", ] diff --git a/processing_services/example/api/test_global_moth_classifier.py b/processing_services/example/api/test_global_moth_classifier.py new file mode 100644 index 000000000..f127cca04 --- /dev/null +++ b/processing_services/example/api/test_global_moth_classifier.py @@ -0,0 +1,148 @@ +""" +Tests for the global moth classifier using a tiny stand-in model and a four-class labels file, +so nothing is downloaded and the tests run in about a second. +""" + +import json +import os +import pathlib +import tempfile +import unittest +from unittest import mock + +import PIL.Image +import torch + +from . import global_moth_classifier +from .algorithms import SAVED_MODELS +from .global_moth_classifier import GlobalMothClassifier, get_model_url +from .schemas import AlgorithmReference, BoundingBox, ClassificationResponse, Detection, DetectionResponse, SourceImage + +LABELS = {"Moth a": 0, "Moth b": 1, "Moth c": 2, "Moth d": 3} + + +class TinyClassifier(GlobalMothClassifier): + """The real classifier with the ResNet replaced by a single linear layer on 8x8 images.""" + + input_size = 8 + + def get_model(self) -> torch.nn.Module: + torch.manual_seed(0) + model = torch.nn.Sequential(torch.nn.Flatten(), torch.nn.Linear(3 * 8 * 8, self.num_classes)) + return model.to(self.device).eval() + + +def make_detection(detection_id: str, color: tuple[int, int, int] | None) -> Detection: + detection = Detection( + id=detection_id, + url="https://example.org/source.jpg", + source_image=SourceImage(id="1", url="https://example.org/source.jpg"), + bbox=BoundingBox(x1=0, y1=0, x2=8, y2=8), + algorithm=AlgorithmReference(name="Detector", key="detector"), + ) + if color is not None: + detection._pil = PIL.Image.new("RGB", (16, 16), color) + return detection + + +class TestGlobalMothClassifier(unittest.TestCase): + def setUp(self): + tmp = tempfile.TemporaryDirectory() + self.addCleanup(tmp.cleanup) + self.labels_path = pathlib.Path(tmp.name) / "labels.json" + self.labels_path.write_text(json.dumps(LABELS)) + self.addCleanup(SAVED_MODELS.clear) + self.addCleanup(global_moth_classifier._CATEGORY_MAPS.clear) + + def make_classifier(self) -> TinyClassifier: + classifier = TinyClassifier(weights_path=None, labels_path=str(self.labels_path), device="cpu") + classifier.compile() + return classifier + + def expected_logits(self, classifier: TinyClassifier, color: tuple[int, int, int]) -> list[float]: + image = PIL.Image.new("RGB", (16, 16), color) + with torch.no_grad(): + return classifier.model(classifier.transforms(image).unsqueeze(0))[0].tolist() + + def test_category_map_is_built_from_the_labels_file(self): + classifier = TinyClassifier(weights_path=None, labels_path=str(self.labels_path)) + category_map = classifier.algorithm_config_response.category_map + + self.assertEqual(category_map.labels, ["Moth a", "Moth b", "Moth c", "Moth d"]) + self.assertEqual([row["index"] for row in category_map.data], [0, 1, 2, 3]) + self.assertEqual(category_map.data[2]["label"], "Moth c") + + def test_class_count_comes_from_the_labels_file(self): + self.assertEqual(self.make_classifier().num_classes, len(LABELS)) + + def test_labels_with_gaps_are_rejected(self): + self.labels_path.write_text(json.dumps({"Moth a": 0, "Moth b": 5})) + classifier = TinyClassifier(weights_path=None, labels_path=str(self.labels_path)) + with self.assertRaises(ValueError): + classifier.compile() + + def test_returns_all_logits_and_the_argmax_label_without_scores(self): + classifier = self.make_classifier() + color = (200, 30, 90) + + (result,) = classifier.run([make_detection("a", color)]) + classification = result.classifications[0] + + expected = self.expected_logits(classifier, color) + self.assertEqual(len(classification.logits), len(LABELS)) + for actual, wanted in zip(classification.logits, expected, strict=True): + self.assertAlmostEqual(actual, wanted, places=3) + self.assertEqual(classification.classification, list(LABELS)[expected.index(max(expected))]) + self.assertEqual(classification.scores, []) + self.assertIsNone(classification.labels) + + def test_detection_without_an_image_does_not_shift_the_other_predictions(self): + classifier = self.make_classifier() + colors = {"a": (250, 0, 0), "c": (0, 0, 250)} + detections = [make_detection("a", colors["a"]), make_detection("b", None), make_detection("c", colors["c"])] + + results = classifier.run(detections) + + self.assertEqual([r.id for r in results], ["a", "c"]) + for result in results: + expected = self.expected_logits(classifier, colors[result.id]) + for actual, wanted in zip(result.classifications[0].logits, expected, strict=True): + self.assertAlmostEqual(actual, wanted, places=3) + + def test_no_detections_returns_nothing(self): + self.assertEqual(self.make_classifier().run([]), []) + + def test_running_before_compile_raises(self): + classifier = TinyClassifier(weights_path=None, labels_path=str(self.labels_path)) + with self.assertRaises(RuntimeError): + classifier.run([make_detection("a", (1, 2, 3))]) + + def test_loaded_model_is_shared_between_instances(self): + first = self.make_classifier() + second = self.make_classifier() + self.assertIs(first.model, second.model) + + def test_model_base_url_can_be_overridden(self): + with mock.patch.dict(os.environ, {"MODEL_BASE_URL": "https://mirror.example.org/models/"}): + self.assertEqual(get_model_url("moths/a.pth"), "https://mirror.example.org/models/moths/a.pth") + with mock.patch.dict(os.environ, {}, clear=False): + os.environ.pop("MODEL_BASE_URL", None) + self.assertTrue(get_model_url("moths/a.pth").startswith(global_moth_classifier.DEFAULT_MODEL_BASE_URL)) + + +class TestResponseSchemas(unittest.TestCase): + def test_logits_only_classification_and_detection_without_bbox_are_valid(self): + algorithm = AlgorithmReference(name="Classifier", key="classifier") + classification = ClassificationResponse( + classification="Moth a", logits=[0.5, -1.0], algorithm=algorithm, timestamp="2026-01-01T00:00:00" + ) + detection = DetectionResponse( + source_image_id="1", algorithm=algorithm, timestamp="2026-01-01T00:00:00", classifications=[classification] + ) + + self.assertIsNone(detection.bbox) + self.assertEqual(detection.classifications[0].scores, []) + + +if __name__ == "__main__": + unittest.main() diff --git a/processing_services/example/api/test_utils.py b/processing_services/example/api/test_utils.py new file mode 100644 index 000000000..82fd3a459 --- /dev/null +++ b/processing_services/example/api/test_utils.py @@ -0,0 +1,76 @@ +import pathlib +import tempfile +import unittest +from unittest import mock + +from . import utils + + +class FakeResponse: + def __init__(self, chunks: list[bytes], content_length: int | None = None, error: Exception | None = None): + self._chunks = chunks + self._error = error + length = sum(len(c) for c in chunks) if content_length is None else content_length + self.headers = {"content-length": str(length)} + + def raise_for_status(self): + pass + + def iter_content(self, chunk_size): + yield from self._chunks + if self._error: + raise self._error + + +class TestGetOrDownloadFile(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + self.cache_root = pathlib.Path(self.tmp.name) / "cache" + patcher = mock.patch.object(utils, "get_cache_root", return_value=self.cache_root) + patcher.start() + self.addCleanup(patcher.stop) + + def test_local_path_is_returned_without_a_request(self): + local_file = pathlib.Path(self.tmp.name) / "weights.pth" + local_file.write_bytes(b"abc") + with mock.patch.object(utils.requests, "get") as get: + self.assertEqual(utils.get_or_download_file(str(local_file)), local_file) + self.assertEqual(utils.get_or_download_file(local_file), local_file) + get.assert_not_called() + + def test_missing_local_path_raises(self): + with self.assertRaises(FileNotFoundError): + utils.get_or_download_file(str(pathlib.Path(self.tmp.name) / "missing.pth")) + + def test_download_uses_a_timeout_and_caches_the_result(self): + url = "https://example.org/models/weights.pth?token=abc" + with mock.patch.object(utils.requests, "get", return_value=FakeResponse([b"ab", b"cd"])) as get: + path = utils.get_or_download_file(url, tempdir_prefix="models") + again = utils.get_or_download_file(url, tempdir_prefix="models") + + get.assert_called_once() + self.assertEqual(get.call_args.kwargs["timeout"], utils.DOWNLOAD_TIMEOUT) + self.assertEqual(path.read_bytes(), b"abcd") + self.assertEqual(path.name, "weights.pth") + self.assertEqual(path.stat().st_mode & 0o777, 0o644) + self.assertEqual(again.resolve(), path.resolve()) + self.assertEqual([p.name for p in (self.cache_root / "models").iterdir()], ["weights.pth"]) + + def test_interrupted_download_leaves_nothing_in_the_cache(self): + response = FakeResponse([b"ab"], error=ConnectionError("dropped")) + with mock.patch.object(utils.requests, "get", return_value=response): + with self.assertRaises(ConnectionError): + utils.get_or_download_file("https://example.org/weights.pth", tempdir_prefix="models") + self.assertEqual(list((self.cache_root / "models").iterdir()), []) + + def test_truncated_download_is_rejected(self): + response = FakeResponse([b"ab"], content_length=10) + with mock.patch.object(utils.requests, "get", return_value=response): + with self.assertRaises(OSError): + utils.get_or_download_file("https://example.org/weights.pth", tempdir_prefix="models") + self.assertEqual(list((self.cache_root / "models").iterdir()), []) + + +if __name__ == "__main__": + unittest.main() diff --git a/processing_services/example/api/utils.py b/processing_services/example/api/utils.py index a7fcb6a75..98bbbc912 100644 --- a/processing_services/example/api/utils.py +++ b/processing_services/example/api/utils.py @@ -10,6 +10,7 @@ import PIL.Image import PIL.ImageFile import requests +import torch logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG) @@ -38,11 +39,26 @@ def is_base64(s: str) -> bool: return False +# (connect, read) timeouts in seconds. The read timeout applies between received chunks, not to the whole download, +# so large model weights are fine as long as data keeps arriving. +DOWNLOAD_TIMEOUT = (10, 60) + +# Log download progress roughly this often (in bytes) instead of once per chunk. +DOWNLOAD_LOG_INTERVAL_BYTES = 10 * 1024 * 1024 + + +def get_cache_root() -> pathlib.Path: + """Directory holding downloaded model files and other cached downloads.""" + return pathlib.Path("/app/cache") if pathlib.Path("/app").exists() else pathlib.Path.home() / ".antenna_cache" + + def get_or_download_file(path_or_url, tempdir_prefix="antenna") -> pathlib.Path: """ - Fetch a file from a URL or local path. If the path is a URL, download the file. - If the URL has already been downloaded, return the existing local path. - If the path is a local path, return the path. + Fetch a file from a URL or local path. + + A local path is returned as-is (and must exist). A URL is downloaded into the cache directory under + `tempdir_prefix`, or reused if a file with the same name was already downloaded. The download is streamed to a + temporary file and renamed into place, so an interrupted download never leaves a truncated file in the cache. >>> filepath = get_or_download_file("https://example.uk/images/31-20230919033000-snapshot.jpg?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=451d406b7eb1113e1bb05c083ce51481%2F20240429%2F") # noqa: E501 >>> filepath.name @@ -52,33 +68,62 @@ def get_or_download_file(path_or_url, tempdir_prefix="antenna") -> pathlib.Path: '31-20230919033000-snapshot.jpg' """ if not path_or_url: - raise Exception("Specify a URL or path to fetch file from.") + raise ValueError("Specify a URL or path to fetch file from.") + + path_or_url = str(path_or_url) + if not is_url(path_or_url): + local_path = pathlib.Path(path_or_url) + if not local_path.exists(): + raise FileNotFoundError(f"Local file does not exist: {local_path}") + logger.info(f"Using local file: {local_path}") + return local_path + + destination_dir = get_cache_root() / tempdir_prefix + destination_dir.mkdir(parents=True, exist_ok=True) + local_filepath = destination_dir / pathlib.Path(urlparse(path_or_url).path).name + + if local_filepath.exists(): + logger.info(f"Using cached file: {local_filepath}") + return local_filepath - # If path is a local path instead of a URL then urlretrieve will just return that path + logger.info(f"Downloading {path_or_url} to {local_filepath}") + response = requests.get(path_or_url, stream=True, headers={"User-Agent": USER_AGENT}, timeout=DOWNLOAD_TIMEOUT) + response.raise_for_status() + total_size = int(response.headers.get("content-length", 0)) - destination_dir = pathlib.Path(tempfile.mkdtemp(prefix=tempdir_prefix)) - fname = pathlib.Path(urlparse(path_or_url).path).name - if not destination_dir.exists(): - destination_dir.mkdir(parents=True, exist_ok=True) - local_filepath = pathlib.Path(destination_dir) / fname + with tempfile.NamedTemporaryFile(dir=destination_dir, prefix=".download-", delete=False) as partial_file: + partial_path = pathlib.Path(partial_file.name) + try: + downloaded = 0 + next_log_at = DOWNLOAD_LOG_INTERVAL_BYTES + for chunk in response.iter_content(chunk_size=1024 * 1024): + partial_file.write(chunk) + downloaded += len(chunk) + if downloaded >= next_log_at: + percent = f" ({downloaded / total_size:.0%})" if total_size else "" + logger.info(f" Downloaded {downloaded // (1024 * 1024)} MB{percent}") + next_log_at += DOWNLOAD_LOG_INTERVAL_BYTES + if total_size and downloaded != total_size: + raise OSError(f"Incomplete download of {path_or_url}: got {downloaded} of {total_size} bytes") + except BaseException: + partial_path.unlink(missing_ok=True) + raise - if local_filepath and local_filepath.exists(): - logger.info(f"Using existing {local_filepath}") - return local_filepath + partial_path.chmod(0o644) # NamedTemporaryFile creates owner-only files, which breaks shared cache mounts. + partial_path.replace(local_filepath) + logger.info(f"Download completed: {local_filepath.resolve()}") + return local_filepath.resolve() + +def get_best_device() -> str: + """ + Returns the best available device for running the model. + MPS is not supported by the current algorithms. + """ + if torch.cuda.is_available(): + return f"cuda:{torch.cuda.current_device()}" else: - logger.info(f"Downloading {path_or_url} to {local_filepath}") - headers = {"User-Agent": USER_AGENT} - response = requests.get(path_or_url, stream=True, headers=headers) - response.raise_for_status() # Raise an exception for HTTP errors - - with open(local_filepath, "wb") as f: - for chunk in response.iter_content(chunk_size=8192): - f.write(chunk) - - resulting_filepath = pathlib.Path(local_filepath).resolve() - logger.info(f"Downloaded to {resulting_filepath}") - return resulting_filepath + return "cpu" def open_image(fp: str | bytes | pathlib.Path | io.BytesIO, raise_exception: bool = True) -> PIL.Image.Image | None: diff --git a/processing_services/example/docker-compose.yml b/processing_services/example/docker-compose.yml index 83db6ccfa..38469734e 100644 --- a/processing_services/example/docker-compose.yml +++ b/processing_services/example/docker-compose.yml @@ -4,8 +4,11 @@ services: context: . volumes: - ./:/app:z + - ./cache:/app/cache # Model weights and labels cache - ./huggingface_cache:/root/.cache/huggingface - ./pytorch_cache:/root/.cache/torch + environment: + - MODEL_BASE_URL # Optional: mirror or local directory for model weights ports: - "2003:2000" extra_hosts: diff --git a/processing_services/example/requirements.txt b/processing_services/example/requirements.txt index eccbee47a..251693690 100644 --- a/processing_services/example/requirements.txt +++ b/processing_services/example/requirements.txt @@ -6,4 +6,5 @@ requests==2.32.4 transformers==4.50.3 torch==2.6.0 torchvision==0.21.0 +timm==1.0.11 scipy==1.16.0