#!/usr/bin/env python3 """ BackgroundFX Pro - Model Loading & Utilities (Hardened) ====================================================== - Avoids heavy CUDA/Hydra work at import time - Adds timeouts to subprocess probes - Safer sys.path wiring for third_party repos - MatAnyone loader is probe-only here; actual run happens in matanyone_loader.MatAnyoneSession Changes (2025-09-16): - Aligned with torch==2.3.1+cu121 and MatAnyone v1.0.0 - Updated load_matany to apply T=1 squeeze patch before InferenceCore import - Added patch status logging and MatAnyone version - Added InferenceCore attributes logging for debugging - Fixed InferenceCore import path to matanyone.inference.inference_core """ from __future__ import annotations import os import sys import cv2 import subprocess import inspect import logging import importlib.metadata from pathlib import Path from typing import Optional, Tuple, Dict, Any, Union, Callable import numpy as np import yaml # Import torch for GPU memory monitoring try: import torch except ImportError: torch = None # -------------------------------------------------------------------------------------- # Logging (ensure a handler exists very early) # -------------------------------------------------------------------------------------- logger = logging.getLogger("backgroundfx_pro") if not logger.handlers: _h = logging.StreamHandler() _h.setFormatter(logging.Formatter("[%(asctime)s] %(levelname)s: %(message)s")) logger.addHandler(_h) logger.setLevel(logging.INFO) # Pin OpenCV threads (helps libgomp stability in Spaces) try: cv_threads = int(os.environ.get("CV_THREADS", "1")) if hasattr(cv2, "setNumThreads"): cv2.setNumThreads(cv_threads) except Exception: pass # -------------------------------------------------------------------------------------- # Optional dependencies # -------------------------------------------------------------------------------------- try: import mediapipe as mp # type: ignore _HAS_MEDIAPIPE = True except Exception: _HAS_MEDIAPIPE = False # -------------------------------------------------------------------------------------- # Path setup for third_party repos # -------------------------------------------------------------------------------------- ROOT = Path(__file__).resolve().parent.parent # project root TP_SAM2 = Path(os.environ.get("THIRD_PARTY_SAM2_DIR", ROOT / "third_party" / "sam2")).resolve() TP_MATANY = Path(os.environ.get("THIRD_PARTY_MATANY_DIR", ROOT / "third_party" / "matanyone")).resolve() def _add_sys_path(p: Path) -> None: if p.exists(): p_str = str(p) if p_str not in sys.path: sys.path.insert(0, p_str) else: logger.warning(f"third_party path not found: {p}") _add_sys_path(TP_SAM2) _add_sys_path(TP_MATANY) # -------------------------------------------------------------------------------------- # Safe Torch accessors (no top-level import) # -------------------------------------------------------------------------------------- def _torch(): try: import torch # local import avoids early CUDA init during module import return torch except Exception as e: logger.warning(f"[models.safe-torch] import failed: {e}") return None def _has_cuda() -> bool: t = _torch() if t is None: return False try: return bool(t.cuda.is_available()) except Exception as e: logger.warning(f"[models.safe-torch] cuda.is_available() failed: {e}") return False def _pick_device(env_key: str) -> str: requested = os.environ.get(env_key, "").strip().lower() has_cuda = _has_cuda() # Log all CUDA-related environment variables cuda_env_vars = { 'FORCE_CUDA_DEVICE': os.environ.get('FORCE_CUDA_DEVICE', ''), 'CUDA_MEMORY_FRACTION': os.environ.get('CUDA_MEMORY_FRACTION', ''), 'PYTORCH_CUDA_ALLOC_CONF': os.environ.get('PYTORCH_CUDA_ALLOC_CONF', ''), 'REQUIRE_CUDA': os.environ.get('REQUIRE_CUDA', ''), 'SAM2_DEVICE': os.environ.get('SAM2_DEVICE', ''), 'MATANY_DEVICE': os.environ.get('MATANY_DEVICE', ''), } logger.info(f"CUDA environment variables: {cuda_env_vars}") logger.info(f"_pick_device({env_key}): requested='{requested}', has_cuda={has_cuda}") # Force CUDA if available (empty string counts as no explicit CPU request) if has_cuda and requested not in {"cpu"}: logger.info(f"FORCING CUDA device (GPU available, requested='{requested}')") return "cuda" elif requested in {"cuda", "cpu"}: logger.info(f"Using explicitly requested device: {requested}") return requested result = "cuda" if has_cuda else "cpu" logger.info(f"Auto-selected device: {result}") return result # -------------------------------------------------------------------------------------- # Basic Utilities # -------------------------------------------------------------------------------------- def _ffmpeg_bin() -> str: return os.environ.get("FFMPEG_BIN", "ffmpeg") def _probe_ffmpeg(timeout: int = 2) -> bool: try: subprocess.run([_ffmpeg_bin(), "-version"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True, timeout=timeout) return True except Exception: return False def _ensure_dir(p: Path) -> None: p.mkdir(parents=True, exist_ok=True) def _cv_read_first_frame(video_path: Union[str, Path]) -> Tuple[Optional[np.ndarray], int, Tuple[int, int]]: cap = cv2.VideoCapture(str(video_path)) if not cap.isOpened(): return None, 0, (0, 0) fps = int(round(cap.get(cv2.CAP_PROP_FPS) or 25)) ok, frame = cap.read() w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0) h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0) cap.release() if not ok: return None, fps, (w, h) return frame, fps, (w, h) def _save_mask_png(mask: np.ndarray, path: Union[str, Path]) -> str: if mask.dtype == bool: mask = (mask.astype(np.uint8) * 255) elif mask.dtype != np.uint8: mask = np.clip(mask, 0, 255).astype(np.uint8) cv2.imwrite(str(path), mask) return str(path) def _resize_keep_ar(image: np.ndarray, target_wh: Tuple[int, int]) -> np.ndarray: tw, th = target_wh h, w = image.shape[:2] if h == 0 or w == 0 or tw == 0 or th == 0: return image scale = min(tw / w, th / h) nw, nh = max(1, int(round(w * scale))), max(1, int(round(h * scale))) resized = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_CUBIC) canvas = np.zeros((th, tw, 3), dtype=resized.dtype) x0 = (tw - nw) // 2 y0 = (th - nh) // 2 canvas[y0:y0+nh, x0:x0+nw] = resized return canvas def _video_writer(out_path: Path, fps: int, size: Tuple[int, int]) -> cv2.VideoWriter: fourcc = cv2.VideoWriter_fourcc(*"mp4v") return cv2.VideoWriter(str(out_path), fourcc, max(1, fps), size) def _mux_audio(src_video: Union[str, Path], silent_video: Union[str, Path], out_path: Union[str, Path]) -> bool: """Copy video from silent_video + audio from src_video into out_path (AAC).""" try: cmd = [ _ffmpeg_bin(), "-y", "-i", str(silent_video), "-i", str(src_video), "-map", "0:v:0", "-map", "1:a:0?", "-c:v", "copy", "-c:a", "aac", "-b:a", "192k", "-shortest", str(out_path) ] subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) return True except Exception as e: logger.warning(f"Audio mux failed; returning silent video. Reason: {e}") return False # -------------------------------------------------------------------------------------- # Compositing & Image Processing # -------------------------------------------------------------------------------------- def _refine_alpha(alpha: np.ndarray, erode_px: int = 1, dilate_px: int = 2, blur_px: float = 1.5) -> np.ndarray: if alpha.dtype != np.float32: a = alpha.astype(np.float32) if a.max() > 1.0: a = a / 255.0 else: a = alpha.copy() a_u8 = np.clip(np.round(a * 255.0), 0, 255).astype(np.uint8) if erode_px > 0: k = max(1, int(erode_px)) a_u8 = cv2.erode(a_u8, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k)), iterations=1) if dilate_px > 0: k = max(1, int(dilate_px)) a_u8 = cv2.dilate(a_u8, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k)), iterations=1) a = a_u8.astype(np.float32) / 255.0 if blur_px and blur_px > 0: rad = max(1, int(round(blur_px))) a = cv2.GaussianBlur(a, (rad | 1, rad | 1), 0) return np.clip(a, 0.0, 1.0) def _to_linear(rgb: np.ndarray, gamma: float = 2.2) -> np.ndarray: x = np.clip(rgb.astype(np.float32) / 255.0, 0.0, 1.0) return np.power(x, gamma) def _to_srgb(lin: np.ndarray, gamma: float = 2.2) -> np.ndarray: x = np.clip(lin, 0.0, 1.0) return np.clip(np.power(x, 1.0 / gamma) * 255.0, 0, 255).astype(np.uint8) def _light_wrap(bg_rgb: np.ndarray, alpha01: np.ndarray, radius: int = 5, amount: float = 0.18) -> np.ndarray: r = max(1, int(radius)) inv = 1.0 - alpha01 inv_blur = cv2.GaussianBlur(inv, (r | 1, r | 1), 0) lw = (bg_rgb.astype(np.float32) * inv_blur[..., None] * float(amount)) return lw def _despill_edges(fg_rgb: np.ndarray, alpha01: np.ndarray, amount: float = 0.35) -> np.ndarray: w = 1.0 - 2.0 * np.abs(alpha01 - 0.5) w = np.clip(w, 0.0, 1.0) hsv = cv2.cvtColor(fg_rgb.astype(np.uint8), cv2.COLOR_RGB2HSV).astype(np.float32) H, S, V = cv2.split(hsv) S = S * (1.0 - amount * w) hsv2 = cv2.merge([H, np.clip(S, 0, 255), V]) out = cv2.cvtColor(hsv2.astype(np.uint8), cv2.COLOR_HSV2RGB) return out def _composite_frame_pro( fg_rgb: np.ndarray, alpha: np.ndarray, bg_rgb: np.ndarray, erode_px: int = None, dilate_px: int = None, blur_px: float = None, lw_radius: int = None, lw_amount: float = None, despill_amount: float = None ) -> np.ndarray: erode_px = erode_px if erode_px is not None else int(os.environ.get("EDGE_ERODE", "1")) dilate_px = dilate_px if dilate_px is not None else int(os.environ.get("EDGE_DILATE", "2")) blur_px = blur_px if blur_px is not None else float(os.environ.get("EDGE_BLUR", "1.5")) lw_radius = lw_radius if lw_radius is not None else int(os.environ.get("LIGHTWRAP_RADIUS", "5")) lw_amount = lw_amount if lw_amount is not None else float(os.environ.get("LIGHTWRAP_AMOUNT", "0.18")) despill_amount = despill_amount if despill_amount is not None else float(os.environ.get("DESPILL_AMOUNT", "0.35")) a = _refine_alpha(alpha, erode_px=erode_px, dilate_px=dilate_px, blur_px=blur_px) fg_rgb = _despill_edges(fg_rgb, a, amount=despill_amount) fg_lin = _to_linear(fg_rgb) bg_lin = _to_linear(bg_rgb) lw = _light_wrap(bg_rgb, a, radius=lw_radius, amount=lw_amount) lw_lin = _to_linear(np.clip(lw, 0, 255).astype(np.uint8)) comp_lin = fg_lin * a[..., None] + bg_lin * (1.0 - a[..., None]) + lw_lin comp = _to_srgb(comp_lin) return comp # -------------------------------------------------------------------------------------- # SAM2 Integration # -------------------------------------------------------------------------------------- def _resolve_sam2_cfg(cfg_str: str) -> str: """Resolve SAM2 config path - return relative path for Hydra compatibility.""" logger.info(f"_resolve_sam2_cfg called with cfg_str={cfg_str}") # Get the third-party SAM2 directory tp_sam2 = os.environ.get("THIRD_PARTY_SAM2_DIR", "/home/user/app/third_party/sam2") logger.info(f"TP_SAM2 = {tp_sam2}") # Check if the full path exists candidate = os.path.join(tp_sam2, cfg_str) logger.info(f"Candidate path: {candidate}") logger.info(f"Candidate exists: {os.path.exists(candidate)}") if os.path.exists(candidate): # For Hydra compatibility, return just the relative path within sam2 package if cfg_str.startswith("sam2/configs/"): relative_path = cfg_str.replace("sam2/configs/", "configs/") else: relative_path = cfg_str logger.info(f"Returning Hydra-compatible relative path: {relative_path}") return relative_path # If not found, try some fallback paths fallbacks = [ os.path.join(tp_sam2, "sam2", cfg_str), os.path.join(tp_sam2, "configs", cfg_str), ] for fallback in fallbacks: logger.info(f"Trying fallback: {fallback}") if os.path.exists(fallback): # Extract relative path for Hydra if "configs/" in fallback: relative_path = "configs/" + fallback.split("configs/")[-1] logger.info(f"Returning fallback relative path: {relative_path}") return relative_path logger.warning(f"Config not found, returning original: {cfg_str}") return cfg_str def _find_hiera_config_if_hieradet(cfg_path: str) -> Optional[str]: """If config references 'hieradet', try to find a 'hiera' config.""" try: with open(cfg_path, "r") as f: data = yaml.safe_load(f) model = data.get("model", {}) or {} enc = model.get("image_encoder") or {} trunk = enc.get("trunk") or {} target = trunk.get("_target_") or trunk.get("target") if isinstance(target, str) and "hieradet" in target: for y in TP_SAM2.rglob("*.yaml"): try: with open(y, "r") as f2: d2 = yaml.safe_load(f2) or {} e2 = (d2.get("model", {}) or {}).get("image_encoder") or {} t2 = (e2.get("trunk") or {}) tgt2 = t2.get("_target_") or t2.get("target") if isinstance(tgt2, str) and ".hiera." in tgt2: logger.info(f"SAM2: switching config from 'hieradet' → 'hiera': {y}") return str(y) except Exception: continue except Exception: pass return None def load_sam2() -> Tuple[Optional[object], bool, Dict[str, Any]]: """Robust SAM2 loader with config resolution and error handling.""" meta = {"sam2_import_ok": False, "sam2_init_ok": False} try: from sam2.build_sam import build_sam2 # type: ignore from sam2.sam2_image_predictor import SAM2ImagePredictor # type: ignore meta["sam2_import_ok"] = True except Exception as e: logger.warning(f"SAM2 import failed: {e}") return None, False, meta # Check GPU memory before loading if torch and torch.cuda.is_available(): mem_before = torch.cuda.memory_allocated() / 1024**3 logger.info(f"🔍 GPU memory before SAM2 load: {mem_before:.2f}GB") device = _pick_device("SAM2_DEVICE") cfg_env = os.environ.get("SAM2_MODEL_CFG", "sam2/configs/sam2/sam2_hiera_l.yaml") cfg = _resolve_sam2_cfg(cfg_env) ckpt = os.environ.get("SAM2_CHECKPOINT", "") def _try_build(cfg_path: str): logger.info(f"_try_build called with cfg_path: {cfg_path}") params = set(inspect.signature(build_sam2).parameters.keys()) logger.info(f"build_sam2 parameters: {list(params)}") kwargs = {} if "config_file" in params: kwargs["config_file"] = cfg_path logger.info(f"Using config_file parameter: {cfg_path}") elif "model_cfg" in params: kwargs["model_cfg"] = cfg_path logger.info(f"Using model_cfg parameter: {cfg_path}") if ckpt: if "checkpoint" in params: kwargs["checkpoint"] = ckpt elif "ckpt_path" in params: kwargs["ckpt_path"] = ckpt elif "weights" in params: kwargs["weights"] = ckpt if "device" in params: kwargs["device"] = device try: logger.info(f"Calling build_sam2 with kwargs: {kwargs}") result = build_sam2(**kwargs) logger.info(f"build_sam2 succeeded with kwargs") # Log actual device of the model if hasattr(result, 'device'): logger.info(f"SAM2 model device: {result.device}") elif hasattr(result, 'image_encoder') and hasattr(result.image_encoder, 'device'): logger.info(f"SAM2 model device: {result.image_encoder.device}") return result except TypeError as e: logger.info(f"build_sam2 kwargs failed: {e}, trying positional args") pos = [cfg_path] if ckpt: pos.append(ckpt) if "device" not in kwargs: pos.append(device) logger.info(f"Calling build_sam2 with positional args: {pos}") result = build_sam2(*pos) logger.info(f"build_sam2 succeeded with positional args") return result try: try: sam = _try_build(cfg) except Exception: alt_cfg = _find_hiera_config_if_hieradet(cfg) if alt_cfg: sam = _try_build(alt_cfg) else: raise if sam is not None: predictor = SAM2ImagePredictor(sam) meta["sam2_init_ok"] = True meta["sam2_device"] = device return predictor, True, meta else: return None, False, meta except Exception as e: logger.error(f"SAM2 loading failed: {e}") return None, False, meta def run_sam2_mask(predictor: object, first_frame_bgr: np.ndarray, point: Optional[Tuple[int, int]] = None, auto: bool = False) -> Tuple[Optional[np.ndarray], bool]: """Return (mask_uint8_0_255, ok).""" if predictor is None: return None, False try: rgb = cv2.cvtColor(first_frame_bgr, cv2.COLOR_BGR2RGB) predictor.set_image(rgb) if auto: h, w = rgb.shape[:2] box = np.array([int(0.05*w), int(0.05*h), int(0.95*w), int(0.95*h)]) masks, _, _ = predictor.predict(box=box) elif point is not None: x, y = int(point[0]), int(point[1]) pts = np.array([[x, y]], dtype=np.int32) labels = np.array([1], dtype=np.int32) masks, _, _ = predictor.predict(point_coords=pts, point_labels=labels) else: h, w = rgb.shape[:2] box = np.array([int(0.1*w), int(0.1*h), int(0.9*w), int(0.9*h)]) masks, _, _ = predictor.predict(box=box) if masks is None or len(masks) == 0: return None, False m = masks[0].astype(np.uint8) * 255 return m, True except Exception as e: logger.warning(f"SAM2 mask failed: {e}") return None, False def _refine_mask_grabcut(image_bgr: np.ndarray, mask_u8: np.ndarray, iters: int = None, trimap_erode: int = None, trimap_dilate: int = None) -> np.ndarray: """Use SAM2 seed as initialization for GrabCut refinement.""" iters = int(os.environ.get("REFINE_GRABCUT_ITERS", "2")) if iters is None else int(iters) e = int(os.environ.get("REFINE_TRIMAP_ERODE", "3")) if trimap_erode is None else int(trimap_erode) d = int(os.environ.get("REFINE_TRIMAP_DILATE", "6")) if trimap_dilate is None else int(trimap_dilate) h, w = mask_u8.shape[:2] m = (mask_u8 > 127).astype(np.uint8) * 255 sure_fg = cv2.erode(m, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (max(1, e), max(1, e))), iterations=1) sure_bg = cv2.erode(255 - m, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (max(1, d), max(1, d))), iterations=1) gc_mask = np.full((h, w), cv2.GC_PR_BGD, dtype=np.uint8) gc_mask[sure_bg > 0] = cv2.GC_BGD gc_mask[sure_fg > 0] = cv2.GC_FGD bgdModel = np.zeros((1, 65), np.float64) fgdModel = np.zeros((1, 65), np.float64) try: cv2.grabCut(image_bgr, gc_mask, None, bgdModel, fgdModel, iters, cv2.GC_INIT_WITH_MASK) out = np.where((gc_mask == cv2.GC_FGD) | (gc_mask == cv2.GC_PR_FGD), 255, 0).astype(np.uint8) out = cv2.medianBlur(out, 5) return out except Exception as e: logger.warning(f"GrabCut refinement failed; using original mask. Reason: {e}") return m # -------------------------------------------------------------------------------------- # MatAnyone Integration # -------------------------------------------------------------------------------------- def load_matany() -> Tuple[Optional[object], bool, Dict[str, Any]]: """ Probe MatAnyone availability with T=1 squeeze patch for conv2d compatibility. Returns (None, available, meta); actual instantiation happens in MatAnyoneSession. """ meta = {"matany_import_ok": False, "matany_init_ok": False} enable_env = os.environ.get("ENABLE_MATANY", "1").strip().lower() if enable_env in {"0", "false", "off", "no"}: logger.info("MatAnyone disabled by ENABLE_MATANY=0.") meta["disabled"] = True return None, False, meta # Apply T=1 squeeze patch before importing InferenceCore try: from .matany_compat_patch import apply_matany_t1_squeeze_guard if apply_matany_t1_squeeze_guard(): logger.info("[MatAnyCompat] T=1 squeeze guard applied") meta["patch_applied"] = True else: logger.warning("[MatAnyCompat] T=1 squeeze patch failed; conv2d errors may occur") meta["patch_applied"] = False except Exception as e: logger.warning(f"[MatAnyCompat] Patch import failed: {e}") meta["patch_applied"] = False try: from matanyone.inference.inference_core import InferenceCore # type: ignore meta["matany_import_ok"] = True # Log MatAnyone version and InferenceCore attributes try: version = importlib.metadata.version("matanyone") logger.info(f"[MATANY] MatAnyone version: {version}") except Exception: logger.info("[MATANY] MatAnyone version unknown") logger.debug(f"[MATANY] InferenceCore attributes: {dir(InferenceCore)}") device = _pick_device("MATANY_DEVICE") repo_id = os.environ.get("MATANY_REPO_ID", "PeiqingYang/MatAnyone") meta["matany_repo_id"] = repo_id meta["matany_device"] = device return None, True, meta except Exception as e: logger.warning(f"MatAnyone import failed: {e}") return None, False, meta # -------------------------------------------------------------------------------------- # Fallback Functions # -------------------------------------------------------------------------------------- def fallback_mask(first_frame_bgr: np.ndarray) -> np.ndarray: """Prefer MediaPipe; fallback to GrabCut. Returns uint8 mask 0/255.""" h, w = first_frame_bgr.shape[:2] if _HAS_MEDIAPIPE: try: mp_selfie = mp.solutions.selfie_segmentation with mp_selfie.SelfieSegmentation(model_selection=1) as segmenter: rgb = cv2.cvtColor(first_frame_bgr, cv2.COLOR_BGR2RGB) res = segmenter.process(rgb) m = (np.clip(res.segmentation_mask, 0, 1) > 0.5).astype(np.uint8) * 255 m = cv2.medianBlur(m, 5) return m except Exception as e: logger.warning(f"MediaPipe fallback failed: {e}") # Ultimate fallback: GrabCut mask = np.zeros((h, w), np.uint8) rect = (int(0.1*w), int(0.1*h), int(0.8*w), int(0.8*h)) bgdModel = np.zeros((1, 65), np.float64) fgdModel = np.zeros((1, 65), np.float64) try: cv2.grabCut(first_frame_bgr, mask, rect, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_RECT) mask_bin = np.where((mask == cv2.GC_FGD) | (mask == cv2.GC_PR_FGD), 255, 0).astype(np.uint8) return mask_bin except Exception as e: logger.warning(f"GrabCut failed: {e}") return np.zeros((h, w), dtype=np.uint8) def composite_video(fg_path: Union[str, Path], alpha_path: Union[str, Path], bg_image_path: Union[str, Path], out_path: Union[str, Path], fps: int, size: Tuple[int, int]) -> bool: """Blend MatAnyone FG+ALPHA over background using pro compositor.""" fg_cap = cv2.VideoCapture(str(fg_path)) al_cap = cv2.VideoCapture(str(alpha_path)) if not fg_cap.isOpened() or not al_cap.isOpened(): return False w, h = size bg = cv2.imread(str(bg_image_path), cv2.IMREAD_COLOR) if bg is None: bg = np.full((h, w, 3), 127, dtype=np.uint8) bg_f = _resize_keep_ar(bg, (w, h)) if _probe_ffmpeg(): tmp_out = Path(str(out_path) + ".tmp.mp4") writer = _video_writer(tmp_out, fps, (w, h)) post_h264 = True else: writer = _video_writer(Path(out_path), fps, (w, h)) post_h264 = False ok_any = False try: while True: ok_fg, fg = fg_cap.read() ok_al, al = al_cap.read() if not ok_fg or not ok_al: break fg = cv2.resize(fg, (w, h), interpolation=cv2.INTER_CUBIC) al_gray = cv2.cvtColor(cv2.resize(al, (w, h)), cv2.COLOR_BGR2GRAY) comp = _composite_frame_pro( cv2.cvtColor(fg, cv2.COLOR_BGR2RGB), al_gray, cv2.cvtColor(bg_f, cv2.COLOR_BGR2RGB) ) writer.write(cv2.cvtColor(comp, cv2.COLOR_RGB2BGR)) ok_any = True finally: fg_cap.release() al_cap.release() writer.release() if post_h264 and ok_any: try: cmd = [ _ffmpeg_bin(), "-y", "-i", str(tmp_out), "-c:v", "libx264", "-pix_fmt", "yuv420p", "-movflags", "+faststart", str(out_path) ] subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) tmp_out.unlink(missing_ok=True) except Exception as e: logger.warning(f"ffmpeg finalize failed: {e}") Path(out_path).unlink(missing_ok=True) tmp_out.replace(out_path) return ok_any def fallback_composite(video_path: Union[str, Path], mask_path: Union[str, Path], bg_image_path: Union[str, Path], out_path: Union[str, Path]) -> bool: """Static-mask compositing using pro compositor.""" mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE) cap = cv2.VideoCapture(str(video_path)) if mask is None or not cap.isOpened(): return False w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0) h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0) fps = int(round(cap.get(cv2.CAP_PROP_FPS) or 25)) bg = cv2.imread(str(bg_image_path), cv2.IMREAD_COLOR) if bg is None: bg = np.full((h, w, 3), 127, dtype=np.uint8) mask_resized = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST) bg_f = _resize_keep_ar(bg, (w, h)) if _probe_ffmpeg(): tmp_out = Path(str(out_path) + ".tmp.mp4") writer = _video_writer(tmp_out, fps, (w, h)) use_post_ffmpeg = True else: writer = _video_writer(Path(out_path), fps, (w, h)) use_post_ffmpeg = False ok_any = False try: while True: ok, frame = cap.read() if not ok: break comp = _composite_frame_pro( cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), mask_resized, cv2.cvtColor(bg_f, cv2.COLOR_BGR2RGB) ) writer.write(cv2.cvtColor(comp, cv2.COLOR_RGB2BGR)) ok_any = True finally: cap.release() writer.release() if use_post_ffmpeg and ok_any: try: cmd = [ _ffmpeg_bin(), "-y", "-i", str(tmp_out), "-c:v", "libx264", "-pix_fmt", "yuv420p", "-movflags", "+faststart", str(out_path) ] subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) tmp_out.unlink(missing_ok=True) except Exception as e: logger.warning(f"ffmpeg H.264 finalize failed: {e}") Path(out_path).unlink(missing_ok=True) tmp_out.replace(out_path) return ok_any # -------------------------------------------------------------------------------------- # Stage-A (Transparent Export) Functions # -------------------------------------------------------------------------------------- def _checkerboard_bg(w: int, h: int, tile: int = 32) -> np.ndarray: y, x = np.mgrid[0:h, 0:w] c = ((x // tile) + (y // tile)) % 2 a = np.where(c == 0, 200, 150).astype(np.uint8) return np.stack([a, a, a], axis=-1) def _build_stage_a_rgba_vp9_from_fg_alpha( fg_path: Union[str, Path], alpha_path: Union[str, Path], out_webm: Union[str, Path], fps: int, size: Tuple[int, int], src_audio: Optional[Union[str, Path]] = None, ) -> bool: if not _probe_ffmpeg(): return False w, h = size try: cmd = [_ffmpeg_bin(), "-y", "-i", str(fg_path), "-i", str(alpha_path)] if src_audio: cmd += ["-i", str(src_audio)] fcx = f"[1:v]format=gray,scale={w}:{h},fps={fps}[al];" \ f"[0:v]scale={w}:{h},fps={fps}[fg];" \ f"[fg][al]alphamerge[outv]" cmd += ["-filter_complex", fcx, "-map", "[outv]"] if src_audio: cmd += ["-map", "2:a:0?", "-c:a", "libopus", "-b:a", "128k"] cmd += [ "-c:v", "libvpx-vp9", "-pix_fmt", "yuva420p", "-crf", os.environ.get("STAGEA_VP9_CRF", "28"), "-b:v", "0", "-row-mt", "1", "-shortest", str(out_webm), ] subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) return True except Exception as e: logger.warning(f"Stage-A VP9(alpha) build failed: {e}") return False def _build_stage_a_rgba_vp9_from_mask( video_path: Union[str, Path], mask_png: Union[str, Path], out_webm: Union[str, Path], fps: int, size: Tuple[int, int], ) -> bool: if not _probe_ffmpeg(): return False w, h = size try: cmd = [ _ffmpeg_bin(), "-y", "-i", str(video_path), "-loop", "1", "-i", str(mask_png), "-filter_complex", f"[1:v]format=gray,scale={w}:{h},fps={fps}[al];" f"[0:v]scale={w}:{h},fps={fps}[fg];" f"[fg][al]alphamerge[outv]", "-map", "[outv]", "-c:v", "libvpx-vp9", "-pix_fmt", "yuva420p", "-crf", os.environ.get("STAGEA_VP9_CRF", "28"), "-b:v", "0", "-row-mt", "1", "-shortest", str(out_webm), ] subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) return True except Exception as e: logger.warning(f"Stage-A VP9(alpha) (mask) build failed: {e}") return False def _build_stage_a_checkerboard_from_fg_alpha( fg_path: Union[str, Path], alpha_path: Union[str, Path], out_mp4: Union[str, Path], fps: int, size: Tuple[int, int], ) -> bool: fg_cap = cv2.VideoCapture(str(fg_path)) al_cap = cv2.VideoCapture(str(alpha_path)) if not fg_cap.isOpened() or not al_cap.isOpened(): return False w, h = size writer = _video_writer(Path(out_mp4), fps, (w, h)) bg = _checkerboard_bg(w, h) ok_any = False try: while True: okf, fg = fg_cap.read() oka, al = al_cap.read() if not okf or not oka: break fg = cv2.resize(fg, (w, h)) al = cv2.cvtColor(cv2.resize(al, (w, h)), cv2.COLOR_BGR2GRAY) comp = _composite_frame_pro(cv2.cvtColor(fg, cv2.COLOR_BGR2RGB), al, bg) writer.write(cv2.cvtColor(comp, cv2.COLOR_RGB2BGR)) ok_any = True finally: fg_cap.release() al_cap.release() writer.release() return ok_any def _build_stage_a_checkerboard_from_mask( video_path: Union[str, Path], mask_png: Union[str, Path], out_mp4: Union[str, Path], fps: int, size: Tuple[int, int], ) -> bool: cap = cv2.VideoCapture(str(video_path)) if not cap.isOpened(): return False w, h = size mask = cv2.imread(str(mask_png), cv2.IMREAD_GRAYSCALE) if mask is None: return False mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST) writer = _video_writer(Path(out_mp4), fps, (w, h)) bg = _checkerboard_bg(w, h) ok_any = False try: while True: ok, frame = cap.read() if not ok: break frame = cv2.resize(frame, (w, h)) comp = _composite_frame_pro(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), mask, bg) writer.write(cv2.cvtColor(comp, cv2.COLOR_RGB2BGR)) ok_any = True finally: cap.release() writer.release() return ok_any # -------------------------------------------------------------------------------------- # MatAnyone Integration # -------------------------------------------------------------------------------------- def run_matany( video_path: Union[str, Path], mask_path: Optional[Union[str, Path]], out_dir: Union[str, Path], device: Optional[str] = None, progress_callback: Optional[Callable[[float, str], None]] = None, ) -> Tuple[Path, Path]: """ Run MatAnyone streaming matting via our shape-guarded adapter. Returns (alpha_mp4_path, fg_mp4_path). Raises MatAnyError on failure. """ from .matanyone_loader import MatAnyoneSession, MatAnyError session = MatAnyoneSession(device=device, precision="auto") alpha_p, fg_p = session.process_stream( video_path=Path(video_path), seed_mask_path=Path(mask_path) if mask_path else None, out_dir=Path(out_dir), progress_cb=progress_callback, ) return alpha_p, fg_p