YOLO
This commit is contained in:
@@ -0,0 +1,371 @@
|
||||
"""
|
||||
YOLOv11 Face Detector using ONNX Runtime with GPU support.
|
||||
|
||||
This module provides high-performance face detection using
|
||||
YOLOv11-face model with CUDA acceleration.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import List, Tuple, Optional
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
|
||||
|
||||
class YOLOFaceDetector:
|
||||
"""
|
||||
YOLOv11 face detector with ONNX Runtime GPU support.
|
||||
|
||||
Features:
|
||||
- CUDA GPU acceleration
|
||||
- High accuracy face detection
|
||||
- NMS for overlapping detections
|
||||
"""
|
||||
|
||||
# Default model path relative to this file
|
||||
DEFAULT_MODEL = "yolov11n-face.onnx"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_path: Optional[str] = None,
|
||||
conf_threshold: float = 0.25,
|
||||
iou_threshold: float = 0.45,
|
||||
input_size: Tuple[int, int] = (640, 640),
|
||||
):
|
||||
"""
|
||||
Initialize the YOLO face detector.
|
||||
|
||||
Args:
|
||||
model_path: Path to ONNX model file. If None, uses default model.
|
||||
conf_threshold: Confidence threshold for detections
|
||||
iou_threshold: IoU threshold for NMS
|
||||
input_size: Model input size (width, height)
|
||||
"""
|
||||
self.conf_threshold = conf_threshold
|
||||
self.iou_threshold = iou_threshold
|
||||
self.input_size = input_size
|
||||
self._session = None
|
||||
self._model_path = model_path
|
||||
|
||||
@property
|
||||
def session(self):
|
||||
"""Lazy-load ONNX Runtime session."""
|
||||
if self._session is None:
|
||||
import onnxruntime as ort
|
||||
|
||||
# Determine model path
|
||||
if self._model_path is None:
|
||||
# Assuming models are in ../models relative to server/detector.py
|
||||
models_dir = Path(__file__).parent.parent / "models"
|
||||
model_path = str(models_dir / self.DEFAULT_MODEL)
|
||||
else:
|
||||
model_path = self._model_path
|
||||
|
||||
if not os.path.exists(model_path):
|
||||
raise FileNotFoundError(f"Model not found: {model_path}")
|
||||
|
||||
# Configure providers (prefer CUDA)
|
||||
providers = []
|
||||
if 'CUDAExecutionProvider' in ort.get_available_providers():
|
||||
providers.append('CUDAExecutionProvider')
|
||||
print("[FaceMask] Using CUDA GPU for inference")
|
||||
providers.append('CPUExecutionProvider')
|
||||
|
||||
# Create session
|
||||
sess_options = ort.SessionOptions()
|
||||
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||||
|
||||
self._session = ort.InferenceSession(
|
||||
model_path,
|
||||
sess_options=sess_options,
|
||||
providers=providers,
|
||||
)
|
||||
|
||||
print(f"[FaceMask] YOLO model loaded: {model_path}")
|
||||
print(f"[FaceMask] Providers: {self._session.get_providers()}")
|
||||
|
||||
return self._session
|
||||
|
||||
def detect(self, frame: np.ndarray) -> List[Tuple[int, int, int, int, float]]:
|
||||
"""
|
||||
Detect faces in a frame.
|
||||
|
||||
Args:
|
||||
frame: BGR image as numpy array (H, W, C)
|
||||
|
||||
Returns:
|
||||
List of detections as (x, y, width, height, confidence)
|
||||
"""
|
||||
import cv2
|
||||
|
||||
original_height, original_width = frame.shape[:2]
|
||||
|
||||
input_tensor = self._preprocess(frame)
|
||||
# print(f"[DEBUG] Input tensor shape: {input_tensor.shape}, Range: [{input_tensor.min():.3f}, {input_tensor.max():.3f}]", flush=True)
|
||||
|
||||
# Run inference
|
||||
input_name = self.session.get_inputs()[0].name
|
||||
outputs = self.session.run(None, {input_name: input_tensor})
|
||||
|
||||
raw_output = outputs[0]
|
||||
# print(f"[DEBUG] Raw output shape: {raw_output.shape}, Range: [{raw_output.min():.3f}, {raw_output.max():.3f}]", flush=True)
|
||||
|
||||
# Postprocess
|
||||
detections = self._postprocess(
|
||||
raw_output,
|
||||
original_width,
|
||||
original_height,
|
||||
)
|
||||
# print(f"[DEBUG] Detections found: {len(detections)}", flush=True)
|
||||
|
||||
return detections
|
||||
|
||||
def _preprocess(self, frame: np.ndarray) -> np.ndarray:
|
||||
"""Preprocess frame for YOLO input with letterbox resizing."""
|
||||
import cv2
|
||||
|
||||
# Letterbox resize
|
||||
shape = frame.shape[:2] # current shape [height, width]
|
||||
new_shape = self.input_size
|
||||
|
||||
# Scale ratio (new / old)
|
||||
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
|
||||
|
||||
# Compute padding
|
||||
ratio = r, r # width, height ratios
|
||||
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
|
||||
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
|
||||
|
||||
dw /= 2 # divide padding into 2 sides
|
||||
dh /= 2
|
||||
|
||||
if shape[::-1] != new_unpad: # resize
|
||||
frame = cv2.resize(frame, new_unpad, interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
|
||||
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
|
||||
|
||||
# Add border
|
||||
frame = cv2.copyMakeBorder(frame, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
|
||||
|
||||
# Store metadata for postprocessing
|
||||
self._last_letterbox_meta = {'ratio': ratio, 'dwdh': (dw, dh)}
|
||||
|
||||
# Convert BGR to RGB
|
||||
rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
|
||||
# Normalize to [0, 1]
|
||||
normalized = rgb.astype(np.float32) / 255.0
|
||||
|
||||
# Transpose to CHW format
|
||||
transposed = np.transpose(normalized, (2, 0, 1))
|
||||
|
||||
# Add batch dimension
|
||||
batched = np.expand_dims(transposed, axis=0)
|
||||
|
||||
return batched
|
||||
|
||||
def _postprocess(
|
||||
self,
|
||||
output: np.ndarray,
|
||||
original_width: int,
|
||||
original_height: int,
|
||||
) -> List[Tuple[int, int, int, int, float]]:
|
||||
"""
|
||||
Postprocess YOLO output to get detections.
|
||||
"""
|
||||
# Output shape: [1, num_detections, 5+] where 5 = x_center, y_center, w, h, conf
|
||||
|
||||
# Handle different output formats
|
||||
if output.shape[1] < output.shape[2]:
|
||||
# Format: [1, 5+, num_detections] - transpose
|
||||
output = np.transpose(output[0], (1, 0))
|
||||
else:
|
||||
output = output[0]
|
||||
|
||||
# Debug confidence stats
|
||||
# if output.shape[1] >= 5:
|
||||
# max_conf = output[:, 4].max()
|
||||
# print(f"[DEBUG] Max confidence in raw output: {max_conf:.4f}", flush=True)
|
||||
|
||||
# Filter by confidence
|
||||
confidences = output[:, 4]
|
||||
mask = confidences > self.conf_threshold
|
||||
filtered = output[mask]
|
||||
|
||||
if len(filtered) == 0:
|
||||
return []
|
||||
|
||||
# Get letterbox metadata
|
||||
if hasattr(self, '_last_letterbox_meta') and self._last_letterbox_meta:
|
||||
ratio = self._last_letterbox_meta['ratio']
|
||||
dw, dh = self._last_letterbox_meta['dwdh']
|
||||
|
||||
# Extract coordinates
|
||||
x_center = filtered[:, 0]
|
||||
y_center = filtered[:, 1]
|
||||
width = filtered[:, 2]
|
||||
height = filtered[:, 3]
|
||||
confidences = filtered[:, 4]
|
||||
|
||||
# Convert center to corner
|
||||
x1 = x_center - width / 2
|
||||
y1 = y_center - height / 2
|
||||
x2 = x_center + width / 2
|
||||
y2 = y_center + height / 2
|
||||
|
||||
# Adjust for letterbox padding
|
||||
x1 -= dw
|
||||
y1 -= dh
|
||||
x2 -= dw
|
||||
y2 -= dh
|
||||
|
||||
# Adjust for resizing
|
||||
x1 /= ratio[0]
|
||||
y1 /= ratio[1]
|
||||
x2 /= ratio[0]
|
||||
y2 /= ratio[1]
|
||||
|
||||
# Clip to image bounds
|
||||
x1 = np.clip(x1, 0, original_width)
|
||||
y1 = np.clip(y1, 0, original_height)
|
||||
x2 = np.clip(x2, 0, original_width)
|
||||
y2 = np.clip(y2, 0, original_height)
|
||||
|
||||
# Convert back to x, y, w, h
|
||||
final_x = x1
|
||||
final_y = y1
|
||||
final_w = x2 - x1
|
||||
final_h = y2 - y1
|
||||
|
||||
else:
|
||||
# Fallback for non-letterbox (legacy)
|
||||
scale_x = original_width / self.input_size[0]
|
||||
scale_y = original_height / self.input_size[1]
|
||||
|
||||
x_center = filtered[:, 0] * scale_x
|
||||
y_center = filtered[:, 1] * scale_y
|
||||
width = filtered[:, 2] * scale_x
|
||||
height = filtered[:, 3] * scale_y
|
||||
confidences = filtered[:, 4]
|
||||
|
||||
final_x = x_center - width / 2
|
||||
final_y = y_center - height / 2
|
||||
final_w = width
|
||||
final_h = height
|
||||
|
||||
# Apply NMS
|
||||
boxes = np.stack([final_x, final_y, final_w, final_h], axis=1)
|
||||
indices = self._nms(boxes, confidences, self.iou_threshold)
|
||||
|
||||
# Format output
|
||||
detections = []
|
||||
for i in indices:
|
||||
x = int(final_x[i])
|
||||
y = int(final_y[i])
|
||||
w = int(final_w[i])
|
||||
h = int(final_h[i])
|
||||
conf = float(confidences[i])
|
||||
detections.append((x, y, w, h, conf))
|
||||
|
||||
return detections
|
||||
|
||||
def _nms(
|
||||
self,
|
||||
boxes: np.ndarray,
|
||||
scores: np.ndarray,
|
||||
iou_threshold: float,
|
||||
) -> List[int]:
|
||||
"""Non-Maximum Suppression."""
|
||||
x1 = boxes[:, 0]
|
||||
y1 = boxes[:, 1]
|
||||
x2 = x1 + boxes[:, 2]
|
||||
y2 = y1 + boxes[:, 3]
|
||||
|
||||
areas = boxes[:, 2] * boxes[:, 3]
|
||||
order = scores.argsort()[::-1]
|
||||
|
||||
keep = []
|
||||
while len(order) > 0:
|
||||
i = order[0]
|
||||
keep.append(i)
|
||||
|
||||
if len(order) == 1:
|
||||
break
|
||||
|
||||
xx1 = np.maximum(x1[i], x1[order[1:]])
|
||||
yy1 = np.maximum(y1[i], y1[order[1:]])
|
||||
xx2 = np.minimum(x2[i], x2[order[1:]])
|
||||
yy2 = np.minimum(y2[i], y2[order[1:]])
|
||||
|
||||
w = np.maximum(0, xx2 - xx1)
|
||||
h = np.maximum(0, yy2 - yy1)
|
||||
inter = w * h
|
||||
|
||||
iou = inter / (areas[i] + areas[order[1:]] - inter)
|
||||
|
||||
inds = np.where(iou <= iou_threshold)[0]
|
||||
order = order[inds + 1]
|
||||
|
||||
return keep
|
||||
|
||||
def generate_mask(
|
||||
self,
|
||||
frame_shape: Tuple[int, int, int],
|
||||
detections: List[Tuple[int, int, int, int, float]],
|
||||
mask_scale: float = 1.5,
|
||||
feather_radius: int = 20,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Generate a mask image from face detections.
|
||||
|
||||
Args:
|
||||
frame_shape: Shape of the original frame (height, width, channels)
|
||||
detections: List of face detections (x, y, w, h, conf)
|
||||
mask_scale: Scale factor for mask region
|
||||
feather_radius: Radius for edge feathering
|
||||
|
||||
Returns:
|
||||
Grayscale mask image (white = blur, black = keep)
|
||||
"""
|
||||
import cv2
|
||||
|
||||
height, width = frame_shape[:2]
|
||||
mask = np.zeros((height, width), dtype=np.uint8)
|
||||
|
||||
for (x, y, w, h, conf) in detections:
|
||||
# Scale the bounding box
|
||||
center_x = x + w // 2
|
||||
center_y = y + h // 2
|
||||
|
||||
scaled_w = int(w * mask_scale)
|
||||
scaled_h = int(h * mask_scale)
|
||||
|
||||
# Draw ellipse for natural face shape
|
||||
cv2.ellipse(
|
||||
mask,
|
||||
(center_x, center_y),
|
||||
(scaled_w // 2, scaled_h // 2),
|
||||
0, # angle
|
||||
0, 360, # arc
|
||||
255, # color (white)
|
||||
-1, # filled
|
||||
)
|
||||
|
||||
# Apply Gaussian blur for feathering
|
||||
if feather_radius > 0 and len(detections) > 0:
|
||||
kernel_size = feather_radius * 2 + 1
|
||||
mask = cv2.GaussianBlur(mask, (kernel_size, kernel_size), 0)
|
||||
|
||||
return mask
|
||||
|
||||
|
||||
# Singleton instance
|
||||
_detector: Optional[YOLOFaceDetector] = None
|
||||
|
||||
|
||||
def get_detector(**kwargs) -> YOLOFaceDetector:
|
||||
"""Get or create the global YOLO detector instance."""
|
||||
global _detector
|
||||
if _detector is None:
|
||||
_detector = YOLOFaceDetector(**kwargs)
|
||||
return _detector
|
||||
+180
@@ -0,0 +1,180 @@
|
||||
"""
|
||||
Face Detection Inference Server.
|
||||
|
||||
This FastAPI application runs in a separate process to handle
|
||||
GPU-accelerated face detection using ONNX Runtime.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import uuid
|
||||
import queue
|
||||
import traceback
|
||||
from typing import Dict, Optional, List
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi import FastAPI, HTTPException, BackgroundTasks
|
||||
from pydantic import BaseModel
|
||||
import uvicorn
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
# Add project root to path for imports if needed
|
||||
sys.path.append(str(Path(__file__).parent.parent))
|
||||
|
||||
from server.detector import YOLOFaceDetector, get_detector
|
||||
|
||||
app = FastAPI(title="Face Mask Inference Server")
|
||||
|
||||
# Task storage
|
||||
class TaskStatus:
|
||||
PENDING = "pending"
|
||||
PROCESSING = "processing"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
class Task(BaseModel):
|
||||
id: str
|
||||
status: str
|
||||
progress: int = 0
|
||||
total: int = 0
|
||||
message: Optional[str] = None
|
||||
result_path: Optional[str] = None
|
||||
|
||||
# In-memory storage
|
||||
tasks: Dict[str, Task] = {}
|
||||
cancel_events: Dict[str, threading.Event] = {}
|
||||
|
||||
class GenerateRequest(BaseModel):
|
||||
video_path: str
|
||||
output_dir: str
|
||||
start_frame: int
|
||||
end_frame: int
|
||||
conf_threshold: float = 0.5
|
||||
iou_threshold: float = 0.45
|
||||
mask_scale: float = 1.5
|
||||
|
||||
def process_video_task(task_id: str, req: GenerateRequest):
|
||||
"""Background task to process video."""
|
||||
try:
|
||||
tasks[task_id].status = TaskStatus.PROCESSING
|
||||
cancel_event = cancel_events.get(task_id)
|
||||
|
||||
# Verify video exists
|
||||
if not os.path.exists(req.video_path):
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
tasks[task_id].message = f"Video not found: {req.video_path}"
|
||||
return
|
||||
|
||||
# Initialize detector (will load model on first run)
|
||||
print(f"Loading detector for task {task_id}...")
|
||||
detector = get_detector(
|
||||
conf_threshold=req.conf_threshold,
|
||||
iou_threshold=req.iou_threshold
|
||||
)
|
||||
# Ensure session is loaded
|
||||
_ = detector.session
|
||||
|
||||
# Open video
|
||||
cap = cv2.VideoCapture(req.video_path)
|
||||
if not cap.isOpened():
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
tasks[task_id].message = "Failed to open video"
|
||||
return
|
||||
|
||||
# Determine frame range
|
||||
total_video_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
end_frame = min(req.end_frame, total_video_frames - 1)
|
||||
frames_to_process = end_frame - req.start_frame + 1
|
||||
|
||||
tasks[task_id].total = frames_to_process
|
||||
|
||||
# Ensure output directory exists
|
||||
os.makedirs(req.output_dir, exist_ok=True)
|
||||
|
||||
print(f"Starting processing: {req.video_path} ({frames_to_process} frames)")
|
||||
|
||||
# Process loop
|
||||
current_count = 0
|
||||
for frame_idx in range(req.start_frame, end_frame + 1):
|
||||
if cancel_event and cancel_event.is_set():
|
||||
tasks[task_id].status = TaskStatus.CANCELLED
|
||||
tasks[task_id].message = "Cancelled by user"
|
||||
break
|
||||
|
||||
# Read frame
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
|
||||
ret, frame = cap.read()
|
||||
|
||||
if ret:
|
||||
# Detect
|
||||
detections = detector.detect(frame)
|
||||
|
||||
# Generate mask
|
||||
mask = detector.generate_mask(
|
||||
frame.shape,
|
||||
detections,
|
||||
mask_scale=req.mask_scale
|
||||
)
|
||||
|
||||
# Save
|
||||
mask_filename = f"mask_{current_count:06d}.png" # Note: using relative index for filename
|
||||
mask_path = os.path.join(req.output_dir, mask_filename)
|
||||
cv2.imwrite(mask_path, mask)
|
||||
|
||||
# Update progress
|
||||
current_count += 1
|
||||
tasks[task_id].progress = current_count
|
||||
|
||||
cap.release()
|
||||
|
||||
if tasks[task_id].status == TaskStatus.PROCESSING:
|
||||
tasks[task_id].status = TaskStatus.COMPLETED
|
||||
tasks[task_id].result_path = req.output_dir
|
||||
tasks[task_id].message = "Processing completed successfully"
|
||||
print(f"Task {task_id} completed.")
|
||||
|
||||
except Exception as e:
|
||||
tasks[task_id].status = TaskStatus.FAILED
|
||||
tasks[task_id].message = str(e)
|
||||
print(f"Error in task {task_id}: {e}")
|
||||
traceback.print_exc()
|
||||
finally:
|
||||
# Cleanup
|
||||
if task_id in cancel_events:
|
||||
del cancel_events[task_id]
|
||||
|
||||
@app.get("/status")
|
||||
def get_status():
|
||||
return {"status": "running", "gpu_available": True} # TODO: check GPU
|
||||
|
||||
@app.post("/generate", response_model=Task)
|
||||
def generate_mask_endpoint(req: GenerateRequest, background_tasks: BackgroundTasks):
|
||||
task_id = str(uuid.uuid4())
|
||||
task = Task(id=task_id, status=TaskStatus.PENDING)
|
||||
tasks[task_id] = task
|
||||
cancel_events[task_id] = threading.Event()
|
||||
|
||||
background_tasks.add_task(process_video_task, task_id, req)
|
||||
return task
|
||||
|
||||
@app.get("/tasks/{task_id}", response_model=Task)
|
||||
def get_task(task_id: str):
|
||||
if task_id not in tasks:
|
||||
raise HTTPException(status_code=404, detail="Task not found")
|
||||
return tasks[task_id]
|
||||
|
||||
@app.post("/tasks/{task_id}/cancel")
|
||||
def cancel_task(task_id: str):
|
||||
if task_id not in tasks:
|
||||
raise HTTPException(status_code=404, detail="Task not found")
|
||||
|
||||
if task_id in cancel_events:
|
||||
cancel_events[task_id].set()
|
||||
|
||||
return {"message": "Cancellation requested"}
|
||||
|
||||
if __name__ == "__main__":
|
||||
uvicorn.run(app, host="127.0.0.1", port=8181)
|
||||
Reference in New Issue
Block a user