face-detect-crop (#7)

- update (fafb3f65f4a4bc8c3b20993ea51eb5e4b0b2d868)

face_processing.py

@@ -0,0 +1,518 @@
+"""Face detection and region extraction for lipsync optimization"""
+
+import os
+import math
+import logging
+from typing import List, Dict, Tuple, Optional
+
+import cv2
+import mediapipe as mp
+import numpy as np
+import subprocess
+
+from video_processing import get_video_info
+
+logger = logging.getLogger(__name__)
+
+
+class FaceDetectionError(Exception):
+    """Custom exception for face detection errors"""
+
+    pass
+
+
+def sample_frames_from_video(
+    video_path: str, output_dir: str, sample_count: int = 20
+) -> List[Tuple[int, str]]:
+    """Extract uniform sample frames from video
+
+    Args:
+        video_path: Path to video
+        output_dir: Directory to save extracted frames
+        sample_count: Number of frames to sample
+
+    Returns:
+        List of (frame_index, frame_path) tuples
+    """
+    video_info = get_video_info(video_path)
+    fps = video_info["fps"]
+    duration = video_info["duration"]
+    total_frames = int(duration * fps)
+
+    frames_dir = os.path.join(output_dir, "sampled_frames")
+    os.makedirs(frames_dir, exist_ok=True)
+
+    if total_frames <= sample_count:
+        frame_indices = list(range(total_frames))
+    else:
+        frame_indices = [
+            int(i * total_frames / sample_count) for i in range(sample_count)
+        ]
+
+    extracted_frames = []
+    for idx, frame_idx in enumerate(frame_indices):
+        frame_path = os.path.join(frames_dir, f"frame_{idx:04d}.jpg")
+
+        cmd = [
+            "ffmpeg",
+            "-loglevel",
+            "error",
+            "-y",
+            "-hwaccel",
+            "cuda",
+            "-i",
+            video_path,
+            "-vframes",
+            "1",
+            "-ss",
+            f"{frame_idx / fps:.3f}",
+            frame_path,
+        ]
+        subprocess.run(cmd, check=True)
+
+        extracted_frames.append((frame_idx, frame_path))
+
+    logger.info(f"Extracted {len(extracted_frames)} frames from {video_path}")
+    return extracted_frames
+
+
+def detect_faces_in_frames(
+    extracted_frames: List[Tuple[int, str]],
+    min_confidence: float = 0.5,
+    min_face_pixels: int = 100,
+) -> List[Dict]:
+    """Detect faces in all sampled frames using MediaPipe
+
+    Args:
+        extracted_frames: List of (frame_index, frame_path) tuples
+        min_confidence: Minimum detection confidence (0-1)
+        min_face_pixels: Minimum face size in pixels
+
+    Returns:
+        List of detections: [{"frame_idx", "confidence", "bbox": (x, y, w, h)}]
+    """
+    mp_face_detection = mp.solutions.face_detection
+    detections = []
+
+    with mp_face_detection.FaceDetection(
+        model_selection=0, min_detection_confidence=min_confidence
+    ) as face_detector:
+        for frame_idx, frame_path in extracted_frames:
+            frame = cv2.imread(frame_path)
+            if frame is None:
+                logger.warning(f"Failed to read frame: {frame_path}")
+                continue
+
+            h, w = frame.shape[:2]
+            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+            results = face_detector.process(rgb_frame)
+
+            if results.detections:
+                for detection in results.detections:
+                    bbox = detection.location_data.relative_bounding_box
+
+                    abs_x = int(bbox.xmin * w)
+                    abs_y = int(bbox.ymin * h)
+                    abs_w = int(bbox.width * w)
+                    abs_h = int(bbox.height * h)
+
+                    if abs_w >= min_face_pixels and abs_h >= min_face_pixels:
+                        detections.append(
+                            {
+                                "frame_idx": frame_idx,
+                                "confidence": float(detection.score[0]),
+                                "bbox": (abs_x, abs_y, abs_w, abs_h),
+                            }
+                        )
+
+    logger.info(f"Detected {len(detections)} faces in {len(extracted_frames)} faces")
+    return detections
+
+
+def cluster_face_detections(
+    detections: List[Dict], max_distance: int = 100
+) -> List[List[Dict]]:
+    """Group face detections belonging to the same person using clustering
+
+    Args:
+        detections: List of face detections
+        max_distance: Maximum distance (pixels) to consider detections as same person
+
+    Returns:
+        List of clusters (each cluster is a list of detections)
+    """
+    if not detections:
+        return []
+
+    clusters = []
+    visited = set()
+
+    for i, det_i in enumerate(detections):
+        if i in visited:
+            continue
+
+        x_i, y_i, w_i, h_i = det_i["bbox"]
+        center_i = (x_i + w_i / 2, y_i + h_i / 2)
+
+        cluster = [det_i]
+        visited.add(i)
+
+        for j, det_j in enumerate(detections):
+            if j in visited:
+                continue
+
+            x_j, y_j, w_j, h_j = det_j["bbox"]
+            center_j = (x_j + w_j / 2, y_j + h_j / 2)
+
+            distance = math.sqrt(
+                (center_i[0] - center_j[0]) ** 2 + (center_i[1] - center_j[1]) ** 2
+            )
+
+            if distance < max_distance:
+                cluster.append(det_j)
+                visited.add(j)
+
+        clusters.append(cluster)
+
+    logger.info(f"Clustered {len(detections)} detections into {len(clusters)} clusters")
+    return clusters
+
+
+def select_best_cluster(clusters: List[List[Dict]]) -> List[Dict] | None:
+    """Select the best face cluster (highest frequency)
+
+    Args:
+        clusters: List of clusters
+
+    Returns:
+        Best cluster (most frequent) or None
+    """
+    if not clusters:
+        return None
+
+    scored_clusters = [(len(cluster), cluster) for cluster in clusters]
+    scored_clusters.sort(key=lambda x: x[0], reverse=True)
+
+    best_cluster = scored_clusters[0][1]
+    logger.info(f"Selected best cluster with {len(best_cluster)} detections")
+    return best_cluster
+
+
+def verify_face_stability(
+    cluster: List[Dict], max_movement_percent: float = 0.3
+) -> bool:
+    """Verify face doesn't move too much between frames
+
+    Args:
+        cluster: Face detections for the same person
+        max_movement_percent: Max movement as percentage of average face size
+
+    Returns:
+        True if face is stable, False otherwise
+    """
+    if len(cluster) < 2:
+        return True
+
+    centers = []
+    sizes = []
+
+    for det in cluster:
+        x, y, w, h = det["bbox"]
+        centers.append((x + w / 2, y + h / 2))
+        sizes.append(w * h)
+
+    avg_size = sum(sizes) / len(sizes)
+    avg_face_dim = math.sqrt(avg_size)
+    max_allowed_movement = avg_face_dim * max_movement_percent
+
+    for i in range(len(centers) - 1):
+        dx = abs(centers[i + 1][0] - centers[i][0])
+        dy = abs(centers[i + 1][1] - centers[i][1])
+        movement = math.sqrt(dx**2 + dy**2)
+
+        if movement > max_allowed_movement:
+            logger.warning(
+                f"Face movement {movement:.1f}px > {max_allowed_movement:.1f}px"
+            )
+            return False
+
+    return True
+
+
+def calculate_face_bbox_from_cluster(cluster: List[Dict]) -> Dict:
+    """Calculate average face bounding box from cluster
+
+    Args:
+        cluster: Face detections for the same person
+
+    Returns:
+        Dict: {"x", "y", "width", "height"}
+    """
+    weighted_x = 0
+    weighted_y = 0
+    weighted_w = 0
+    weighted_h = 0
+    total_weight = 0
+
+    for det in cluster:
+        x, y, w, h = det["bbox"]
+        weight = det["confidence"]
+        weighted_x += x * weight
+        weighted_y += y * weight
+        weighted_w += w * weight
+        weighted_h += h * weight
+        total_weight += weight
+
+    avg_bbox = {
+        "x": int(weighted_x / total_weight),
+        "y": int(weighted_y / total_weight),
+        "width": int(weighted_w / total_weight),
+        "height": int(weighted_h / total_weight),
+    }
+
+    return avg_bbox
+
+
+def calculate_safe_crop_512x512(
+    face_bbox: Dict, video_width: int, video_height: int
+) -> Dict:
+    """Calculate safe 512x512 crop region ensuring face is inside
+
+    Args:
+        face_bbox: Face bounding box {"x", "y", "width", "height"}
+        video_width: Video width
+        video_height: Video height
+
+    Returns:
+        Dict: {"x", "y", "width": 512, "height": 512}
+    """
+    crop_size = 512
+    crop_half = crop_size // 2
+
+    face_center_x = face_bbox["x"] + face_bbox["width"] / 2
+    face_center_y = face_bbox["y"] + face_bbox["height"] / 2
+
+    crop_x = int(face_center_x - crop_half)
+    crop_y = int(face_center_y - crop_half)
+
+    crop_x = max(0, crop_x)
+    crop_y = max(0, crop_y)
+    crop_x = min(video_width - crop_size, crop_x)
+    crop_y = min(video_height - crop_size, crop_y)
+
+    face_right = face_bbox["x"] + face_bbox["width"]
+    face_bottom = face_bbox["y"] + face_bbox["height"]
+    crop_right = crop_x + crop_size
+    crop_bottom = crop_y + crop_size
+
+    if (
+        face_bbox["x"] < crop_x
+        or face_bbox["y"] < crop_y
+        or face_right > crop_right
+        or face_bottom > crop_bottom
+    ):
+        if face_bbox["x"] < crop_x:
+            crop_x = face_bbox["x"]
+        elif face_right > crop_right:
+            crop_x = face_right - crop_size
+
+        if face_bbox["y"] < crop_y:
+            crop_y = face_bbox["y"]
+        elif face_bottom > crop_bottom:
+            crop_y = face_bottom - crop_size
+
+        crop_x = max(0, crop_x)
+        crop_y = max(0, crop_y)
+        crop_x = min(video_width - crop_size, crop_x)
+        crop_y = min(video_height - crop_size, crop_y)
+
+    return {"x": crop_x, "y": crop_y, "width": crop_size, "height": crop_size}
+
+
+def detect_face_region(
+    video_path: str,
+    output_dir: str,
+    sample_count: int = 20,
+    min_confidence: float = 0.5,
+    min_face_pixels: int = 100,
+    max_face_movement_percent: float = 0.3,
+) -> Dict:
+    """Main function: Detect face and calculate safe 512x512 crop
+
+    Args:
+        video_path: Path to video
+        output_dir: Directory for temporary files
+        sample_count: Number of frames to sample (default: 20)
+        min_confidence: Minimum detection confidence
+        min_face_pixels: Minimum face size in pixels
+        max_face_movement_percent: Max allowed face movement
+
+    Returns:
+        Dict: {"x", "y", "width": 512, "height": 512, "face_bbox"}
+
+    Raises:
+        FaceDetectionError: If face detection fails
+    """
+    try:
+        video_info = get_video_info(video_path)
+        video_w, video_h = video_info["width"], video_info["height"]
+
+        if video_w < 512 or video_h < 512:
+            raise FaceDetectionError(
+                f"Video resolution {video_w}x{video_h} < 512x512. "
+                f"Please upload higher resolution video."
+            )
+
+        extracted_frames = sample_frames_from_video(
+            video_path, output_dir, sample_count
+        )
+        detections = detect_faces_in_frames(
+            extracted_frames, min_confidence, min_face_pixels
+        )
+
+        if not detections:
+            raise FaceDetectionError(
+                f"No face detected in {sample_count} sampled frames. "
+                f"Please upload a video with a visible face."
+            )
+
+        frames_with_face = len(set(d["frame_idx"] for d in detections))
+        face_coverage = frames_with_face / len(extracted_frames)
+
+        if face_coverage < 0.5:
+            raise FaceDetectionError(
+                f"Face detected in only {frames_with_face}/{len(extracted_frames)} frames "
+                f"({face_coverage * 100:.1f}%). "
+                f"Please upload a video with a visible face."
+            )
+
+        clusters = cluster_face_detections(detections)
+        best_cluster = select_best_cluster(clusters)
+
+        if best_cluster is None:
+            raise FaceDetectionError(
+                f"Failed to identify main face in video. "
+                f"Please upload a video with a clear, visible face."
+            )
+
+        if not verify_face_stability(best_cluster, max_face_movement_percent):
+            raise FaceDetectionError(
+                f"Face moves too much between frames. "
+                f"Please upload a video with a stable face position."
+            )
+
+        face_bbox = calculate_face_bbox_from_cluster(best_cluster)
+        crop_bbox = calculate_safe_crop_512x512(face_bbox, video_w, video_h)
+
+        crop_bbox["face_bbox"] = face_bbox
+
+        logger.info(
+            f"Face detected at ({face_bbox['x']}, {face_bbox['y']}) "
+            f"size {face_bbox['width']}x{face_bbox['height']}, "
+            f"crop at ({crop_bbox['x']}, {crop_bbox['y']})"
+        )
+
+        return crop_bbox
+
+    except FaceDetectionError:
+        raise
+    except Exception as e:
+        logger.error(f"Face detection failed: {e}")
+        raise FaceDetectionError(f"Face detection failed: {str(e)}")
+
+
+def crop_video_to_512x512(video_path: str, crop_bbox: Dict, output_dir: str) -> str:
+    """Crop video to 512x512 using calculated bbox
+
+    Args:
+        video_path: Path to input video
+        crop_bbox: Crop region {"x", "y", "width", "height"}
+        output_dir: Directory to save output
+
+    Returns:
+        Path to cropped video
+    """
+    output_path = os.path.join(output_dir, "face_cropped_512x512.mp4")
+
+    cmd = [
+        "ffmpeg",
+        "-loglevel",
+        "error",
+        "-y",
+        "-hwaccel",
+        "cuda",
+        "-i",
+        video_path,
+        "-vf",
+        f"crop={crop_bbox['width']}:{crop_bbox['height']}:{crop_bbox['x']}:{crop_bbox['y']}",
+        "-c:v",
+        "h264_nvenc",
+        "-preset",
+        "p6",
+        "-cq",
+        "18",
+        "-c:a",
+        "copy",
+        output_path,
+    ]
+
+    subprocess.run(cmd, check=True)
+    logger.info(f"Cropped video to 512x512: {output_path}")
+    return output_path
+
+
+def blend_face_into_original(
+    original_video: str,
+    face_video: str,
+    crop_bbox: Dict,
+    output_dir: str,
+    feather: int = 10,
+) -> str:
+    """Blend face video back into original video
+
+    Args:
+        original_video: Path to original video
+        face_video: Path to face-cropped lipsynced video
+        crop_bbox: Crop region {"x", "y", "width", "height"}
+        output_dir: Directory to save output
+        feather: Feather radius for smooth blending
+
+    Returns:
+        Path to blended video
+    """
+    output_path = os.path.join(output_dir, "face_blended.mp4")
+
+    cmd = [
+        "ffmpeg",
+        "-loglevel",
+        "error",
+        "-y",
+        "-hwaccel",
+        "cuda",
+        "-hwaccel_output_format",
+        "cuda",
+        "-i",
+        original_video,
+        "-i",
+        face_video,
+        "-filter_complex",
+        (
+            f"[1:v]format=rgba,split[A][B];"
+            f"[A]colorchannelmixer=aa=0.8[B];"
+            f"[0:v][1:v]overlay={crop_bbox['x']}:{crop_bbox['y']}"
+        ),
+        "-c:v",
+        "h264_nvenc",
+        "-preset",
+        "p6",
+        "-cq",
+        "18",
+        "-c:a",
+        "copy",
+        output_path,
+    ]
+
+    subprocess.run(cmd, check=True)
+    logger.info(f"Blended face into original: {output_path}")
+    return output_path

processing.py

@@ -12,6 +12,12 @@ import torch
 
 from audio_processing import get_audio_duration, prepare_target_audio
 from config import PROCESSED_RESULTS_DIR
+from face_processing import (
+    blend_face_into_original,
+    crop_video_to_512x512,
+    detect_face_region,
+    FaceDetectionError,
+)
 from lipsync_processing import apply_lipsync_to_video
 from time_util import timer
 from utils import setup_output_dir
@@ -97,7 +103,27 @@ def process_lipsync_with_audio_target(
         gc.collect()
         logger.info(f"Memory after video loop: {get_memory_usage()}")
 
-        progress(0.3, desc="🎵 Đang xử lý audio...")
+        progress(0.2, desc="👤 Đang phát hiện khuôn mặt...")
+        with timer("Detecting face"):
+            try:
+                face_bbox = detect_face_region(video_looped, output_dir)
+            except FaceDetectionError as e:
+                logger.error(f"Face detection failed: {e}")
+                raise gr.Error(str(e))
+
+        gc.collect()
+        logger.info(f"Memory after face detection: {get_memory_usage()}")
+
+        progress(0.3, desc="✂️ Đang crop video 512x512...")
+        with timer("Cropping video to 512x512"):
+            face_cropped_512 = crop_video_to_512x512(
+                video_looped, face_bbox, output_dir
+            )
+
+        gc.collect()
+        logger.info(f"Memory after crop: {get_memory_usage()}")
+
+        progress(0.4, desc="🎵 Đang xử lý audio...")
         logger.info(f"Memory before audio prep: {get_memory_usage()}")
 
         with timer("Preparing target audio"):
@@ -106,17 +132,17 @@ def process_lipsync_with_audio_target(
         gc.collect()
         logger.info(f"Memory after audio prep: {get_memory_usage()}")
 
-        progress(0.5, desc="👄 Đang lipsync...")
+        progress(0.6, desc="👄 Đang lipsync...")
 
         logger.info(
-            f"Starting lipsync: video={video_looped}, audio_16k={audio_16k}, output={output_dir}"
+            f"Starting lipsync: video={face_cropped_512}, audio_16k={audio_16k}, output={output_dir}"
         )
         logger.info(f"Memory before lipsync: {get_memory_usage()}")
 
         with timer("Applying lipsync"):
             try:
-                lipsynced_video = apply_lipsync_to_video(
-                    video_looped, audio_16k, output_dir
+                lipsynced_face = apply_lipsync_to_video(
+                    face_cropped_512, audio_16k, output_dir
                 )
             except Exception as e:
                 logger.error(f"Lipsync failed with error: {type(e).__name__}: {e}")
@@ -127,9 +153,17 @@ def process_lipsync_with_audio_target(
         gc.collect()
         logger.info(f"Memory after lipsync: {get_memory_usage()}")
 
-        progress(0.8, desc="🔗 Đang ghép audio...")
+        progress(0.8, desc="🔀 Đang ghép video...")
+        with timer("Blending face into original"):
+            lipsynced_full = blend_face_into_original(
+                video_looped, lipsynced_face, face_bbox, output_dir
+            )
+
+        gc.collect()
+        logger.info(f"Memory after blend: {get_memory_usage()}")
 
-        final_video = merge_audio_video(lipsynced_video, audio_upsampled, output_dir)
+        progress(0.9, desc="🔗 Đang ghép audio...")
+        final_video = merge_audio_video(lipsynced_full, audio_upsampled, output_dir)
 
         progress(1.0, desc="✅ Hoàn tất!")
         logger.info(f"Memory at end: {get_memory_usage()}")

requirements.txt

@@ -21,3 +21,7 @@ git+https://github.com/techio-dev/LatentSync-batch.git
 
 # --- Flash Attention (HuggingFace Spaces) ---
 https://huggingface.co/alexnasa/flash-attn-3/resolve/main/128/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl
+
+# --- Face Detection ---
+mediapipe>=0.10.0
+opencv-python>=4.8.0

video_processing.py

@@ -23,6 +23,8 @@ def crop_video_duration(video_path: str, duration: int, output_dir: str) -> str:
         "-loglevel",
         "error",
         "-y",
+        "-hwaccel",
+        "cuda",
         "-i",
         video_path,
         "-t",
@@ -52,6 +54,8 @@ def merge_audio_video(video_path: str, audio_path: str, output_dir: str) -> str:
         "-loglevel",
         "error",
         "-y",
+        "-hwaccel",
+        "cuda",
         "-i",
         video_path,
         "-i",
@@ -127,6 +131,8 @@ def loop_video_to_match_audio(
             "-loglevel",
             "error",
             "-y",
+            "-hwaccel",
+            "cuda",
             "-i",
             video_path,
             "-t",
@@ -142,6 +148,8 @@ def loop_video_to_match_audio(
             "-loglevel",
             "error",
             "-y",
+            "-hwaccel",
+            "cuda",
             "-stream_loop",
             f"{loop_count}",
             "-i",