After story 4

2026-08-10 15:22:22 -04:00
parent 4fdafeef91
commit 612fdd5789
6 changed files with 707 additions and 0 deletions
@@ -0,0 +1,68 @@
+"""Dynamic batching utilities for face detection inference."""
+
+import logging
+import time
+from collections import deque
+from typing import Callable, Deque, Generic, List, Optional, TypeVar
+
+from gpu_manager import GPUMemoryManager
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+class DynamicBatcher(Generic[T, R]):
+    """Accumulate items into batches and flush based on size or timeout."""
+
+    def __init__(
+        self,
+        process_fn: Callable[[List[T]], List[R]],
+        gpu_manager: Optional[GPUMemoryManager] = None,
+        max_batch_size: int = 16,
+        batch_timeout_ms: float = 100.0,
+        min_batch_size: int = 1,
+    ):
+        self.process_fn = process_fn
+        self.gpu_manager = gpu_manager
+        self.max_batch_size = max_batch_size
+        self.batch_timeout_ms = batch_timeout_ms
+        self.min_batch_size = min_batch_size
+        self._queue: Deque[T] = deque()
+        self._last_flush = time.monotonic()
+
+    def add(self, item: T) -> List[R]:
+        """Add an item and return any results if a batch was flushed."""
+        self._queue.append(item)
+        if len(self._queue) >= self._current_batch_size():
+            return self.flush()
+
+        elapsed_ms = (time.monotonic() - self._last_flush) * 1000
+        if elapsed_ms >= self.batch_timeout_ms and len(self._queue) >= self.min_batch_size:
+            return self.flush()
+        return []
+
+    def flush(self) -> List[R]:
+        """Flush all queued items through the processor."""
+        if not self._queue:
+            return []
+
+        batch = [self._queue.popleft() for _ in range(min(len(self._queue), self._current_batch_size()))]
+        results = self.process_fn(batch)
+        self._last_flush = time.monotonic()
+
+        if self.gpu_manager is not None:
+            self.gpu_manager.adjust_batch_size()
+            self.gpu_manager.empty_cache()
+
+        return results
+
+    def _current_batch_size(self) -> int:
+        if self.gpu_manager is not None:
+            return min(self.gpu_manager.current_batch_size, self.max_batch_size)
+        return self.max_batch_size
+
+    @property
+    def queued_count(self) -> int:
+        return len(self._queue)
@@ -0,0 +1,104 @@
+"""Export a YOLOv8n face-detection model to ONNX and build a TensorRT engine."""
+
+import argparse
+import json
+import logging
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+def export_onnx(output_dir: Path, input_size: int = 640) -> Path:
+    """Export YOLOv8n to ONNX with the required input shape."""
+    output_dir.mkdir(parents=True, exist_ok=True)
+    onnx_path = output_dir / "face_detector.onnx"
+    metadata_path = output_dir / "model.json"
+
+    try:
+        from ultralytics import YOLO
+
+        model = YOLO("yolov8n.pt")
+        model.export(
+            format="onnx",
+            imgsz=input_size,
+            half=False,
+            simplify=True,
+            dynamic=False,
+        )
+        exported = Path("yolov8n.onnx")
+        if exported.exists():
+            exported.rename(onnx_path)
+        logger.info("Exported ONNX model to %s", onnx_path)
+    except Exception as exc:
+        logger.error("Failed to export ONNX model: %s", exc)
+        raise
+
+    metadata = {
+        "input_shape": [1, 3, input_size, input_size],
+        "mean": [0.485, 0.456, 0.406],
+        "std": [0.229, 0.224, 0.225],
+        "confidence_threshold": 0.25,
+        "iou_threshold": 0.45,
+    }
+    metadata_path.write_text(json.dumps(metadata, indent=2))
+    return onnx_path
+
+
+def build_tensorrt_engine(onnx_path: Path, output_dir: Path, max_batch_size: int = 32) -> Path:
+    """Build and serialize a TensorRT FP32 engine from an ONNX file."""
+    engine_path = output_dir / "face_detector.trt"
+
+    try:
+        import tensorrt as trt
+
+        logger = trt.Logger(trt.Logger.WARNING)
+        builder = trt.Builder(logger)
+        network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
+        parser = trt.OnnxParser(network, logger)
+
+        onnx_data = onnx_path.read_bytes()
+        if not parser.parse(onnx_data):
+            for error in range(parser.num_errors):
+                logger.error(parser.get_error(error))
+            raise RuntimeError("ONNX parsing failed")
+
+        config = builder.create_builder_config()
+        config.max_workspace_size = 4 << 30  # 4GB
+        config.set_flag(trt.BuilderFlag.FP32)
+
+        profile = builder.create_optimization_profile()
+        input_name = network.get_input(0).name
+        profile.set_shape(
+            input_name,
+            (1, 3, 640, 640),
+            (max_batch_size // 2, 3, 640, 640),
+            (max_batch_size, 3, 640, 640),
+        )
+        config.add_optimization_profile(profile)
+
+        engine = builder.build_engine(network, config)
+        if engine is None:
+            raise RuntimeError("TensorRT engine build failed")
+
+        engine_path.write_bytes(engine.serialize())
+        logger.info("Serialized TensorRT engine to %s", engine_path)
+        return engine_path
+    except Exception as exc:
+        logger.error("Failed to build TensorRT engine: %s", exc)
+        raise
+
+
+def main():
+    logging.basicConfig(level=logging.INFO)
+    parser = argparse.ArgumentParser(description="Export YOLOv8n face detector to ONNX/TensorRT")
+    parser.add_argument("--output-dir", type=Path, default=Path("models/face_detector"))
+    parser.add_argument("--input-size", type=int, default=640)
+    parser.add_argument("--max-batch-size", type=int, default=32)
+    args = parser.parse_args()
+
+    onnx_path = export_onnx(args.output_dir, args.input_size)
+    build_tensorrt_engine(onnx_path, args.output_dir, args.max_batch_size)
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,295 @@
+"""
+Face detection runner.
+
+Wraps a YOLOv8n model exported to ONNX/TensorRT, runs batched inference on
+sampled frames, applies NMS, and produces 224×224 face crops for the classifier.
+"""
+
+import logging
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+class Detection:
+    """A single face detection result."""
+
+    def __init__(
+        self,
+        frame_path: str,
+        x1: float,
+        y1: float,
+        x2: float,
+        y2: float,
+        confidence: float,
+        crop_path: Optional[str] = None,
+    ):
+        self.frame_path = frame_path
+        self.x1 = x1
+        self.y1 = y1
+        self.x2 = x2
+        self.y2 = y2
+        self.confidence = confidence
+        self.crop_path = crop_path
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "frame_path": self.frame_path,
+            "x1": self.x1,
+            "y1": self.y1,
+            "x2": self.x2,
+            "y2": self.y2,
+            "confidence": self.confidence,
+            "crop_path": self.crop_path,
+        }
+
+
+class FaceDetector:
+    """Run batched face detection on video frames."""
+
+    def __init__(
+        self,
+        engine_path: str,
+        input_size: int = 640,
+        confidence_threshold: float = 0.25,
+        iou_threshold: float = 0.45,
+        max_faces_per_frame: int = 10,
+        max_faces_per_video: int = 100,
+        device: str = "cuda",
+    ):
+        self.engine_path = engine_path
+        self.input_size = input_size
+        self.confidence_threshold = confidence_threshold
+        self.iou_threshold = iou_threshold
+        self.max_faces_per_frame = max_faces_per_frame
+        self.max_faces_per_video = max_faces_per_video
+        self.device = device
+        self._session: Optional[Any] = None
+        self._load_model()
+
+    def _load_model(self):
+        """Load the inference backend."""
+        path = Path(self.engine_path)
+        if not path.exists():
+            logger.warning("Face detector 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 face detector from %s", self.engine_path)
+            except Exception as exc:
+                logger.warning("Failed to load ONNX face detector: %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 face detector from %s", self.engine_path)
+            except Exception as exc:
+                logger.warning("Failed to load TensorRT face detector: %s", exc)
+        else:
+            logger.warning("Unsupported face detector format: %s", suffix)
+
+    def detect_faces(
+        self,
+        frame_paths: List[str],
+        batch_size: int = 16,
+    ) -> List[List[Detection]]:
+        """Detect faces in a list of frame image paths."""
+        all_results: List[List[Detection]] = []
+        for i in range(0, len(frame_paths), batch_size):
+            batch = frame_paths[i : i + batch_size]
+            batch_results = self._detect_batch(batch)
+            all_results.extend(batch_results)
+        return all_results
+
+    def _detect_batch(self, frame_paths: List[str]) -> List[List[Detection]]:
+        """Run detection on one batch and return detections per frame."""
+        preprocessed = []
+        for path in frame_paths:
+            try:
+                preprocessed.append(self._preprocess(path))
+            except Exception as exc:
+                logger.warning("Failed to preprocess %s: %s", path, exc)
+                preprocessed.append(np.zeros((3, self.input_size, self.input_size), dtype=np.float32))
+
+        batch_input = np.stack(preprocessed, axis=0)
+
+        outputs = self._infer(batch_input)
+
+        results = []
+        for idx, path in enumerate(frame_paths):
+            try:
+                detections = self._parse_outputs(outputs, idx, path)
+                detections = self._nms(detections)
+                detections = detections[: self.max_faces_per_frame]
+                results.append(detections)
+            except Exception as exc:
+                logger.warning("Failed to parse detections for %s: %s", path, exc)
+                results.append([])
+        return results
+
+    def _preprocess(self, frame_path: str) -> np.ndarray:
+        """Load and normalize a frame for the detector."""
+        from PIL import Image
+
+        image = Image.open(frame_path).convert("RGB")
+        image = image.resize((self.input_size, self.input_size), Image.Resampling.BILINEAR)
+        arr = np.array(image, dtype=np.float32) / 255.0
+        mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+        std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+        arr = (arr - mean) / std
+        return np.transpose(arr, (2, 0, 1))  # HWC -> CHW
+
+    def _infer(self, batch_input: np.ndarray) -> Any:
+        """Run inference on the preprocessed batch."""
+        if self._session is None:
+            return self._placeholder_outputs(batch_input.shape[0])
+
+        try:
+            if hasattr(self._session, "run"):
+                input_name = self._session.get_inputs()[0].name
+                return self._session.run(None, {input_name: batch_input})
+
+            # TensorRT execution
+            import pycuda.driver as cuda
+            import pycuda.autoinit  # noqa: F401
+
+            context = self._session.create_execution_context()
+            output_shape = (batch_input.shape[0], 84, 8400)  # YOLOv8n default shape
+            d_input = cuda.mem_alloc(batch_input.nbytes)
+            d_output = cuda.mem_alloc(np.prod(output_shape) * np.dtype(np.float32).itemsize)
+            stream = cuda.Stream()
+            cuda.memcpy_htod_async(d_input, batch_input, stream)
+            context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
+            output = np.empty(output_shape, dtype=np.float32)
+            cuda.memcpy_dtoh_async(output, d_output, stream)
+            stream.synchronize()
+            return output
+        except Exception as exc:
+            logger.warning("Face detector inference failed: %s", exc)
+            return self._placeholder_outputs(batch_input.shape[0])
+
+    def _placeholder_outputs(self, batch_size: int) -> np.ndarray:
+        """Return an empty output tensor when no model is loaded."""
+        return np.zeros((batch_size, 84, 8400), dtype=np.float32)
+
+    def _parse_outputs(self, outputs: Any, batch_index: int, frame_path: str) -> List[Detection]:
+        """Parse raw inference outputs into Detection objects."""
+        if isinstance(outputs, list):
+            raw = outputs[0][batch_index]  # (84, 8400)
+        else:
+            raw = outputs[batch_index]
+
+        # YOLOv8 output layout: (cx, cy, w, h, cls scores...)
+        raw = raw.T  # (8400, 84)
+        scores = raw[:, 4:].max(axis=1)
+        mask = scores >= self.confidence_threshold
+        candidates = raw[mask]
+        scores = scores[mask]
+
+        detections = []
+        for row, score in zip(candidates, scores):
+            cx, cy, w, h = row[:4]
+            x1 = cx - w / 2
+            y1 = cy - h / 2
+            x2 = cx + w / 2
+            y2 = cy + h / 2
+            detections.append(
+                Detection(
+                    frame_path=frame_path,
+                    x1=float(x1),
+                    y1=float(y1),
+                    x2=float(x2),
+                    y2=float(y2),
+                    confidence=float(score),
+                )
+            )
+        return detections
+
+    def _nms(self, detections: List[Detection]) -> List[Detection]:
+        """Apply greedy Non-Maximum Suppression."""
+        if not detections:
+            return []
+
+        sorted_dets = sorted(detections, key=lambda d: d.confidence, reverse=True)
+        kept: List[Detection] = []
+
+        while sorted_dets:
+            current = sorted_dets.pop(0)
+            kept.append(current)
+            sorted_dets = [
+                det
+                for det in sorted_dets
+                if self._iou(current, det) <= self.iou_threshold
+            ]
+            if len(kept) >= self.max_faces_per_frame:
+                break
+
+        return kept
+
+    @staticmethod
+    def _iou(a: Detection, b: Detection) -> float:
+        """Compute IoU between two detections."""
+        x1 = max(a.x1, b.x1)
+        y1 = max(a.y1, b.y1)
+        x2 = min(a.x2, b.x2)
+        y2 = min(a.y2, b.y2)
+
+        inter_area = max(0, x2 - x1) * max(0, y2 - y1)
+        area_a = (a.x2 - a.x1) * (a.y2 - a.y1)
+        area_b = (b.x2 - b.x1) * (b.y2 - b.y1)
+        union_area = area_a + area_b - inter_area
+        if union_area == 0:
+            return 0.0
+        return inter_area / union_area
+
+    def extract_crops(
+        self,
+        detections: List[Detection],
+        output_dir: str,
+        crop_size: Tuple[int, int] = (224, 224),
+    ) -> List[Detection]:
+        """Extract resized face crops from original frames and update detections."""
+        from PIL import Image
+
+        output_dir_path = Path(output_dir)
+        output_dir_path.mkdir(parents=True, exist_ok=True)
+        cropped: List[Detection] = []
+
+        for idx, det in enumerate(detections):
+            try:
+                image = Image.open(det.frame_path).convert("RGB")
+                width, height = image.size
+
+                x1 = int(max(0, det.x1 * width / self.input_size))
+                y1 = int(max(0, det.y1 * height / self.input_size))
+                x2 = int(min(width, det.x2 * width / self.input_size))
+                y2 = int(min(height, det.y2 * height / self.input_size))
+
+                if x2 <= x1 or y2 <= y1:
+                    image.close()
+                    continue
+
+                crop = image.crop((x1, y1, x2, y2)).resize(crop_size, Image.Resampling.BILINEAR)
+                image.close()
+                crop_path = output_dir_path / f"crop_{int(time.time() * 1000)}_{idx}.jpg"
+                crop.save(crop_path, "JPEG", quality=95)
+                crop.close()
+                det.crop_path = str(crop_path)
+                cropped.append(det)
+            except Exception as exc:
+                logger.warning("Failed to extract crop for %s: %s", det.frame_path, exc)
+
+        return cropped
@@ -0,0 +1,101 @@
+"""GPU memory management and batch-size auto-tuning helpers."""
+
+import logging
+from typing import Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+class GPUMemoryManager:
+    """Track GPU memory and recommend safe batch sizes."""
+
+    def __init__(
+        self,
+        device: int = 0,
+        max_memory_gb: float = 18.0,
+        reduce_threshold_gb: float = 16.0,
+        increase_threshold_gb: float = 10.0,
+        initial_batch_size: int = 16,
+        min_batch_size: int = 1,
+        max_batch_size: int = 32,
+    ):
+        self.device = device
+        self.max_memory_gb = max_memory_gb
+        self.reduce_threshold_gb = reduce_threshold_gb
+        self.increase_threshold_gb = increase_threshold_gb
+        self.batch_size = initial_batch_size
+        self.min_batch_size = min_batch_size
+        self.max_batch_size = max_batch_size
+
+    def get_memory_stats(self) -> Tuple[float, float]:
+        """Return (allocated_gb, reserved_gb) for the managed device."""
+        try:
+            import torch
+
+            if not torch.cuda.is_available():
+                return 0.0, 0.0
+
+            allocated = torch.cuda.memory_allocated(self.device) / (1024**3)
+            reserved = torch.cuda.memory_reserved(self.device) / (1024**3)
+            return allocated, reserved
+        except Exception as exc:
+            logger.debug("Could not query GPU memory: %s", exc)
+            return 0.0, 0.0
+
+    def adjust_batch_size(self) -> int:
+        """Adjust the current batch size based on memory pressure."""
+        allocated, _ = self.get_memory_stats()
+
+        if allocated >= self.reduce_threshold_gb:
+            new_batch_size = max(self.min_batch_size, int(self.batch_size * 0.75))
+            if new_batch_size < self.batch_size:
+                logger.info(
+                    "GPU memory high (%.2f GB), reducing batch size %d -> %d",
+                    allocated,
+                    self.batch_size,
+                    new_batch_size,
+                )
+                self.batch_size = new_batch_size
+        elif allocated <= self.increase_threshold_gb:
+            new_batch_size = min(self.max_batch_size, int(self.batch_size * 1.25))
+            if new_batch_size > self.batch_size:
+                logger.info(
+                    "GPU memory low (%.2f GB), increasing batch size %d -> %d",
+                    allocated,
+                    self.batch_size,
+                    new_batch_size,
+                )
+                self.batch_size = new_batch_size
+
+        return self.batch_size
+
+    def is_memory_critical(self) -> bool:
+        """Return True if allocated memory is close to the hard limit."""
+        allocated, _ = self.get_memory_stats()
+        return allocated >= self.max_memory_gb
+
+    def empty_cache(self):
+        """Try to release unused cached GPU memory."""
+        try:
+            import torch
+
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        except Exception as exc:
+            logger.debug("Could not empty GPU cache: %s", exc)
+
+    @property
+    def current_batch_size(self) -> int:
+        return self.batch_size
+
+
+def get_available_vram_gb(device: int = 0) -> float:
+    """Return total available VRAM in GB, or 0.0 if CUDA is unavailable."""
+    try:
+        import torch
+
+        if not torch.cuda.is_available():
+            return 0.0
+        return torch.cuda.get_device_properties(device).total_memory / (1024**3)
+    except Exception:
+        return 0.0
@@ -13,7 +13,10 @@ from enum import Enum
 from pathlib import Path
 from typing import Dict, List, Optional

+from batcher import DynamicBatcher
+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

@@ -81,6 +84,22 @@ class WorkerPool:
        self._jobs_failed = 0
        self._sampling_config = (config or {}).get("sampling", {})
        self._storage_config = (config or {}).get("storage", {})
+        self._face_detection_config = (config or {}).get("face_detection", {})
+        self._batching_config = (config or {}).get("batching", {})
+        self._gpu_manager = GPUMemoryManager(
+            max_memory_gb=config.get("gpu", {}).get("max_memory_gb", 18.0),
+            reduce_threshold_gb=self._batching_config.get("vram_reduce_threshold_gb", 16.0),
+            increase_threshold_gb=self._batching_config.get("vram_increase_threshold_gb", 10.0),
+            initial_batch_size=self._batching_config.get("max_batch_size", 16),
+        )
+        self._face_detector = FaceDetector(
+            engine_path=self._face_detection_config.get("model_path", "/models/face_detector/face_detector.trt"),
+            input_size=int(self._face_detection_config.get("input_size", 640)),
+            confidence_threshold=float(self._face_detection_config.get("confidence_threshold", 0.25)),
+            iou_threshold=float(self._face_detection_config.get("iou_threshold", 0.45)),
+            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)),
+        )

    def start(self):
        """Start the worker pool."""
@@ -189,6 +208,28 @@ class WorkerPool:
                scratch_manager.cleanup()
                raise RuntimeError("No frames were extracted")

+            # Face detection on sampled frames
+            batch_size = self._gpu_manager.current_batch_size
+            detections_per_frame = self._face_detector.detect_faces(extracted_frames, batch_size=batch_size)
+
+            # Flatten and cap total faces per video
+            all_detections: List[Detection] = []
+            for frame_dets in detections_per_frame:
+                all_detections.extend(frame_dets)
+            all_detections = sorted(all_detections, key=lambda d: d.confidence, reverse=True)
+            all_detections = all_detections[: self._face_detection_config.get("max_faces_per_video", 100)]
+
+            # Extract face crops
+            crop_dir = scratch_manager.frame_dir.parent / "crops"
+            cropped_detections = self._face_detector.extract_crops(
+                all_detections,
+                output_dir=str(crop_dir),
+                crop_size=(224, 224),
+            )
+
+            if not cropped_detections:
+                logger.info("No faces detected for video %s", job.video_id)
+
            self._complete_job(job, frame_count=len(extracted_frames))
            scratch_manager.cleanup()
            self._jobs_processed += 1