Upload inference.py

inference.py

@@ -0,0 +1,262 @@
+"""
+NAM-ADS-v1 YOLOv8 Classification Model - Custom Inference Script
+
+This script provides model-specific inference code for the Namibian Desert
+species classifier (30 classes). It follows the AddaxAI-WebUI interface contract
+for custom classification models.
+
+Model: Namibian Desert v1 (namib_desert_v1.pt)
+Framework: YOLOv8 (Ultralytics)
+Classes: 30 species from Skeleton Coast National Park, Namibia
+Training data: 850,000+ images
+
+Original source: streamlit-AddaxAI/classification/model_types/addax-yolov8/classify_detections.py
+Adapted by: Claude Code on 2026-01-11
+"""
+
+from __future__ import annotations
+
+import pathlib
+import platform
+from pathlib import Path
+
+import torch
+from PIL import Image, ImageFile, ImageOps
+from ultralytics import YOLO
+
+# Module-level variables (injected by AddaxAI framework)
+MODEL_DIR: Path | None = None  # Set by CustomInferenceLoader
+MODEL_PATH: Path | None = None  # Set by CustomInferenceLoader
+
+# Module-level model instance (loaded once at startup)
+animal_model: YOLO | None = None
+
+# Don't freak out over truncated images
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+# Make sure Windows-trained models work on Unix
+plt = platform.system()
+if plt != 'Windows':
+    pathlib.WindowsPath = pathlib.PosixPath
+
+
+def check_gpu() -> bool:
+    """
+    Check GPU availability for YOLOv8 inference.
+
+    Checks both Apple Metal Performance Shaders (MPS) and CUDA availability.
+
+    Returns:
+        True if GPU available, False otherwise
+    """
+    # Check Apple MPS (Apple Silicon)
+    try:
+        if torch.backends.mps.is_built() and torch.backends.mps.is_available():
+            return True
+    except Exception:
+        pass
+
+    # Check CUDA (NVIDIA)
+    return torch.cuda.is_available()
+
+
+def load_model() -> None:
+    """
+    Load YOLOv8 classification model into memory.
+
+    This function is called once during worker initialization.
+    The model is stored in the global `animal_model` variable and reused
+    for all subsequent classification requests.
+
+    Raises:
+        RuntimeError: If model loading fails
+        FileNotFoundError: If MODEL_PATH is invalid
+    """
+    global animal_model, MODEL_PATH
+
+    if MODEL_PATH is None:
+        raise RuntimeError("MODEL_PATH not set - must be injected by framework")
+
+    if not MODEL_PATH.exists():
+        raise FileNotFoundError(f"Model file not found: {MODEL_PATH}")
+
+    try:
+        animal_model = YOLO(str(MODEL_PATH))
+    except Exception as e:
+        raise RuntimeError(f"Failed to load YOLOv8 model from {MODEL_PATH}: {e}") from e
+
+
+def get_crop(image: Image.Image, bbox: tuple[float, float, float, float]) -> Image.Image:
+    """
+    Crop image using model-specific preprocessing.
+
+    This cropping method was developed by Dan Morris for MegaDetector and is
+    designed to:
+    1. Square the bounding box (max of width/height)
+    2. Add padding to prevent over-enlargement of small animals
+    3. Center the detection within the crop
+    4. Pad with black (0) to maintain square aspect ratio
+
+    Args:
+        image: PIL Image (full resolution)
+        bbox: Normalized bounding box (x, y, width, height) in range [0.0, 1.0]
+
+    Returns:
+        Cropped and padded PIL Image ready for classification
+
+    Raises:
+        ValueError: If bbox is invalid (zero size)
+    """
+    img_w, img_h = image.size
+
+    # Denormalize bbox coordinates
+    xmin = int(bbox[0] * img_w)
+    ymin = int(bbox[1] * img_h)
+    box_w = int(bbox[2] * img_w)
+    box_h = int(bbox[3] * img_h)
+
+    # Square the box (use max dimension)
+    box_size = max(box_w, box_h)
+
+    # Add padding (prevents over-enlargement of small animals)
+    box_size = _pad_crop(box_size)
+
+    # Center the detection within the squared crop
+    xmin = max(0, min(
+        xmin - int((box_size - box_w) / 2),
+        img_w - box_w
+    ))
+    ymin = max(0, min(
+        ymin - int((box_size - box_h) / 2),
+        img_h - box_h
+    ))
+
+    # Clip to image boundaries
+    box_w = min(img_w, box_size)
+    box_h = min(img_h, box_size)
+
+    if box_w == 0 or box_h == 0:
+        raise ValueError(f"Invalid bbox size: {box_w}x{box_h}")
+
+    # Crop and pad to square
+    crop = image.crop(box=[xmin, ymin, xmin + box_w, ymin + box_h])
+    crop = ImageOps.pad(crop, size=(box_size, box_size), color=0)
+
+    return crop
+
+
+def _pad_crop(box_size: int) -> int:
+    """
+    Calculate padded crop size to prevent over-enlargement of small animals.
+
+    YOLOv8 expects 224x224 input. This function ensures small detections aren't
+    excessively upscaled while adding consistent padding to larger detections.
+
+    Args:
+        box_size: Original bounding box size (max of width/height)
+
+    Returns:
+        Padded box size
+    """
+    input_size_network = 224
+    default_padding = 30
+
+    if box_size >= input_size_network:
+        # Large detection: add default padding
+        return box_size + default_padding
+    else:
+        # Small detection: ensure minimum size without excessive enlargement
+        diff_size = input_size_network - box_size
+        if diff_size < default_padding:
+            return box_size + default_padding
+        else:
+            return input_size_network
+
+
+def get_classification(crop: Image.Image) -> list[tuple[str, float]]:
+    """
+    Run YOLOv8 classification on cropped image.
+
+    Args:
+        crop: Cropped and preprocessed PIL Image
+
+    Returns:
+        List of (class_id, confidence) tuples for ALL classes, sorted by confidence.
+        Example: [("14", 0.99985), ("7", 0.00003), ...]
+
+    Raises:
+        RuntimeError: If model not loaded or inference fails
+    """
+    global animal_model
+
+    if animal_model is None:
+        raise RuntimeError("Model not loaded - call load_model() first")
+
+    try:
+        # Run YOLOv8 classification (verbose=False suppresses progress bar)
+        results = animal_model(crop, verbose=False)
+
+        # Extract class names dict: {0: "porcupine", 1: "elephant", ...}
+        names_dict = results[0].names
+
+        # Extract probabilities: [0.0001, 0.0002, ..., 0.9998, ...]
+        probs = results[0].probs.data.tolist()
+
+        # Build list of (class_id, confidence) tuples
+        # Class IDs are 0-indexed in YOLOv8 but we output as strings
+        classifications = []
+        for idx, class_name in names_dict.items():
+            # YOLOv8 uses 0-based indexing, but we need 1-based for compatibility
+            # with the taxonomy.csv and expected JSON output
+            class_id_str = str(idx + 1)  # Convert 0-indexed to 1-indexed
+            confidence = probs[idx]
+            classifications.append((class_id_str, confidence))
+
+        # Sort by confidence descending (already sorted by YOLOv8, but ensure it)
+        classifications.sort(key=lambda x: x[1], reverse=True)
+
+        return classifications
+
+    except Exception as e:
+        raise RuntimeError(f"YOLOv8 classification failed: {e}") from e
+
+
+def get_class_names() -> dict[str, str]:
+    """
+    Get mapping of class IDs to species names from taxonomy.csv.
+
+    Reads taxonomy.csv from the model directory and extracts the model_class
+    (common name) for each species.
+
+    Returns:
+        Dict mapping class ID (1-indexed string) to common name
+        Example: {"1": "porcupine", "2": "elephant", ...}
+
+    Raises:
+        FileNotFoundError: If taxonomy.csv not found
+        RuntimeError: If parsing fails
+    """
+    global animal_model, MODEL_DIR
+
+    if MODEL_DIR is None:
+        raise RuntimeError("MODEL_DIR not set - must be injected by framework")
+
+    # YOLOv8 models have class names built-in
+    # We'll use those directly since they should match taxonomy.csv
+    if animal_model is None:
+        raise RuntimeError("Model not loaded - call load_model() first")
+
+    try:
+        # YOLOv8 names dict: {0: "porcupine", 1: "elephant", ...}
+        yolo_names = animal_model.names
+
+        # Convert to 1-indexed dict for compatibility
+        class_names = {}
+        for idx, name in yolo_names.items():
+            class_id_str = str(idx + 1)  # 1-indexed
+            class_names[class_id_str] = name
+
+        return class_names
+
+    except Exception as e:
+        raise RuntimeError(f"Failed to extract class names: {e}") from e