From e2d19cfbf088d4aaa428d148e0fb29a301c9cd62 Mon Sep 17 00:00:00 2001 From: fiatcode Date: Fri, 27 Feb 2026 21:38:27 +0700 Subject: [PATCH] refactor: update bbox format from (x,y,w,h) to (x1,y1,x2,y2) --- src/faceblur/output.py | 109 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 109 insertions(+) create mode 100644 src/faceblur/output.py diff --git a/src/faceblur/output.py b/src/faceblur/output.py new file mode 100644 index 0000000..3921a5e --- /dev/null +++ b/src/faceblur/output.py @@ -0,0 +1,109 @@ +"""Output generation module.""" + +import colorsys +from pathlib import Path +from typing import List, Dict, Tuple + +import cv2 +import numpy as np + +from .video import Frame +from .detect import FaceData +from .cluster import Cluster + + +def generate_output( + frames: List[Frame], faces: List[FaceData], clusters: List[Cluster], output_dir: str +) -> None: + """Generate output with bounding boxes and face crops. + + Args: + frames: List of extracted frames + faces: All detected faces + clusters: Clustered faces + output_dir: Output directory + """ + output_dir = Path(output_dir) + + face_to_cluster: Dict[int, int] = {} + for cluster in clusters: + for face in cluster.faces: + face_to_cluster[face.id] = cluster.id + + faces_by_frame: Dict[int, List[FaceData]] = {} + for face in faces: + if face.frame_index not in faces_by_frame: + faces_by_frame[face.frame_index] = [] + faces_by_frame[face.frame_index].append(face) + + frames_dir = output_dir / "frames" + frames_dir.mkdir(parents=True, exist_ok=True) + + for frame in frames: + frame_faces = faces_by_frame.get(frame.index, []) + _draw_frame_with_boxes(frame.path, frame_faces, face_to_cluster, frames_dir) + + for i, cluster in enumerate(clusters): + if cluster.id == -1: + cluster_dir = output_dir / "unclustered" + else: + cluster_dir = output_dir / f"cluster_{i:02d}" + cluster_dir.mkdir(parents=True, exist_ok=True) + + for face in cluster.faces: + _extract_face_crop(face, cluster_dir) + + +def _get_cluster_color(cluster_id: int) -> Tuple[int, int, int]: + """Generate a consistent color for a cluster.""" + if cluster_id < 0: + return (128, 128, 128) + + hue = (cluster_id * 0.618033988749895) % 1.0 + rgb = colorsys.hsv_to_rgb(hue, 0.8, 0.9) + return (int(rgb[2] * 255), int(rgb[1] * 255), int(rgb[0] * 255)) + + +def _draw_frame_with_boxes( + frame_path: Path, + faces: List[FaceData], + face_to_cluster: Dict[int, int], + output_dir: Path, +) -> None: + """Draw bounding boxes on a frame.""" + image = cv2.imread(str(frame_path)) + if image is None: + return + + for face in faces: + x1, y1, x2, y2 = face.bbox + cluster_id = face_to_cluster.get(face.id, -1) + color = _get_cluster_color(cluster_id) + + cv2.rectangle(image, (x1, y1), (x2, y2), color, 2) + + if cluster_id >= 0: + label = f"C{cluster_id}" + cv2.putText( + image, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1 + ) + + output_path = output_dir / f"frame_{frame_path.stem.split('_')[1]}.jpg" + cv2.imwrite(str(output_path), image) + + +def _extract_face_crop(face: FaceData, output_dir: Path) -> None: + """Extract and save a face crop.""" + image = cv2.imread(str(face.frame_path)) + if image is None: + return + + x1, y1, x2, y2 = face.bbox + face_img = image[y1:y2, x1:x2] + + if face_img.size == 0: + return + + filename = f"face_{face.frame_index:04d}_{face.id % 100:02d}.jpg" + output_path = output_dir / filename + cv2.imwrite(str(output_path), face_img)