After story 4
This commit is contained in:
@@ -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 pathlib import Path
|
||||||
from typing import Dict, List, Optional
|
from typing import Dict, List, Optional
|
||||||
|
|
||||||
|
from batcher import DynamicBatcher
|
||||||
|
from face_detector import Detection, FaceDetector
|
||||||
from frame_sampler import FrameSampler
|
from frame_sampler import FrameSampler
|
||||||
|
from gpu_manager import GPUMemoryManager
|
||||||
from prober import VideoProber
|
from prober import VideoProber
|
||||||
from scratch_manager import ScratchManager
|
from scratch_manager import ScratchManager
|
||||||
|
|
||||||
@@ -81,6 +84,22 @@ class WorkerPool:
|
|||||||
self._jobs_failed = 0
|
self._jobs_failed = 0
|
||||||
self._sampling_config = (config or {}).get("sampling", {})
|
self._sampling_config = (config or {}).get("sampling", {})
|
||||||
self._storage_config = (config or {}).get("storage", {})
|
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):
|
def start(self):
|
||||||
"""Start the worker pool."""
|
"""Start the worker pool."""
|
||||||
@@ -189,6 +208,28 @@ class WorkerPool:
|
|||||||
scratch_manager.cleanup()
|
scratch_manager.cleanup()
|
||||||
raise RuntimeError("No frames were extracted")
|
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))
|
self._complete_job(job, frame_count=len(extracted_frames))
|
||||||
scratch_manager.cleanup()
|
scratch_manager.cleanup()
|
||||||
self._jobs_processed += 1
|
self._jobs_processed += 1
|
||||||
|
|||||||
@@ -0,0 +1,98 @@
|
|||||||
|
import sys
|
||||||
|
import tempfile
|
||||||
|
import unittest
|
||||||
|
from pathlib import Path
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
|
||||||
|
|
||||||
|
from batcher import DynamicBatcher
|
||||||
|
from face_detector import Detection, FaceDetector
|
||||||
|
from gpu_manager import GPUMemoryManager
|
||||||
|
|
||||||
|
|
||||||
|
class Story04FaceDetectionTests(unittest.TestCase):
|
||||||
|
def test_gpu_manager_respects_memory_thresholds(self):
|
||||||
|
manager = GPUMemoryManager(
|
||||||
|
max_memory_gb=18.0,
|
||||||
|
reduce_threshold_gb=16.0,
|
||||||
|
increase_threshold_gb=10.0,
|
||||||
|
initial_batch_size=16,
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch.object(manager, "get_memory_stats", return_value=(17.0, 18.0)):
|
||||||
|
manager.adjust_batch_size()
|
||||||
|
self.assertLess(manager.current_batch_size, 16)
|
||||||
|
|
||||||
|
with patch.object(manager, "get_memory_stats", return_value=(8.0, 10.0)):
|
||||||
|
manager.adjust_batch_size()
|
||||||
|
self.assertGreater(manager.current_batch_size, manager.min_batch_size)
|
||||||
|
|
||||||
|
def test_dynamic_batcher_flushes_when_batch_full(self):
|
||||||
|
processed_batches = []
|
||||||
|
|
||||||
|
def process_fn(batch):
|
||||||
|
processed_batches.append(batch)
|
||||||
|
return [len(batch)]
|
||||||
|
|
||||||
|
batcher = DynamicBatcher(
|
||||||
|
process_fn=process_fn,
|
||||||
|
max_batch_size=3,
|
||||||
|
batch_timeout_ms=10000.0,
|
||||||
|
)
|
||||||
|
|
||||||
|
for i in range(5):
|
||||||
|
batcher.add(i)
|
||||||
|
|
||||||
|
self.assertEqual(len(processed_batches), 1)
|
||||||
|
self.assertEqual(processed_batches[0], [0, 1, 2])
|
||||||
|
self.assertEqual(batcher.queued_count, 2)
|
||||||
|
|
||||||
|
def test_face_detector_nms_removes_overlapping_boxes(self):
|
||||||
|
detector = FaceDetector(
|
||||||
|
engine_path="/nonexistent/model.trt",
|
||||||
|
confidence_threshold=0.1,
|
||||||
|
iou_threshold=0.45,
|
||||||
|
max_faces_per_frame=10,
|
||||||
|
)
|
||||||
|
|
||||||
|
duplicates = [
|
||||||
|
Detection("frame.jpg", 10, 10, 50, 50, 0.9),
|
||||||
|
Detection("frame.jpg", 12, 12, 48, 48, 0.8),
|
||||||
|
Detection("frame.jpg", 100, 100, 150, 150, 0.75),
|
||||||
|
]
|
||||||
|
|
||||||
|
kept = detector._nms(duplicates)
|
||||||
|
self.assertEqual(len(kept), 2)
|
||||||
|
self.assertAlmostEqual(kept[0].confidence, 0.9, places=5)
|
||||||
|
|
||||||
|
def test_face_detector_extracts_and_resizes_crops(self):
|
||||||
|
detector = FaceDetector(
|
||||||
|
engine_path="/nonexistent/model.trt",
|
||||||
|
input_size=640,
|
||||||
|
confidence_threshold=0.25,
|
||||||
|
)
|
||||||
|
|
||||||
|
with tempfile.TemporaryDirectory() as tmpdir:
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
frame_path = Path(tmpdir) / "frame.jpg"
|
||||||
|
image = Image.new("RGB", (640, 480), color=(100, 150, 200))
|
||||||
|
image.save(frame_path)
|
||||||
|
|
||||||
|
detections = [
|
||||||
|
Detection(str(frame_path), 0, 0, 640, 480, 0.95),
|
||||||
|
]
|
||||||
|
cropped = detector.extract_crops(detections, output_dir=tmpdir, crop_size=(224, 224))
|
||||||
|
|
||||||
|
self.assertEqual(len(cropped), 1)
|
||||||
|
self.assertTrue(Path(cropped[0].crop_path).exists())
|
||||||
|
|
||||||
|
with Image.open(cropped[0].crop_path) as crop:
|
||||||
|
self.assertEqual(crop.size, (224, 224))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
Reference in New Issue
Block a user