From ac8fc5ed7cbc67baf0bb95decce36736a2035677 Mon Sep 17 00:00:00 2001 From: Ryan Shpeherd Date: Mon, 10 Aug 2026 15:51:25 -0400 Subject: [PATCH] Story 5 --- src/aggregator.py | 39 ++++++++++++ src/classifier.py | 138 +++++++++++++++++++++++++++++++++++++++++ src/orchestrator.py | 81 ++++++++++++++++++------ src/router.py | 26 ++++++++ tests/test_story_05.py | 83 +++++++++++++++++++++++++ 5 files changed, 349 insertions(+), 18 deletions(-) create mode 100644 src/aggregator.py create mode 100644 src/classifier.py create mode 100644 src/router.py create mode 100644 tests/test_story_05.py diff --git a/src/aggregator.py b/src/aggregator.py new file mode 100644 index 0000000..3bb8233 --- /dev/null +++ b/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()) diff --git a/src/classifier.py b/src/classifier.py new file mode 100644 index 0000000..ec089f9 --- /dev/null +++ b/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) diff --git a/src/orchestrator.py b/src/orchestrator.py index 85dba0e..dcbecae 100644 --- a/src/orchestrator.py +++ b/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,26 +291,29 @@ class WorkerPool: self._jobs_failed += 1 return False - def _complete_job(self, job: Job, frame_count: Optional[int] = None): - """Mark a job as completed.""" + def _complete_job( + 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 - WHERE id = %s""", - (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, - ) + self.db.execute( + """UPDATE videos + SET status = 'COMPLETED', + frame_count = %s, + confidence_score = %s, + routing_decision = %s, + updated_at = %s + WHERE id = %s""", + (frame_count, confidence, routing, now, job.video_id), + transaction=True, + ) logger.info("Job completed: %s", job) def _fail_job(self, job: Job): diff --git a/src/router.py b/src/router.py new file mode 100644 index 0000000..0d6199b --- /dev/null +++ b/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 diff --git a/tests/test_story_05.py b/tests/test_story_05.py new file mode 100644 index 0000000..13b26a9 --- /dev/null +++ b/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()