Story 5

src/aggregator.py

@@ -0,0 +1,39 @@
+"""Confidence aggregation strategies for video-level scoring."""
+
+import logging
+from typing import List
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+def aggregate(
+    confidences: List[float],
+    strategy: str = "max",
+    alpha: float = 1.0,
+    beta: float = 0.1,
+    top_k: int = 3,
+) -> float:
+    """Aggregate per-crop confidences into a single video-level score C ∈ [0,1]."""
+    if not confidences:
+        return 0.0
+
+    arr = np.array(confidences, dtype=np.float64)
+
+    if strategy == "max":
+        return float(arr.max())
+
+    if strategy == "weighted_mean":
+        mean = float(arr.mean())
+        var  = float(arr.var())
+        raw  = alpha * mean + beta * var
+        # Clamp to [0,1]; the formula is a linear combination, not a softmax
+        return float(np.clip(raw, 0.0, 1.0))
+
+    if strategy == "top_k_mean":
+        k = max(1, min(top_k, len(arr)))
+        return float(np.partition(arr, -k)[-k:].mean())
+
+    logger.warning("Unknown aggregation strategy '%s', falling back to max", strategy)
+    return float(arr.max())

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)

src/orchestrator.py

@@ -14,11 +14,14 @@ from pathlib import Path
 from typing import Dict, List, Optional
 
 from batcher import DynamicBatcher
+from classifier import FaceClassifier
 from face_detector import Detection, FaceDetector
 from frame_sampler import FrameSampler
 from gpu_manager import GPUMemoryManager
 from prober import VideoProber
 from scratch_manager import ScratchManager
+import aggregator
+import router
 
 logger = logging.getLogger(__name__)
 
@@ -100,6 +103,14 @@ class WorkerPool:
             max_faces_per_frame=int(self._face_detection_config.get("max_faces_per_frame", 10)),
             max_faces_per_video=int(self._face_detection_config.get("max_faces_per_video", 100)),
         )
+        classifier_config = (config or {}).get("classifier", {})
+        self._classifier = FaceClassifier(
+            engine_path=classifier_config.get("model_path", "/models/classifier/classifier.trt"),
+            temperature=float(classifier_config.get("temperature", 1.0)),
+            input_size=int(classifier_config.get("input_size", 224)),
+        )
+        self._aggregation_config = (config or {}).get("aggregation", {})
+        self._routing_config = (config or {}).get("routing", {})
 
     def start(self):
         """Start the worker pool."""
@@ -228,9 +239,40 @@ class WorkerPool:
             )
 
             if not cropped_detections:
-                logger.info("No faces detected for video %s", job.video_id)
+                logger.info("No faces detected for video %s; routing to SKIP", job.video_id)
+                routing_decision = router.SKIP
+                video_confidence = 0.0
+                frame_confidences: List[float] = []
+            else:
+                crop_paths = [d.crop_path for d in cropped_detections if d.crop_path]
+                frame_confidences = self._classifier.classify(
+                    crop_paths, batch_size=self._gpu_manager.current_batch_size
+                )
+                video_confidence = aggregator.aggregate(
+                    frame_confidences,
+                    strategy=self._aggregation_config.get("strategy", "max"),
+                    alpha=float(self._aggregation_config.get("alpha", 1.0)),
+                    beta=float(self._aggregation_config.get("beta", 0.1)),
+                    top_k=int(self._aggregation_config.get("top_k", 3)),
+                )
+                routing_decision = router.route(
+                    video_confidence,
+                    t_high=float(self._routing_config.get("T_high", 0.75)),
+                    t_low=float(self._routing_config.get("T_low", 0.45)),
+                )
 
-            self._complete_job(job, frame_count=len(extracted_frames))
+            logger.info(
+                "Video %s: C=%.4f routing=%s faces=%d frames=%d",
+                job.video_id, video_confidence, routing_decision,
+                len(cropped_detections), len(extracted_frames),
+            )
+
+            self._complete_job(
+                job,
+                frame_count=len(extracted_frames),
+                confidence=video_confidence,
+                routing=routing_decision,
+            )
             scratch_manager.cleanup()
             self._jobs_processed += 1
             return True
@@ -249,24 +291,27 @@ class WorkerPool:
                 self._jobs_failed += 1
                 return False
 
-    def _complete_job(self, job: Job, frame_count: Optional[int] = None):
+    def _complete_job(
-        """Mark a job as completed."""
+        self,
+        job: Job,
+        frame_count: Optional[int] = None,
+        confidence: Optional[float] = None,
+        routing: Optional[str] = None,
+    ):
+        """Mark a job as completed, persisting confidence and routing decision."""
         now = datetime.now(timezone.utc).replace(tzinfo=None)
         job.status = JobStatus.COMPLETED
         job.completed_at = datetime.now(timezone.utc)
 
-        if frame_count is None:
         self.db.execute(
-                """UPDATE videos SET status = 'COMPLETED', updated_at = %s
+            """UPDATE videos
+               SET status = 'COMPLETED',
+                   frame_count = %s,
+                   confidence_score = %s,
+                   routing_decision = %s,
+                   updated_at = %s
                WHERE id = %s""",
-                (now, job.video_id),
+            (frame_count, confidence, routing, now, job.video_id),
-                transaction=True,
-            )
-        else:
-            self.db.execute(
-                """UPDATE videos SET status = 'COMPLETED', frame_count = %s, updated_at = %s
-                   WHERE id = %s""",
-                (frame_count, now, job.video_id),
             transaction=True,
         )
         logger.info("Job completed: %s", job)

src/router.py

@@ -0,0 +1,26 @@
+"""Threshold-based routing of video-level confidence to MATCH / REVIEW / SKIP."""
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+MATCH  = "MATCH"
+REVIEW = "REVIEW"
+SKIP   = "SKIP"
+
+
+def route(
+    confidence: float,
+    t_high: float = 0.75,
+    t_low: float = 0.45,
+) -> str:
+    """Return the routing decision for a video-level confidence score."""
+    if confidence >= t_high:
+        decision = MATCH
+    elif confidence >= t_low:
+        decision = REVIEW
+    else:
+        decision = SKIP
+
+    logger.debug("route: C=%.4f t_high=%.2f t_low=%.2f → %s", confidence, t_high, t_low, decision)
+    return decision

tests/test_story_05.py

@@ -0,0 +1,83 @@
+import sys
+import unittest
+from pathlib import Path
+
+import numpy as np
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
+
+import aggregator
+import router
+from classifier import FaceClassifier, calibrated_softmax
+
+
+class Story05ClassifierTests(unittest.TestCase):
+    def test_calibrated_softmax_sums_to_one(self):
+        logits = np.array([[2.0, 1.0], [-1.0, 3.0]], dtype=np.float32)
+        probs = calibrated_softmax(logits, temperature=1.0)
+        np.testing.assert_allclose(probs.sum(axis=1), [1.0, 1.0], atol=1e-6)
+
+    def test_temperature_scaling_raises_lower_confidence_entropy(self):
+        logits = np.array([[2.0, 0.5]], dtype=np.float32)
+        sharp = calibrated_softmax(logits, temperature=0.5)
+        soft  = calibrated_softmax(logits, temperature=2.0)
+        # higher temperature → softer distribution (target class prob moves toward 0.5)
+        self.assertGreater(sharp[0, 0], soft[0, 0])
+
+    def test_classifier_placeholder_returns_neutral_probability(self):
+        clf = FaceClassifier(engine_path="/nonexistent/model.trt", temperature=1.0)
+        probs = clf.classify([])
+        self.assertEqual(probs, [])
+
+    def test_classifier_placeholder_single_crop_returns_half(self):
+        import tempfile
+        from PIL import Image
+
+        clf = FaceClassifier(engine_path="/nonexistent/model.trt", temperature=1.0)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            crop = Path(tmpdir) / "crop.jpg"
+            Image.new("RGB", (224, 224)).save(crop)
+            probs = clf.classify([str(crop)])
+        # placeholder logits are all zeros → softmax → 0.5 for each class
+        self.assertAlmostEqual(probs[0], 0.5, places=5)
+
+
+class Story05AggregatorTests(unittest.TestCase):
+    def test_max_strategy(self):
+        self.assertAlmostEqual(aggregator.aggregate([0.3, 0.8, 0.6], strategy="max"), 0.8)
+
+    def test_empty_confidences_returns_zero(self):
+        self.assertEqual(aggregator.aggregate([], strategy="max"), 0.0)
+
+    def test_top_k_mean(self):
+        result = aggregator.aggregate([0.1, 0.9, 0.5, 0.8], strategy="top_k_mean", top_k=2)
+        self.assertAlmostEqual(result, (0.9 + 0.8) / 2, places=5)
+
+    def test_weighted_mean_clamps_to_unit_interval(self):
+        result = aggregator.aggregate([1.0, 1.0], strategy="weighted_mean", alpha=100.0, beta=0.0)
+        self.assertLessEqual(result, 1.0)
+        self.assertGreaterEqual(result, 0.0)
+
+
+class Story05RouterTests(unittest.TestCase):
+    def test_match_at_high_threshold(self):
+        self.assertEqual(router.route(0.75), router.MATCH)
+
+    def test_review_between_thresholds(self):
+        self.assertEqual(router.route(0.60), router.REVIEW)
+
+    def test_skip_below_low_threshold(self):
+        self.assertEqual(router.route(0.44), router.SKIP)
+
+    def test_inclusive_high_threshold_boundary(self):
+        self.assertEqual(router.route(0.75, t_high=0.75, t_low=0.45), router.MATCH)
+
+    def test_inclusive_low_threshold_boundary(self):
+        self.assertEqual(router.route(0.45, t_high=0.75, t_low=0.45), router.REVIEW)
+
+    def test_no_faces_zero_confidence_routes_skip(self):
+        self.assertEqual(router.route(0.0), router.SKIP)
+
+
+if __name__ == "__main__":
+    unittest.main()