Upload inference.py

inference.py

@@ -0,0 +1,292 @@
+"""
+TAS-BB-v1 MEWC-Keras Classification Model - Custom Inference Script
+
+This script provides model-specific inference code for the Tasmania
+species classifier (96 classes). It follows the AddaxAI-WebUI interface contract
+for custom classification models.
+
+Model: Tasmania MEWC Ensemble (tas_ens_mewc.keras)
+Framework: Keras 3 with JAX backend (TensorFlow compatible)
+Classes: 96 classes (Tasmanian terrestrial mammals and birds)
+Training data: 2.5+ million images
+
+MEWC - Mega Efficient Wildlife Classifier - University of Tasmania
+Original source: streamlit-AddaxAI/classification/model_types/mewc-keras/classify_detections.py
+Reference: https://github.com/zaandahl/mewc
+Adapted by: Claude Code on 2026-01-11
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import cv2
+import numpy as np
+import tensorflow as tf
+import yaml
+from keras import saving
+from PIL import Image, ImageFile
+
+# Set Keras backend to JAX (as per original MEWC code)
+os.environ["KERAS_BACKEND"] = "jax"
+
+# 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 = None
+img_size = 384  # MEWC uses 384x384 images
+
+# Class mapping variables
+class_map: dict[str, str] | None = None
+class_ids: list[str] | None = None
+
+# Don't freak out over truncated images
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+def check_gpu() -> bool:
+    """
+    Check GPU availability for TensorFlow/Keras inference.
+
+    TensorFlow can detect GPUs, Metal (Apple Silicon), and CUDA.
+
+    Returns:
+        True if GPU available, False otherwise
+    """
+    try:
+        gpus = tf.config.list_logical_devices('GPU')
+        return len(gpus) > 0
+    except Exception:
+        return False
+
+
+def load_model() -> None:
+    """
+    Load Keras 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.
+
+    Also loads the class_list.yaml file which maps class indices to species names.
+
+    Raises:
+        RuntimeError: If model loading fails
+        FileNotFoundError: If MODEL_PATH or class_list.yaml is invalid
+    """
+    global animal_model, class_map, class_ids, MODEL_PATH, MODEL_DIR
+
+    if MODEL_PATH is None:
+        raise RuntimeError("MODEL_PATH not set - must be injected by framework")
+
+    if MODEL_DIR is None:
+        raise RuntimeError("MODEL_DIR not set - must be injected by framework")
+
+    if not MODEL_PATH.exists():
+        raise FileNotFoundError(f"Model file not found: {MODEL_PATH}")
+
+    # Load the Keras model (without compilation for inference only)
+    try:
+        animal_model = saving.load_model(str(MODEL_PATH), compile=False)
+    except Exception as e:
+        raise RuntimeError(f"Failed to load Keras model from {MODEL_PATH}: {e}") from e
+
+    # Load class_list.yaml
+    class_list_path = MODEL_DIR / "class_list.yaml"
+    if not class_list_path.exists():
+        raise FileNotFoundError(
+            f"class_list.yaml not found: {class_list_path}\n"
+            f"MEWC models require class_list.yaml in the model directory."
+        )
+
+    try:
+        with open(class_list_path, 'r') as f:
+            class_map = yaml.safe_load(f)
+    except Exception as e:
+        raise RuntimeError(f"Failed to load class_list.yaml: {e}") from e
+
+    # Build class_ids list based on YAML format
+    # The YAML can be formatted as either {int: str} or {str: int}
+    inv_class = {v: k for k, v in class_map.items()}
+
+    # Check if keys are numeric (int:label format) or string (label:int format)
+    formatted_int_label = _can_all_keys_be_converted_to_int(class_map)
+
+    if formatted_int_label:
+        # Format: {0: "species1", 1: "species2", ...}
+        class_ids = [class_map[i] for i in sorted(inv_class.values())]
+    else:
+        # Format: {"species1": 0, "species2": 1, ...}
+        class_ids = sorted(inv_class.values())
+
+
+def _can_all_keys_be_converted_to_int(d: dict) -> bool:
+    """
+    Check if all dictionary keys can be converted to integers.
+
+    Used to determine class_list.yaml format.
+
+    Args:
+        d: Dictionary to check
+
+    Returns:
+        True if all keys are convertible to int, False otherwise
+    """
+    for key in d.keys():
+        try:
+            int(key)
+        except ValueError:
+            return False
+    return True
+
+
+def get_crop(image: Image.Image, bbox: tuple[float, float, float, float]) -> Image.Image | None:
+    """
+    Crop image using MEWC-specific preprocessing.
+
+    This cropping method is used by MEWC and follows the MegaDetector
+    visualization_utils approach. It:
+    1. Denormalizes the bbox coordinates
+    2. Clips to image boundaries
+    3. Returns the cropped region (no padding or squaring)
+
+    Reference: https://github.com/zaandahl/mewc-snip/blob/main/src/mewc_snip.py#L29
+    Reference: https://github.com/agentmorris/MegaDetector/blob/main/megadetector/visualization/visualization_utils.py#L352
+
+    Args:
+        image: PIL Image (full resolution)
+        bbox: Normalized bounding box (x, y, width, height) in range [0.0, 1.0]
+
+    Returns:
+        Cropped PIL Image, or None if bbox is invalid
+
+    Raises:
+        None - Returns None for invalid boxes (graceful degradation)
+    """
+    x1, y1, w_box, h_box = bbox
+
+    # Check for invalid bounding boxes (zero or negative dimensions)
+    if w_box <= 0 or h_box <= 0:
+        return None
+
+    # Convert normalized coordinates to pixel coordinates
+    ymin, xmin, ymax, xmax = y1, x1, y1 + h_box, x1 + w_box
+    im_width, im_height = image.size
+
+    # Denormalize
+    left = xmin * im_width
+    right = xmax * im_width
+    top = ymin * im_height
+    bottom = ymax * im_height
+
+    # Clip to image boundaries (ensure non-negative)
+    left = max(left, 0)
+    right = max(right, 0)
+    top = max(top, 0)
+    bottom = max(bottom, 0)
+
+    # Clip to image boundaries (ensure within image)
+    left = min(left, im_width - 1)
+    right = min(right, im_width - 1)
+    top = min(top, im_height - 1)
+    bottom = min(bottom, im_height - 1)
+
+    # Final check - ensure crop has valid dimensions
+    crop_width = right - left
+    crop_height = bottom - top
+
+    if crop_width <= 0 or crop_height <= 0:
+        return None
+
+    # Crop image
+    image_cropped = image.crop((left, top, right, bottom))
+    return image_cropped
+
+
+def get_classification(crop: Image.Image) -> list[tuple[str, float]]:
+    """
+    Run MEWC-Keras classification on cropped image.
+
+    Workflow:
+    1. Convert PIL Image to numpy array
+    2. Resize to 384x384 (MEWC input size)
+    3. Run model prediction
+    4. Return all class probabilities
+
+    Args:
+        crop: Cropped PIL Image
+
+    Returns:
+        List of (class_id, confidence) tuples for ALL classes, in model order.
+        Example: [("10", 0.50674), ("33", 0.46682), ...]
+
+    Raises:
+        RuntimeError: If model not loaded or inference fails
+    """
+    global animal_model, class_ids
+
+    if animal_model is None:
+        raise RuntimeError("Model not loaded - call load_model() first")
+
+    if class_ids is None:
+        raise RuntimeError("Class IDs not loaded - call load_model() first")
+
+    if crop is None:
+        return []
+
+    try:
+        # Convert PIL Image to numpy array
+        img = np.array(crop)
+
+        if img.size == 0:
+            return []
+
+        # Resize to MEWC input size (384x384)
+        img = cv2.resize(img, (img_size, img_size))
+
+        # Add batch dimension
+        img = np.expand_dims(img, axis=0)
+
+        # Run prediction (verbose=0 suppresses progress bar)
+        pred = animal_model.predict(img, verbose=0)[0]
+
+        # Build list of (class_id, confidence) tuples
+        # class_ids is already in the correct order from class_list.yaml
+        classifications = []
+        for i in range(len(pred)):
+            class_id_str = str(i + 1)  # 1-indexed for compatibility
+            confidence = float(pred[i])
+            classifications.append((class_id_str, confidence))
+
+        return classifications
+
+    except Exception as e:
+        raise RuntimeError(f"MEWC-Keras classification failed: {e}") from e
+
+
+def get_class_names() -> dict[str, str]:
+    """
+    Get mapping of class IDs to species names from class_list.yaml.
+
+    Returns:
+        Dict mapping class ID (1-indexed string) to species name
+        Example: {"1": "bait", "2": "unknown_animal", ...}
+
+    Raises:
+        RuntimeError: If class_map not loaded
+    """
+    global class_ids
+
+    if class_ids is None:
+        raise RuntimeError("Class IDs not loaded - call load_model() first")
+
+    # Build 1-indexed mapping
+    class_names = {}
+    for i, class_name in enumerate(class_ids):
+        class_id_str = str(i + 1)  # 1-indexed
+        class_names[class_id_str] = class_name
+
+    return class_names