#!/usr/bin/env python3
"""Genera un matte de persona (alpha) desde un video usando MediaPipe Selfie Segmentation.
Salida: matte.mp4 en escala de grises (blanco=persona, negro=fondo), mismo fps/res/duración.
Uso: matte_person.py <in.mp4> <matte.mp4> [start_s] [end_s]"""
import sys, subprocess, numpy as np, cv2, mediapipe as mp

IN = sys.argv[1]; OUT = sys.argv[2]
START = float(sys.argv[3]) if len(sys.argv) > 3 else 0.0
END = float(sys.argv[4]) if len(sys.argv) > 4 else 1e9

cap = cv2.VideoCapture(IN)
fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
W = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)); H = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
n_total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
seg = mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=1)

# ffmpeg pipe writer (gray -> yuv420p mp4)
ff = subprocess.Popen(
    ["ffmpeg","-y","-f","rawvideo","-pix_fmt","gray","-s",f"{W}x{H}","-r",f"{fps}",
     "-i","pipe:0","-an","-c:v","libx264","-preset","veryfast","-crf","16","-pix_fmt","yuv420p",OUT],
    stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

i = 0; done = 0
while True:
    ok, frame = cap.read()
    if not ok: break
    t = i / fps
    if t < START or t > END:
        # fuera de ventana: matte negro (sin persona) para ahorrar cómputo
        ff.stdin.write(np.zeros((H, W), np.uint8).tobytes()); i += 1; continue
    rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    res = seg.process(rgb)
    m = res.segmentation_mask  # float 0..1
    a = np.clip(m, 0, 1)
    # umbral suave + feather para bordes limpios
    a = np.where(a > 0.5, 1.0, a * 0.3)
    a = cv2.GaussianBlur((a * 255).astype(np.uint8), (5, 5), 0)
    ff.stdin.write(a.tobytes())
    i += 1; done += 1
    if done % 150 == 0: print(f"  {done} frames @ t={t:.1f}s")
cap.release(); ff.stdin.close(); ff.wait()
print(f"OK -> {OUT}  ({i} frames, {done} segmentados)")