didi-lot1-ai/ai_platform/modules/forensic_features/preprocessing.py

"""
preprocessing.py — Shared infrastructure for all forensic tools.

Handles:
  1. Frame extraction from video via ffmpeg
  2. Face detection on extracted frames
  3. Loading frames as numpy arrays for analysis

Every tool imports from this. No tool extracts frames on its own.
"""

import subprocess
import os
import cv2
import numpy as np


def extract_frames(video_path, output_dir, every_n=30):
    """
    Extract every Nth frame from a video as PNG files.

    Uses ffmpeg under the hood:
      ffmpeg -i <video> -vf "select=not(mod(n,N))" -vsync vfr <output>

    Why PNG? Lossless — we don't want to add compression artifacts
    on top of whatever the video already has. Our forensic methods
    need to analyze the pixels as they were encoded, not re-compressed.

    Args:
        video_path: path to the .mp4 file
        output_dir: where to save the frame PNGs
        every_n: extract 1 frame per N (default 30 = ~1fps for 30fps video)

    Returns:
        list of file paths to the extracted frames, sorted by frame number
    """
    os.makedirs(output_dir, exist_ok=True)

    cmd = [
        "ffmpeg", "-i", video_path,
        "-vf", f"select=not(mod(n\\,{every_n}))",
        "-vsync", "vfr",
        os.path.join(output_dir, "frame_%04d.png"),
        "-y"
    ]
    subprocess.run(cmd, capture_output=True, text=True)

    frames = sorted([
        os.path.join(output_dir, f)
        for f in os.listdir(output_dir)
        if f.startswith("frame_") and f.endswith(".png")
    ])
    return frames


def extract_consecutive_frames(video_path, output_dir, start_frame, count=30):
    """
    Extract a block of consecutive frames starting at a specific frame number.

    Used by temporal methods (optical flow, temporal variance, SSIM)
    that need frame-to-frame comparison without gaps.

    Uses ffmpeg:
      ffmpeg -i <video> -vf "select=between(n,start,start+count)" -vsync vfr <output>

    Args:
        video_path: path to the .mp4 file
        output_dir: where to save
        start_frame: which frame number to start from
        count: how many consecutive frames to extract

    Returns:
        list of file paths to the extracted frames
    """
    os.makedirs(output_dir, exist_ok=True)

    end_frame = start_frame + count - 1
    cmd = [
        "ffmpeg", "-i", video_path,
        "-vf", f"select=between(n\\,{start_frame}\\,{end_frame})",
        "-vsync", "vfr",
        os.path.join(output_dir, "consec_%04d.png"),
        "-y"
    ]
    subprocess.run(cmd, capture_output=True, text=True)

    frames = sorted([
        os.path.join(output_dir, f)
        for f in os.listdir(output_dir)
        if f.startswith("consec_") and f.endswith(".png")
    ])
    return frames


def load_frame(frame_path):
    """
    Load a single frame as a BGR numpy array (OpenCV default).

    Args:
        frame_path: path to a PNG file

    Returns:
        numpy array of shape (H, W, 3) in BGR color order
    """
    return cv2.imread(frame_path)


def load_frame_gray(frame_path):
    """
    Load a single frame as grayscale float64.
    Most forensic methods work on grayscale.

    Returns:
        numpy array of shape (H, W) as float64
    """
    return cv2.imread(frame_path, cv2.IMREAD_GRAYSCALE).astype(np.float64)


def detect_faces(frame, min_size=50):
    """
    Detect faces in a frame using Haar cascade.

    This is the same detector I used in all our analyses.
    It's fast and good enough for ROI extraction.
    For landmark-based methods (m33-m40), we'll use MediaPipe instead.

    Args:
        frame: BGR numpy array
        min_size: minimum face size in pixels

    Returns:
        list of (x, y, w, h) tuples, sorted largest first
    """
    face_cascade = cv2.CascadeClassifier(
        cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
    )
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(
        gray, scaleFactor=1.1, minNeighbors=4, minSize=(min_size, min_size)
    )
    if len(faces) == 0:
        return []
    # Sort by area, largest first
    faces = sorted(faces, key=lambda f: f[2] * f[3], reverse=True)
    return [tuple(f) for f in faces]


def get_face_roi(frame, face, margin=0.3):
    """
    Extract a face ROI with margin around it.

    The margin is important — blending boundaries in face swaps
    sit OUTSIDE the face detection box, so we need to include
    some surrounding area.

    Args:
        frame: BGR or grayscale numpy array
        face: (x, y, w, h) tuple from detect_faces
        margin: how much to expand (0.3 = 30% on each side)

    Returns:
        cropped numpy array of the face region
    """
    x, y, w, h = face
    m = int(w * margin)
    H_img = frame.shape[0]
    W_img = frame.shape[1]
    y1 = max(0, y - m)
    y2 = min(H_img, y + h + m)
    x1 = max(0, x - m)
    x2 = min(W_img, x + w + m)
    return frame[y1:y2, x1:x2]