Upload inference.py

inference.py

@@ -0,0 +1,232 @@
+"""
+Inference script for PAM-SDZWA-v1 (Peruvian Amazon Species Classifier)
+
+This model classifies 53 species found in Peruvian Amazon rainforest habitats.
+Developed by Mathias Tobler from the San Diego Zoo Wildlife Alliance Conservation
+Technology Lab using their animl-py framework.
+
+Model: Peru Amazon v0.86
+Input: Variable size (extracted from model config)
+Framework: TensorFlow/Keras (TensorFlow 1.x compatible)
+Classes: 53 Amazonian species and taxonomic groups
+Developer: San Diego Zoo Wildlife Alliance (Mathias Tobler)
+License: MIT
+Info: https://github.com/conservationtechlab
+
+Author: Peter van Lunteren
+Created: 2026-01-14
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import cv2
+import numpy as np
+import tensorflow as tf
+from PIL import Image, ImageFile
+from tensorflow.keras.models import load_model
+
+# Don't freak out over truncated images
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+class ModelInference:
+    """TensorFlow/Keras inference implementation for Peruvian Amazon species classifier."""
+
+    def __init__(self, model_dir: Path, model_path: Path):
+        """
+        Initialize with model paths.
+
+        Args:
+            model_dir: Directory containing model files and class labels
+            model_path: Path to Peru-Amazon_0.86.h5 file
+        """
+        self.model_dir = model_dir
+        self.model_path = model_path
+        self.model = None
+        self.img_size = None
+        self.class_map = {}
+        self.class_ids_sorted = []
+
+    def check_gpu(self) -> bool:
+        """
+        Check GPU availability for TensorFlow inference.
+
+        Returns:
+            True if GPU available, False otherwise
+        """
+        return len(tf.config.list_logical_devices('GPU')) > 0
+
+    def load_model(self) -> None:
+        """
+        Load TensorFlow/Keras model and class labels into memory.
+
+        This function is called once during worker initialization.
+        The model is stored in self.model and reused for all subsequent
+        classification requests.
+
+        Raises:
+            RuntimeError: If model loading fails
+            FileNotFoundError: If model_path or label file is invalid
+        """
+        if not self.model_path.exists():
+            raise FileNotFoundError(f"Model file not found: {self.model_path}")
+
+        try:
+            # Load Keras model
+            self.model = load_model(str(self.model_path))
+
+            # Extract input image size from model config
+            # Model expects square images (e.g., 299x299)
+            self.img_size = self.model.get_config()["layers"][0]["config"]["batch_input_shape"][1]
+
+        except Exception as e:
+            raise RuntimeError(f"Failed to load Keras model from {self.model_path}: {e}") from e
+
+        # Load class labels from text file
+        label_file = self.model_dir / "Peru-Amazon_0.86.txt"
+        if not label_file.exists():
+            raise FileNotFoundError(f"Class label file not found: {label_file}")
+
+        try:
+            with open(label_file, 'r') as file:
+                for line in file:
+                    parts = line.strip().split('"')
+                    if len(parts) >= 4:
+                        identifier = parts[1].strip()
+                        animal_name = parts[3].strip()
+                        if identifier.isdigit():
+                            self.class_map[str(identifier)] = str(animal_name)
+
+            # Create sorted list of class names (sorted by ID)
+            # This ensures consistent ordering for inference results
+            self.class_ids_sorted = sorted(self.class_map.values())
+
+        except Exception as e:
+            raise RuntimeError(f"Failed to load class labels from {label_file}: {e}") from e
+
+    def get_crop(
+        self, image: Image.Image, bbox: tuple[float, float, float, float]
+    ) -> Image.Image:
+        """
+        Crop image using SDZWA animl-py preprocessing.
+
+        This cropping method follows the San Diego Zoo Wildlife Alliance's animl-py
+        framework approach with minimal buffering (0 pixels by default).
+
+        Based on: https://github.com/conservationtechlab/animl-py/blob/main/src/animl/generator.py
+
+        Args:
+            image: PIL Image (full resolution)
+            bbox: Normalized bounding box (x, y, width, height) in range [0.0, 1.0]
+
+        Returns:
+            Cropped PIL Image (not resized - resizing happens in get_classification)
+
+        Raises:
+            ValueError: If bbox is invalid
+        """
+        buffer = 0  # SDZWA uses 0 pixel buffer
+        width, height = image.size
+
+        # Denormalize bbox coordinates
+        bbox1, bbox2, bbox3, bbox4 = bbox
+        left = width * bbox1
+        top = height * bbox2
+        right = width * (bbox1 + bbox3)
+        bottom = height * (bbox2 + bbox4)
+
+        # Apply buffer and clip to image boundaries
+        left = max(0, int(left) - buffer)
+        top = max(0, int(top) - buffer)
+        right = min(width, int(right) + buffer)
+        bottom = min(height, int(bottom) + buffer)
+
+        # Validate crop dimensions
+        if left >= right or top >= bottom:
+            raise ValueError(f"Invalid bbox dimensions after cropping: left={left}, top={top}, right={right}, bottom={bottom}")
+
+        # Crop and return
+        image_cropped = image.crop((left, top, right, bottom))
+        return image_cropped
+
+    def get_classification(self, crop: Image.Image) -> list[list[str, float]]:
+        """
+        Run TensorFlow/Keras classification on cropped image.
+
+        Preprocessing follows SDZWA animl-py framework:
+        - Resize to model input size (extracted from model config)
+        - Convert to numpy array
+        - No normalization or augmentation (except potential horizontal flip during training)
+
+        Args:
+            crop: Cropped PIL Image
+
+        Returns:
+            List of [class_name, confidence] lists for ALL classes, sorted by class ID.
+            Example: [["Black-headed squirrel monkey", 0.001], ["Brazilian rabbit", 0.002], ...]
+            NOTE: Sorting by confidence is handled by classification_worker.py
+
+        Raises:
+            RuntimeError: If model not loaded or inference fails
+        """
+        if self.model is None:
+            raise RuntimeError("Model not loaded - call load_model() first")
+
+        try:
+            # Convert PIL to numpy array
+            img = np.array(crop)
+
+            # Resize to model input size using OpenCV
+            img = cv2.resize(img, (self.img_size, self.img_size))
+
+            # Add batch dimension
+            img = np.expand_dims(img, axis=0)
+
+            # Run inference
+            # Note: According to animl-py, no special preprocessing is needed
+            # except for horizontal flip augmentation during training
+            pred = self.model.predict(img, verbose=0)[0]
+
+            # Build list of [class_name, confidence] pairs
+            # Use sorted class IDs to maintain consistent ordering
+            classifications = []
+            for i in range(len(pred)):
+                class_name = self.class_ids_sorted[i]
+                confidence = float(pred[i])
+                classifications.append([class_name, confidence])
+
+            return classifications
+
+        except Exception as e:
+            raise RuntimeError(f"Keras classification failed: {e}") from e
+
+    def get_class_names(self) -> dict[str, str]:
+        """
+        Get mapping of class IDs to species names.
+
+        Class IDs are 1-indexed and correspond to the sorted order of class names.
+
+        Returns:
+            Dict mapping class ID (1-indexed string) to species name
+            Example: {"1": "Black-headed squirrel monkey", "2": "Brazilian rabbit", ...}
+
+        Raises:
+            RuntimeError: If model not loaded
+        """
+        if self.model is None:
+            raise RuntimeError("Model not loaded - call load_model() first")
+
+        try:
+            # Create 1-indexed mapping of class IDs to names
+            class_names = {}
+            for i, class_name in enumerate(self.class_ids_sorted):
+                class_id_str = str(i + 1)  # 1-indexed
+                class_names[class_id_str] = class_name
+
+            return class_names
+
+        except Exception as e:
+            raise RuntimeError(f"Failed to extract class names: {e}") from e