Upload utils/visualization.py with huggingface_hub

utils/visualization.py

@@ -0,0 +1,62 @@
+"""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