preprocessor/videomask_generator.py

import torch
import cv2
import numpy as np
from torchvision.transforms.functional import normalize
from tqdm import tqdm
from PIL import Image, ImageOps
import random
import os
import requests
from insightface.app import FaceAnalysis
from facexlib.parsing import init_parsing_model
from typing import Union, Optional, Tuple, List

# --- Helper Functions (Unchanged) ---
def tensor_to_cv2_img(tensor_frame: torch.Tensor) -> np.ndarray:
    """Converts a single RGB torch tensor to a BGR OpenCV image."""
    img_np = (tensor_frame.cpu().numpy() * 255).astype(np.uint8)
    return cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)

def tensor_to_cv2_bgra_img(tensor_frame: torch.Tensor) -> np.ndarray:
    """Converts a single RGBA torch tensor to a BGRA OpenCV image."""
    if tensor_frame.shape[2] != 4:
        raise ValueError("Input tensor must be an RGBA image with 4 channels.")
    img_np = (tensor_frame.cpu().numpy() * 255).astype(np.uint8)
    return cv2.cvtColor(img_np, cv2.COLOR_RGBA2BGRA)

def pil_to_tensor(image: Image.Image) -> torch.Tensor:
    """Converts a PIL image to a torch tensor."""
    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0)

class VideoMaskGenerator:
    def __init__(self, antelopv2_path=".", device: Optional[torch.device] = None):
        if device is None:
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        else:
            self.device = device

        print(f"Using device: {self.device}")

        providers = ["CUDAExecutionProvider"] if self.device.type == "cuda" else ["CPUExecutionProvider"]
        
        # Initialize face detection and landmark model (antelopev2 provides both)
        self.detection_model = FaceAnalysis(name="antelopev2", root=antelopv2_path, providers=providers)
        self.detection_model.prepare(ctx_id=0, det_size=(640, 640))

        # Initialize face parsing model
        self.parsing_model = init_parsing_model(model_name="bisenet", device=self.device)
        self.parsing_model.eval()
        
        print("FaceProcessor initialized successfully.")

    def process(
        self,
        video_path: str,
        face_image: Union[str, Image.Image],
        confidence_threshold: float = 0.5,
        face_crop_scale: float = 1.5,
        dilation_kernel_size: int = 10,
        feather_amount: int = 21,
        random_horizontal_flip_chance: float = 0.0,
        match_angle_and_size: bool = True
    ) -> Tuple[np.ndarray, np.ndarray, int, int, int]:
        """
        Processes a video to replace a face with a provided face image.

        Args:
            video_path (str): Path to the input video file.
            face_image (Union[str, Image.Image]): Path or PIL image of the face to paste.
            confidence_threshold (float): Confidence threshold for face detection.
            face_crop_scale (float): Scale factor for cropping the detected face box.
            dilation_kernel_size (int): Kernel size for mask dilation.
            feather_amount (int): Amount of feathering for the mask edges.
            random_horizontal_flip_chance (float): Chance to flip the source face horizontally.
            match_angle_and_size (bool): Whether to use landmark matching for rotation and scale.

        Returns:
            Tuple[np.ndarray, np.ndarray, int, int, int]: 
                - Processed video as a numpy array (F, H, W, C).
                - Generated masks as a numpy array (F, H, W).
                - Width of the processed video.
                - Height of the processed video.
                - Number of frames in the processed video.
        """
        # --- Video Pre-processing ---
        if not os.path.exists(video_path):
            raise FileNotFoundError(f"Video file not found at: {video_path}")

        cap = cv2.VideoCapture(video_path)
        frames = []
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
            frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        cap.release()
        
        if not frames:
            raise ValueError("Could not read any frames from the video.")

        video_np = np.array(frames)
        
        h, w = video_np.shape[1], video_np.shape[2]
        new_h, new_w = (h // 16) * 16, (w // 16) * 16
        
        y_start = (h - new_h) // 2
        x_start = (w - new_w) // 2
        video_cropped = video_np[:, y_start:y_start+new_h, x_start:x_start+new_w, :]

        num_frames = video_cropped.shape[0]
        target_frames = (num_frames // 4) * 4 + 1
        video_trimmed = video_cropped[:target_frames]

        final_h, final_w, final_frames = video_trimmed.shape[1], video_trimmed.shape[2], video_trimmed.shape[0]
        print(f"Video pre-processed: {final_w}x{final_h}, {final_frames} frames.")

        # --- Face Image Pre-processing & Source Landmark Extraction ---
        if isinstance(face_image, str):
            if face_image.startswith("http"):
                face_image = Image.open(requests.get(face_image, stream=True, timeout=10).raw)
            else:
                face_image = Image.open(face_image)
        
        face_image = ImageOps.exif_transpose(face_image).convert("RGBA")
        face_rgba_tensor = pil_to_tensor(face_image)
        face_to_paste_cv2 = tensor_to_cv2_bgra_img(face_rgba_tensor)

        source_kpts = None
        if match_angle_and_size:
            # Use insightface (antelopev2) to get landmarks from the source face image
            source_face_bgr = cv2.cvtColor(face_to_paste_cv2, cv2.COLOR_BGRA2BGR)
            source_faces = self.detection_model.get(source_face_bgr)
            if source_faces:
                # Use the landmarks from the first (and likely only) detected face
                source_kpts = source_faces[0].kps
            else:
                print("[Warning] No face or landmarks found in source image. Disabling angle matching.")
                match_angle_and_size = False
        
        face_to_paste_pil = Image.fromarray((face_rgba_tensor.cpu().numpy() * 255).astype(np.uint8), 'RGBA')

        # --- Main Processing Loop ---
        processed_frames_list = []
        mask_list = []

        for i in tqdm(range(final_frames), desc="Pasting face onto frames"):
            frame_rgb = video_trimmed[i]
            frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
            
            # Use insightface for detection and landmarks
            faces = self.detection_model.get(frame_bgr)
            
            pasted = False
            final_mask = np.zeros((final_h, final_w), dtype=np.uint8)

            if faces:
                largest_face = max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
                
                if largest_face.det_score > confidence_threshold:
                    # **MODIFIED BLOCK**: Use insightface landmarks for affine transform
                    if match_angle_and_size and source_kpts is not None:
                        target_kpts = largest_face.kps # Get landmarks directly from the detected face
                        
                        # Estimate the transformation matrix
                        M, _ = cv2.estimateAffinePartial2D(source_kpts, target_kpts, method=cv2.LMEDS)
                        
                        if M is not None:
                            # Split the RGBA source face for separate warping
                            b, g, r, a = cv2.split(face_to_paste_cv2)
                            source_rgb_cv2 = cv2.merge([r, g, b])
                            
                            # Warp the face and its alpha channel
                            warped_face = cv2.warpAffine(source_rgb_cv2, M, (final_w, final_h))
                            warped_alpha = cv2.warpAffine(a, M, (final_w, final_h))
                            
                            # Blend the warped face onto the frame using the warped alpha channel
                            alpha_float = warped_alpha.astype(np.float32) / 255.0
                            alpha_expanded = np.expand_dims(alpha_float, axis=2)
                            
                            frame_rgb = (1.0 - alpha_expanded) * frame_rgb + alpha_expanded * warped_face
                            frame_rgb = frame_rgb.astype(np.uint8)
                            final_mask = warped_alpha
                            pasted = True

                    # Fallback to simple box-pasting if angle matching is off or fails
                    if not pasted: 
                        x1, y1, x2, y2 = map(int, largest_face.bbox)
                        center_x, center_y = (x1 + x2) // 2, (y1 + y2) // 2
                        side_len = int(max(x2 - x1, y2 - y1) * face_crop_scale)
                        half_side = side_len // 2
                        
                        crop_y1, crop_x1 = max(center_y - half_side, 0), max(center_x - half_side, 0)
                        crop_y2, crop_x2 = min(center_y + half_side, final_h), min(center_x + half_side, final_w)
                        
                        box_w, box_h = crop_x2 - crop_x1, crop_y2 - crop_y1

                        if box_w > 0 and box_h > 0:
                            source_img = face_to_paste_pil.copy()
                            if random.random() < random_horizontal_flip_chance:
                                source_img = source_img.transpose(Image.FLIP_LEFT_RIGHT)
                            
                            face_resized = source_img.resize((box_w, box_h), Image.Resampling.LANCZOS)
                            
                            target_frame_pil = Image.fromarray(frame_rgb)
                            
                            # --- Mask Generation using BiSeNet ---
                            face_crop_bgr = cv2.cvtColor(frame_rgb[crop_y1:crop_y2, crop_x1:crop_x2], cv2.COLOR_RGB2BGR)
                            if face_crop_bgr.size > 0:
                                face_resized_512 = cv2.resize(face_crop_bgr, (512, 512), interpolation=cv2.INTER_AREA)
                                face_rgb_512 = cv2.cvtColor(face_resized_512, cv2.COLOR_BGR2RGB)
                                face_tensor_in = torch.from_numpy(face_rgb_512.astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0).to(self.device)
                                
                                with torch.no_grad():
                                    normalized_face = normalize(face_tensor_in, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                                    parsing_map = self.parsing_model(normalized_face)[0].argmax(dim=1, keepdim=True)
                                
                                parsing_map_np = parsing_map.squeeze().cpu().numpy().astype(np.uint8)
                                parts_to_include = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13] # All face parts
                                final_mask_512 = np.isin(parsing_map_np, parts_to_include).astype(np.uint8) * 255
                                
                                if dilation_kernel_size > 0:
                                    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilation_kernel_size, dilation_kernel_size))
                                    final_mask_512 = cv2.dilate(final_mask_512, kernel, iterations=1)
                                
                                if feather_amount > 0:
                                    if feather_amount % 2 == 0: feather_amount += 1
                                    final_mask_512 = cv2.GaussianBlur(final_mask_512, (feather_amount, feather_amount), 0)
                                
                                mask_resized_to_crop = cv2.resize(final_mask_512, (box_w, box_h), interpolation=cv2.INTER_LINEAR)
                                generated_mask_pil = Image.fromarray(mask_resized_to_crop, mode='L')
                                
                                target_frame_pil.paste(face_resized, (crop_x1, crop_y1), mask=generated_mask_pil)
                                frame_rgb = np.array(target_frame_pil)
                                final_mask[crop_y1:crop_y2, crop_x1:crop_x2] = mask_resized_to_crop

            processed_frames_list.append(frame_rgb)
            mask_list.append(final_mask)

        output_video = np.stack(processed_frames_list)
        # Ensure mask has a channel dimension for consistency
        output_masks = np.stack(mask_list)[..., np.newaxis] 
        
        return (output_video, output_masks, final_w, final_h, final_frames)