| """Visualization utilities for face detection results.""" |
|
|
| import numpy as np |
| import cv2 |
| from typing import List, Optional, Tuple |
|
|
|
|
| def draw_detections(image: np.ndarray, boxes: np.ndarray, |
| scores: Optional[np.ndarray] = None, |
| track_ids: Optional[np.ndarray] = None, |
| landmarks: Optional[np.ndarray] = None, |
| color: Tuple[int, int, int] = (0, 255, 0), |
| thickness: int = 2) -> np.ndarray: |
| """Draw face detections on image.""" |
| vis = image.copy() |
|
|
| for i in range(len(boxes)): |
| x1, y1, x2, y2 = boxes[i].astype(int) |
| cv2.rectangle(vis, (x1, y1), (x2, y2), color, thickness) |
|
|
| label_parts = [] |
| if track_ids is not None: |
| label_parts.append(f"ID:{track_ids[i]}") |
| if scores is not None: |
| label_parts.append(f"{scores[i]:.2f}") |
| label = ' '.join(label_parts) |
|
|
| if label: |
| (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) |
| cv2.rectangle(vis, (x1, y1 - th - 4), (x1 + tw, y1), color, -1) |
| cv2.putText(vis, label, (x1, y1 - 2), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1) |
|
|
| if landmarks is not None and i < len(landmarks): |
| lmk = landmarks[i] |
| colors_lmk = [(0, 0, 255), (0, 255, 255), (0, 255, 0), |
| (255, 255, 0), (255, 0, 0)] |
| for j in range(5): |
| x, y = int(lmk[j*2]), int(lmk[j*2+1]) |
| if x > 0 and y > 0: |
| cv2.circle(vis, (x, y), 3, colors_lmk[j], -1) |
|
|
| return vis |
|
|
|
|
| def create_comparison_grid(images: List[np.ndarray], titles: List[str], |
| cols: int = 3, cell_size: Tuple[int, int] = (400, 400) |
| ) -> np.ndarray: |
| """Create a grid of images with titles for comparison.""" |
| rows = (len(images) + cols - 1) // cols |
| grid = np.zeros((rows * cell_size[1], cols * cell_size[0], 3), dtype=np.uint8) |
|
|
| for idx, (img, title) in enumerate(zip(images, titles)): |
| r, c = idx // cols, idx % cols |
| resized = cv2.resize(img, cell_size) |
| y1, y2 = r * cell_size[1], (r + 1) * cell_size[1] |
| x1, x2 = c * cell_size[0], (c + 1) * cell_size[0] |
| grid[y1:y2, x1:x2] = resized |
| cv2.putText(grid, title, (x1 + 5, y1 + 20), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2) |
|
|
| return grid |
|
|
