Upload inference.py

inference.py

@@ -14,6 +14,7 @@ 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
+Updated: 2026-01-13 - Migrated to class-based interface
 """
 
 from __future__ import annotations
@@ -31,258 +32,253 @@ 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
+class ModelInference:
+    """MEWC-Keras inference implementation for Tasmania species classifier."""
 
-    # MODEL_PATH and MODEL_DIR are injected by framework before this function is called
-    # Check that the paths exist (framework guarantees they're not None)
-    if not MODEL_PATH.exists():
-        raise FileNotFoundError(f"Model file not found: {MODEL_PATH}")
+    def __init__(self, model_dir: Path, model_path: Path):
+        """
+        Initialize with model paths.
 
-    # 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
+        Args:
+            model_dir: Directory containing model files (including class_list.yaml)
+            model_path: Path to tas_ens_mewc.keras file
+        """
+        self.model_dir = model_dir
+        self.model_path = model_path
+        self.model = None
+        self.img_size = 384  # MEWC uses 384x384 images
+        self.class_map: dict[str, str] | None = None
+        self.class_ids: list[str] | None = None
 
-    # 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."
-        )
+    def check_gpu(self) -> bool:
+        """
+        Check GPU availability for TensorFlow/Keras inference.
 
-    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
+        TensorFlow can detect GPUs, Metal (Apple Silicon), and CUDA.
 
-    # 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():
+        Returns:
+            True if GPU available, False otherwise
+        """
         try:
-            int(key)
-        except ValueError:
+            gpus = tf.config.list_logical_devices('GPU')
+            return len(gpus) > 0
+        except Exception:
             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
+    def load_model(self) -> None:
+        """
+        Load Keras classification model into memory.
 
-    Args:
-        image: PIL Image (full resolution)
-        bbox: Normalized bounding box (x, y, width, height) in range [0.0, 1.0]
+        This function is called once during worker initialization.
+        The model is stored in self.model and reused for all subsequent
+        classification requests.
 
-    Returns:
-        Cropped PIL Image, or None if bbox is invalid
+        Also loads the class_list.yaml file which maps class indices to species names.
 
-    Raises:
-        None - Returns None for invalid boxes (graceful degradation)
-    """
-    x1, y1, w_box, h_box = bbox
+        Raises:
+            RuntimeError: If model loading fails
+            FileNotFoundError: If model_path or class_list.yaml is invalid
+        """
+        if not self.model_path.exists():
+            raise FileNotFoundError(f"Model file not found: {self.model_path}")
 
-    # 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_name, confidence) tuples for ALL classes, in model order.
-        Example: [("tasmanian_pademelon", 0.50674), ("bennetts_wallaby", 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)
+        # Load the Keras model (without compilation for inference only)
+        try:
+            self.model = saving.load_model(str(self.model_path), compile=False)
+        except Exception as e:
+            raise RuntimeError(f"Failed to load Keras model from {self.model_path}: {e}") from e
+
+        # Load class_list.yaml
+        class_list_path = self.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."
+            )
 
-        if img.size == 0:
+        try:
+            with open(class_list_path, 'r') as f:
+                self.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 self.class_map.items()}
+
+        # Check if keys are numeric (int:label format) or string (label:int format)
+        formatted_int_label = self._can_all_keys_be_converted_to_int(self.class_map)
+
+        if formatted_int_label:
+            # Format: {0: "species1", 1: "species2", ...}
+            self.class_ids = [self.class_map[i] for i in sorted(inv_class.values())]
+        else:
+            # Format: {"species1": 0, "species2": 1, ...}
+            self.class_ids = sorted(inv_class.values())
+
+    def _can_all_keys_be_converted_to_int(self, 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(
+        self, 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(self, 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_name, confidence) tuples for ALL classes, in model order.
+            Example: [("tasmanian_pademelon", 0.50674), ("bennetts_wallaby", 0.46682), ...]
+
+        Raises:
+            RuntimeError: If model not loaded or inference fails
+        """
+        if self.model is None:
+            raise RuntimeError("Model not loaded - call load_model() first")
+
+        if self.class_ids is None:
+            raise RuntimeError("Class IDs not loaded - call load_model() first")
+
+        if crop is None:
             return []
 
-        # Resize to MEWC input size (384x384)
-        img = cv2.resize(img, (img_size, img_size))
+        try:
+            # Convert PIL Image to numpy array
+            img = np.array(crop)
 
-        # Add batch dimension
-        img = np.expand_dims(img, axis=0)
+            if img.size == 0:
+                return []
 
-        # Run prediction (verbose=0 suppresses progress bar)
-        pred = animal_model.predict(img, verbose=0)[0]
+            # Resize to MEWC input size (384x384)
+            img = cv2.resize(img, (self.img_size, self.img_size))
 
-        # Build list of (class_name, confidence) tuples
-        # class_ids is already in the correct order from class_list.yaml
-        classifications = []
-        for i in range(len(pred)):
-            class_name = class_ids[i]  # Get species name from class_list.yaml
-            confidence = float(pred[i])
-            classifications.append((class_name, confidence))
+            # Add batch dimension
+            img = np.expand_dims(img, axis=0)
 
-        return classifications
+            # Run prediction (verbose=0 suppresses progress bar)
+            pred = self.model.predict(img, verbose=0)[0]
 
-    except Exception as e:
-        raise RuntimeError(f"MEWC-Keras classification failed: {e}") from e
+            # Build list of (class_name, confidence) tuples
+            # class_ids is already in the correct order from class_list.yaml
+            classifications = []
+            for i in range(len(pred)):
+                class_name = self.class_ids[i]  # Get species name from class_list.yaml
+                confidence = float(pred[i])
+                classifications.append((class_name, confidence))
 
+            return classifications
 
-def get_class_names() -> dict[str, str]:
-    """
-    Get mapping of class IDs to species names from class_list.yaml.
+        except Exception as e:
+            raise RuntimeError(f"MEWC-Keras classification failed: {e}") from e
 
-    Returns:
-        Dict mapping class ID (1-indexed string) to species name
-        Example: {"1": "bait", "2": "unknown_animal", ...}
+    def get_class_names(self) -> dict[str, str]:
+        """
+        Get mapping of class IDs to species names from class_list.yaml.
 
-    Raises:
-        RuntimeError: If class_map not loaded
-    """
-    global class_ids
+        Returns:
+            Dict mapping class ID (1-indexed string) to species name
+            Example: {"1": "bait", "2": "unknown_animal", ...}
 
-    if class_ids is None:
-        raise RuntimeError("Class IDs not loaded - call load_model() first")
+        Raises:
+            RuntimeError: If class_ids not loaded
+        """
+        if self.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
+        # Build 1-indexed mapping
+        class_names = {}
+        for i, class_name in enumerate(self.class_ids):
+            class_id_str = str(i + 1)  # 1-indexed
+            class_names[class_id_str] = class_name
 
-    return class_names
+        return class_names