Story 5

src/classifier.py

@@ -0,0 +1,138 @@
+"""
+Face crop classifier.
+
+Loads a MobileNetV3 model via ONNX/TensorRT FP32, applies temperature-scaled
+softmax, and returns per-crop probabilities for the target class.
+"""
+
+import logging
+from pathlib import Path
+from typing import Any, List, Optional
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+# ImageNet normalisation constants
+_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+_STD  = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+
+
+def calibrated_softmax(logits: np.ndarray, temperature: float = 1.0) -> np.ndarray:
+    """Softmax with temperature scaling; returns class probabilities."""
+    scaled = logits / max(temperature, 1e-8)
+    shifted = scaled - scaled.max(axis=-1, keepdims=True)   # numerical stability
+    exp = np.exp(shifted)
+    return exp / exp.sum(axis=-1, keepdims=True)
+
+
+class FaceClassifier:
+    """Classify face crops and return calibrated target-class probabilities."""
+
+    def __init__(
+        self,
+        engine_path: str,
+        temperature: float = 1.0,
+        input_size: int = 224,
+        device: str = "cuda",
+    ):
+        self.engine_path = engine_path
+        self.temperature = temperature
+        self.input_size = input_size
+        self.device = device
+        self._session: Optional[Any] = None
+        self._load_model()
+
+    def _load_model(self):
+        path = Path(self.engine_path)
+        if not path.exists():
+            logger.warning("Classifier engine not found at %s; using placeholder", self.engine_path)
+            return
+
+        suffix = path.suffix.lower()
+        if suffix == ".onnx":
+            try:
+                import onnxruntime as ort
+                providers = (["CUDAExecutionProvider"] if self.device.startswith("cuda")
+                             else ["CPUExecutionProvider"])
+                self._session = ort.InferenceSession(str(path), providers=providers)
+                logger.info("Loaded ONNX classifier from %s", self.engine_path)
+            except Exception as exc:
+                logger.warning("Failed to load ONNX classifier: %s", exc)
+        elif suffix in (".trt", ".engine", ".plan"):
+            try:
+                import tensorrt as trt
+                with trt.Logger() as trt_logger, open(path, "rb") as f:
+                    runtime = trt.Runtime(trt_logger)
+                    self._session = runtime.deserialize_cuda_engine(f.read())
+                logger.info("Loaded TensorRT classifier from %s", self.engine_path)
+            except Exception as exc:
+                logger.warning("Failed to load TensorRT classifier: %s", exc)
+        else:
+            logger.warning("Unsupported classifier format: %s", suffix)
+
+    def classify(self, crop_paths: List[str], batch_size: int = 16) -> List[float]:
+        """Return target-class probabilities (p ∈ [0,1]) for each crop path."""
+        results: List[float] = []
+        for i in range(0, len(crop_paths), batch_size):
+            batch = crop_paths[i : i + batch_size]
+            results.extend(self._classify_batch(batch))
+        return results
+
+    def _classify_batch(self, crop_paths: List[str]) -> List[float]:
+        preprocessed = []
+        for path in crop_paths:
+            try:
+                preprocessed.append(self._preprocess(path))
+            except Exception as exc:
+                logger.warning("Preprocessing failed for %s: %s", path, exc)
+                preprocessed.append(np.zeros((3, self.input_size, self.input_size), dtype=np.float32))
+
+        batch = np.stack(preprocessed, axis=0)  # (N, 3, H, W)
+        logits = self._infer(batch)             # (N, 2)
+        probs = calibrated_softmax(logits, self.temperature)
+        return probs[:, 1].tolist()             # target-class column
+
+    def _preprocess(self, crop_path: str) -> np.ndarray:
+        from PIL import Image
+        image = Image.open(crop_path).convert("RGB")
+        if image.size != (self.input_size, self.input_size):
+            image = image.resize((self.input_size, self.input_size), Image.Resampling.BILINEAR)
+        arr = np.array(image, dtype=np.float32) / 255.0
+        arr = (arr - _MEAN) / _STD
+        image.close()
+        return np.transpose(arr, (2, 0, 1))    # HWC → CHW
+
+    def _infer(self, batch: np.ndarray) -> np.ndarray:
+        if self._session is None:
+            return self._placeholder_logits(batch.shape[0])
+
+        try:
+            if hasattr(self._session, "run"):
+                input_name = self._session.get_inputs()[0].name
+                output = self._session.run(None, {input_name: batch})
+                return np.array(output[0])
+
+            # TensorRT path
+            import pycuda.autoinit  # noqa: F401
+            import pycuda.driver as cuda
+
+            context = self._session.create_execution_context()
+            out_shape = (batch.shape[0], 2)
+            d_in  = cuda.mem_alloc(batch.nbytes)
+            d_out = cuda.mem_alloc(np.prod(out_shape) * np.dtype(np.float32).itemsize)
+            stream = cuda.Stream()
+            cuda.memcpy_htod_async(d_in, batch, stream)
+            context.execute_async_v2(bindings=[int(d_in), int(d_out)], stream_handle=stream.handle)
+            output = np.empty(out_shape, dtype=np.float32)
+            cuda.memcpy_dtoh_async(output, d_out, stream)
+            stream.synchronize()
+            return output
+        except Exception as exc:
+            logger.warning("Classifier inference failed: %s", exc)
+            return self._placeholder_logits(batch.shape[0])
+
+    @staticmethod
+    def _placeholder_logits(n: int) -> np.ndarray:
+        """Return neutral logits when no model is loaded."""
+        return np.zeros((n, 2), dtype=np.float32)