Update processing/two_stage/two_stage_processor.py

processing/two_stage/two_stage_processor.py

@@ -1,279 +1,242 @@
 #!/usr/bin/env python3
 """
-MatAnyone Adapter - Robust interface handler for InferenceCore
-Handles unknown process_video signatures and provides fallback to step-based processing
+TRUE Two-Stage Green-Screen Processing System ✅
+
+Stage 1: Original → SAM2(ref mask/trimap) + MatAnyone → ALPHA PNGs → GREEN master (H.264)
+Stage 2: Key GREEN master → Final composite
+
+Safe for HF Spaces:
+- No ProRes (uses mp4v for GREEN master)
+- No MoviePy apply_to=['mask'] (prevents black output)
+- Alpha stays lossless as PNG sequence
 """
 
-import os
-import cv2
-import inspect
-import logging
+import os, cv2, time, logging, tempfile, traceback
+from pathlib import Path
+from typing import Optional, Tuple, Callable, List
 import numpy as np
-from typing import Optional, Union
+from moviepy.editor import VideoFileClip
 
 logger = logging.getLogger(__name__)
 
-# ---------- Helpers
-
-def _read_mask_01(path: str, wh=None) -> np.ndarray:
-    """Read a mask/trimap and return float32 0..1 alpha guidance."""
-    m = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
-    if m is None:
-        raise ValueError(f"Reference mask/trimap not found: {path}")
-
-    # Trimap heuristic: values clustered at {0, ~128, 255}
-    unique = np.unique(m)
-    if len(unique) <= 4 and 120 in unique or any(110 <= v <= 140 for v in unique):
-        # Treat 128± as 'unknown'
-        a = np.zeros_like(m, dtype=np.float32)
-        a[m >= 200] = 1.0           # FG
-        a[(m > 50) & (m < 200)] = 0.5  # Unknown band → 0.5 as neutral prior
-    else:
-        # Binary / soft mask → normalize
-        a = m.astype(np.float32) / 255.0
-
-    if wh is not None and (m.shape[1], m.shape[0]) != wh:
-        a = cv2.resize(a, wh, interpolation=cv2.INTER_LINEAR)
-    return np.clip(a, 0.0, 1.0).astype(np.float32)
-
-def _to_chw01(img_bgr: np.ndarray) -> np.ndarray:
-    """BGR uint8 HxWx3 -> CHW float32 0..1"""
-    x = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
-    return np.transpose(x, (2, 0, 1)).copy()
-
-def _ensure_dir(d: str) -> str:
-    os.makedirs(d, exist_ok=True)
-    return d
-
-# ---------- Primary entry
-
-def run_matanyone_with_reference_mask(
-    inference_core,               # your InferenceCore instance
-    video_path: str,
-    reference_mask_path: str,     # binary mask or trimap
-    output_dir: str,
-    *,
-    fps_override: Optional[float] = None,
-    write_png: bool = True,
-) -> str:
-    """
-    Run MatAnyone with reference mask using robust interface detection.
-    
-    Returns: path to alpha PNG directory (if write_png) or alpha video produced by process_video
-    """
-    _ensure_dir(output_dir)
-
-    # 1) Try process_video with best-effort argument mapping
-    proc = getattr(inference_core, "process_video", None)
-    if callable(proc):
-        sig = inspect.signature(proc)
-        params = list(sig.parameters.keys())
-        logger.info(f"MatAnyone.process_video signature: {params}")
-
-        # Common arg names we'll look for
-        argmap = {}
-        for p in params:
-            lp = p.lower()
-            if "video" in lp and ("path" in lp or lp == "video"):
-                argmap[p] = video_path
-            elif ("trimap" in lp) or ("mask" in lp and "ref" in lp):
-                argmap[p] = reference_mask_path
-            elif ("mask" == lp) or ("ref_mask" == lp):
-                argmap[p] = reference_mask_path
-            elif ("out" in lp and ("dir" in lp or "path" in lp)):
-                argmap[p] = output_dir
-
-        # If we at least matched video + (mask/trimap) + out, try it
-        has_video = any(k for k,v in argmap.items() if v==video_path)
-        has_mask  = any(k for k,v in argmap.items() if v==reference_mask_path)
-        has_out   = any(k for k,v in argmap.items() if v==output_dir)
-
-        if has_video and has_mask and has_out:
-            try:
-                logger.info(f"Calling process_video with mapped args: {argmap}")
-                result = proc(**argmap)  # may return a path or None
-                # If it produced an alpha video, we'll convert to PNGs for safety
-                produced = _detect_alpha_artifact(output_dir)
-                if produced and produced.lower().endswith(".mp4"):
-                    alpha_dir = os.path.join(output_dir, "alpha_png")
-                    _mp4_to_alpha_png(produced, alpha_dir)
-                    return alpha_dir
-                elif produced and os.path.isdir(produced):
-                    # Already a folder of frames
-                    return produced
-                else:
-                    # Nothing obvious—fall back to manual pipeline
-                    logger.warning("process_video returned no detectable alpha; falling back to step()")
-            except Exception as e:
-                logger.warning(f"process_video failed ({e}); falling back to step().")
-
-    # 2) Fallback: step-based pipeline (canonical)
-    logger.info("Using step-based pipeline: init with first frame (image+mask), then step(image) ...")
-
-    cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened():
-        raise ValueError(f"Cannot open video: {video_path}")
-    fps = fps_override or cap.get(cv2.CAP_PROP_FPS) or 25
-    w   = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    h   = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-
-    # first frame
-    ok, frame = cap.read()
-    if not ok:
-        cap.release()
-        raise ValueError("Could not read first frame")
+# =========================
+# Constants (RGB colors)
+# =========================
+CHROMA_PRESETS = {
+    "office":   {"name": "Professional Office", "color": (240, 240, 240)},
+    "blue":     {"name": "Studio Blue",         "color": (70, 130, 180)},
+    "green":    {"name": "Chroma Green",        "color": (0, 177, 64)},
+    "white":    {"name": "Clean White",         "color": (255, 255, 255)},
+    "gradient": {"name": "Soft Gradient",       "color": (200, 220, 255)}
+}
+PURE_GREEN_RGB = (0, 255, 0)
+DEFAULT_REFERENCE_FRAMES = 3
+TEMP_DIR_PREFIX = "twostage_greenscreen_"
+
+def get_quality_profile():
+    # keep minimal: only keying knobs if you want to expose later
+    class P:
+        chroma_tolerance = 25.0
+        chroma_softness = 3.0
+        spill_suppression = 0.6
+        edge_feather = 8
+    return P()
+
+class TwoStageProcessor:
+    def __init__(self, sam2_handler, matanyone_processor):
+        logger.info("Initializing True Two-Stage Green Screen Processor…")
+        self.sam2_handler = sam2_handler
+        self.matanyone_processor = matanyone_processor
+        self.temp_dir: Optional[str] = None
+        self.quality_profile = get_quality_profile()
+
+    # ------------------------------------------------------------
+
+    def process_video(
+        self,
+        video_path: str,
+        background_image_path: Optional[str] = None,
+        background_preset: str = "office",
+        callback: Optional[Callable] = None,
+        trim_seconds: float = 5.0
+    ) -> Tuple[Optional[str], Optional[str], str]:
+        start = time.time()
+        self.temp_dir = tempfile.mkdtemp(prefix=TEMP_DIR_PREFIX)
+
+        try:
+            logger.info(f"🎬 TRUE Two-Stage: {video_path}")
+            logger.info(f"📁 Temp: {self.temp_dir}")
 
-    # Prepare inputs
-    alpha0 = _read_mask_01(reference_mask_path, wh=(w, h))
-    img0   = _to_chw01(frame)  # CHW 0..1
+            green_master = self._stage1_create_greenscreen_video(video_path, trim_seconds)
+            final_out    = self._stage2_chroma_key_replacement(
+                green_master, background_image_path, background_preset
+            )
 
-    # Initialize model session
-    if hasattr(inference_core, "reset"):
-        inference_core.reset()
+            status = "✅ Two-stage (GREEN master → Key) completed."
+            logger.info(f"{status} Output: {final_out}")
+            return final_out, green_master, status
 
-    # Many video-matting cores use 'step(image, mask=None)' where mask only on first call
-    try:
-        _maybe_step_with_mask(inference_core, img0, alpha0)
-    except Exception as e:
-        cap.release()
-        raise RuntimeError(f"MatAnyone step(init) failed: {e}")
+        except Exception as e:
+            logger.error(f"Processor failed: {e}")
+            logger.error(traceback.format_exc())
+            return None, None, f"Processing failed: {e}"
+
+    # =========================
+    # STAGE 1: GREEN MASTER
+    # =========================
 
-    # Output dir
-    alpha_dir = _ensure_dir(os.path.join(output_dir, "alpha_png"))
+    def _stage1_create_greenscreen_video(self, video_path: str, trim_seconds: float) -> str:
+        logger.info("STAGE 1: Creating GREEN master via SAM2 + MatAnyone …")
 
-    # Write first alpha
-    _write_alpha_png(alpha0, os.path.join(alpha_dir, f"alpha_{0:06d}.png"))
+        ref_mask = self._create_reference_mask(video_path)
+        trimap   = self._create_trimap(ref_mask)
 
-    i = 1
-    while True:
-        ok, frame = cap.read()
-        if not ok:
-            break
-        img = _to_chw01(frame)
+        # Trim source (codec here is fine; this is not the matte)
+        trimmed = os.path.join(self.temp_dir, "trimmed_src.mp4")
+        with VideoFileClip(video_path) as clip:
+            sub = clip.subclip(0, min(trim_seconds, clip.duration))
+            sub.write_videofile(trimmed, codec="libx264", audio_codec="aac",
+                                verbose=False, logger=None)
+
+        # MatAnyone → alpha PNGs (lossless)
         try:
-            alpha_i = _maybe_step(inference_core, img)  # should return HxW float 0..1
+            from .matanyone_adapter import run_matanyone_with_reference_mask
         except Exception as e:
-            cap.release()
-            raise RuntimeError(f"MatAnyone step(frame={i}) failed: {e}")
-
-        _write_alpha_png(alpha_i, os.path.join(alpha_dir, f"alpha_{i:06d}.png"))
-        i += 1
-
-    cap.release()
-    logger.info(f"Wrote {i} alpha PNG frames to {alpha_dir}")
-    return alpha_dir
-
-# ---------- Fallback/IO utilities
-
-def _maybe_step_with_mask(core, img_chw01: np.ndarray, mask_01: np.ndarray):
-    """
-    Initialize MatAnyone with first frame using step(image, mask/trimap).
-    
-    FIXED: Removed encode_all_images branch that was causing signature errors.
-    """
-    step = getattr(core, "step", None)
-    if not callable(step):
-        raise RuntimeError("MatAnyone core has no callable 'step' method")
-    
-    try:
-        # Try step(image, mask) - most common pattern for video matting
-        return step(img_chw01, mask_01)
-    except TypeError as e:
-        logger.warning(f"step(image, mask) failed: {e}")
-        
-        # Try with keyword arguments
-        sig = inspect.signature(step)
-        names = list(sig.parameters.keys())
-        
-        if "image" in names and "mask" in names:
-            return step(image=img_chw01, mask=mask_01)
-        elif "image" in names and "trimap" in names:
-            return step(image=img_chw01, trimap=mask_01)
-        elif "guidance" in names:
-            return step(img_chw01, guidance=mask_01)
-        else:
-            # Last resort: try positional with single image (no mask)
-            logger.warning("Could not pass mask to step(), trying image only")
-            return step(img_chw01)
-
-def _maybe_step(core, img_chw01: np.ndarray) -> np.ndarray:
-    """Call step on later frames (no mask). Expect an alpha or prob map; convert to 0..1 HxW."""
-    out = core.step(img_chw01)
-    # Some implementations return dict or tuple; try to extract alpha/prob
-    if isinstance(out, dict):
-        for k in ("alpha", "mask", "prob", "pha"):
-            if k in out:
-                out = out[k]
-                break
-    out = _prob_to_alpha01(core, out)
-    return out
-
-def _prob_to_alpha01(core, x) -> np.ndarray:
-    """Convert model output to HxW float32 0..1; use core.output_prob_to_mask if present."""
-    fn = getattr(core, "output_prob_to_mask", None)
-    if callable(fn):
-        try:
-            y = fn(x)
-            y = np.asarray(y, dtype=np.float32)
-            if y.ndim == 3 and y.shape[0] in (1,3):
-                # CHW -> HW (use single channel)
-                y = y[0] if y.shape[0] == 1 else np.mean(y, axis=0)
-            return np.clip(y, 0.0, 1.0)
-        except Exception:
-            pass
-    
-    # Fallback: assume x is HW or CHW logits/prob
-    a = np.asarray(x, dtype=np.float32)
-    if a.ndim == 3 and a.shape[0] in (1,3):
-        a = a[0] if a.shape[0] == 1 else np.mean(a, axis=0)
-    
-    # If looks like logits, sigmoid
-    if a.max() > 1.5 or a.min() < 0.0:
-        a = 1.0 / (1.0 + np.exp(-a))
-    return np.clip(a, 0.0, 1.0)
-
-def _write_alpha_png(alpha01: np.ndarray, path: str):
-    """Write alpha as PNG."""
-    a = np.clip(alpha01, 0.0, 1.0).astype(np.float32)
-    a = (a * 255.0 + 0.5).astype(np.uint8)
-    cv2.imwrite(path, a)
-
-def _mp4_to_alpha_png(alpha_mp4: str, out_dir: str) -> str:
-    """Convert alpha MP4 to PNG sequence."""
-    os.makedirs(out_dir, exist_ok=True)
-    cap = cv2.VideoCapture(alpha_mp4)
-    if not cap.isOpened():
-        raise ValueError(f"Could not open alpha mp4: {alpha_mp4}")
-    
-    i = 0
-    while True:
-        ok, frame = cap.read()
-        if not ok:
-            break
-        if frame.ndim == 3:
-            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-        cv2.imwrite(os.path.join(out_dir, f"alpha_{i:06d}.png"), frame)
-        i += 1
-    cap.release()
-    logger.info(f"Converted {i} frames from MP4 to PNG")
-    return out_dir
-
-def _detect_alpha_artifact(output_dir: str) -> Optional[str]:
-    """Find a likely alpha artifact produced by process_video (mp4 or a folder)."""
-    if not os.path.exists(output_dir):
-        return None
-        
-    # prefer *_alpha.mp4 / pha.mp4
-    for name in os.listdir(output_dir):
-        low = name.lower()
-        if low.endswith(".mp4") and ("alpha" in low or "pha" in low):
-            return os.path.join(output_dir, name)
-        if low.startswith("alpha_") and low.endswith(".png"):
-            # frames were written flat in output_dir
-            return output_dir
-    
-    # any mp4?
-    for name in os.listdir(output_dir):
-        if name.lower().endswith(".mp4"):
-            return os.path.join(output_dir, name)
-    return None
+            raise ImportError(f"matanyone_adapter missing or broken: {e}")
+
+        out_dir = os.path.join(self.temp_dir, "matanyone_out")
+        os.makedirs(out_dir, exist_ok=True)
+
+        alpha_dir = run_matanyone_with_reference_mask(
+            inference_core=self.matanyone_processor,
+            video_path=trimmed,
+            reference_mask_path=trimap,
+            output_dir=out_dir,
+            write_png=True
+        )
+        if not alpha_dir or not os.path.isdir(alpha_dir):
+            raise RuntimeError("MatAnyone adapter did not create alpha PNG sequence")
+
+        # Convert alpha PNGs + source RGB → GREEN master (safe mp4v)
+        green_master = self._alpha_png_to_green_master(trimmed, alpha_dir)
+        logger.info(f"✅ GREEN master: {green_master}")
+        return green_master
+
+    def _create_reference_mask(self, video_path: str) -> str:
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
+            raise ValueError(f"Cannot open video: {video_path}")
+        total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        w     = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        h     = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+        masks: List[np.ndarray] = []
+        for i in range(min(DEFAULT_REFERENCE_FRAMES, max(1, total))):
+            cap.set(cv2.CAP_PROP_POS_FRAMES, i)
+            ok, frame_bgr = cap.read()
+            if not ok: continue
+            frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+            m = self.sam2_handler.create_mask(frame_rgb)
+            if m is None: continue
+            m = m.astype(np.float32)
+            if m.max() > 1.01: m /= 255.0
+            m = np.clip(m, 0, 1)
+            masks.append(m)
+            logger.info(f"Reference frame {i+1} processed")
+        cap.release()
+
+        if not masks:
+            # ellipse fallback
+            y, x = np.ogrid[:h, :w]
+            cy, cx = h//2, w//2
+            ry, rx = int(h*0.32), int(w*0.22)
+            fallback = (((x-cx)/rx)**2 + ((y-cy)/ry)**2 <= 1).astype(np.float32)
+            masks = [fallback]
+
+        comp = np.mean(np.stack(masks, 0), axis=0).astype(np.float32)
+        comp = np.clip(comp, 0, 1)
+        path = os.path.join(self.temp_dir, "reference_mask.png")
+        cv2.imwrite(path, (comp * 255.0 + 0.5).astype(np.uint8))
+        logger.info(f"Reference mask saved: {path}")
+        return path
+
+    def _create_trimap(self, mask_png: str) -> str:
+        m = cv2.imread(mask_png, cv2.IMREAD_GRAYSCALE)
+        if m is None:
+            raise ValueError(f"Mask not found: {mask_png}")
+        m01 = (m.astype(np.float32) / 255.0)
+
+        tri = np.zeros_like(m, dtype=np.uint8)
+        tri[m01 > 0.85] = 255
+        tri[m01 < 0.15] = 0
+        tri[(m01 >= 0.15) & (m01 <= 0.85)] = 128
+
+        # widen unknown band around edges
+        k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
+        edges = cv2.Canny(m, 40, 120)
+        tri[cv2.dilate(edges, k, 1) > 0] = 128
+
+        out = os.path.join(self.temp_dir, "trimap.png")
+        cv2.imwrite(out, tri)
+        logger.info(f"Trimap created: {out}")
+        return out
+
+    def _alpha_png_to_green_master(self, rgb_video: str, alpha_dir: str) -> str:
+        # Gather PNGs
+        pngs = sorted(p for p in os.listdir(alpha_dir) if p.startswith("alpha_") and p.endswith(".png"))
+        if not pngs: raise RuntimeError(f"No alpha PNGs in {alpha_dir}")
+
+        cap = cv2.VideoCapture(rgb_video)
+        if not cap.isOpened(): raise ValueError(f"Cannot open {rgb_video}")
+        fps = cap.get(cv2.CAP_PROP_FPS) or 25
+        w   = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        h   = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+        out_path = os.path.join(self.temp_dir, "greenscreen_master.mp4")
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # safe on Spaces
+        out = cv2.VideoWriter(out_path, fourcc, fps, (w, h))
+        if not out.isOpened():
+            raise RuntimeError("Failed to open VideoWriter for GREEN master (mp4v)")
+
+        i = 0
+        last_alpha = np.ones((h, w), dtype=np.float32)
+        while True:
+            ok, frame_bgr = cap.read()
+            if not ok: break
+            # original frame in RGB for math
+            frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
+
+            if i < len(pngs):
+                a = cv2.imread(os.path.join(alpha_dir, pngs[i]), cv2.IMREAD_GRAYSCALE)
+                if a is not None:
+                    if (a.shape[1], a.shape[0]) != (w, h):
+                        a = cv2.resize(a, (w, h), interpolation=cv2.INTER_LINEAR)
+                    last_alpha = np.clip(a.astype(np.float32)/255.0, 0, 1)
+
+            a3 = np.dstack([last_alpha]*3)
+            green_rgb = np.array(PURE_GREEN_RGB, dtype=np.float32)/255.0
+            green_img = np.ones_like(frame_rgb) * green_rgb
+
+            mixed_rgb = a3*frame_rgb + (1.0 - a3)*green_img
+            mixed_bgr = cv2.cvtColor((np.clip(mixed_rgb*255.0,0,255)).astype(np.uint8), cv2.COLOR_RGB2BGR)
+
+            # Debug probe for "black video" diagnosis (first frame)
+            if i == 0:
+                cv2.imwrite(os.path.join(self.temp_dir, "probe_green_000.jpg"), mixed_bgr)
+                logger.info("Saved probe_green_000.jpg")
+
+            out.write(mixed_bgr)
+            i += 1
+
+        cap.release()
+        out.release()
+        logger.info(f"GREEN master written: {out_path}")
+        return out_path
+
+    # =========================
+    # STAGE 2: KEY & COMPOSITE
+    # =========================
+
+    def _stage2_chroma_key_replacement(self, green_master: str, bg_image: Optional[str], preset: str) -> str:
+        # Background as RGB
+        if bg_image and os.path.exists(bg_image):