Update processing/video/video_processor.py

processing/video/video_processor.py

@@ -2,28 +2,31 @@
 """
 Compatibility shim: CoreVideoProcessor
 
-Bridges the legacy import
-    from processing.video.video_processor import CoreVideoProcessor
-to the modern pipeline functions in utils (segment, refine, composite),
-using whatever models provider is passed in (e.g., models.loaders.ModelLoader).
-
-Requirements for the models provider:
-- get_sam2()      -> predictor or None
-- get_matanyone() -> InferenceCore or compatible (or None)
+Adds:
+- Optional NVENC hardware encoding (fast, high quality) with ffmpeg.
+- Optional model-only downscale (max_model_size) to speed up MatAnyOne WITHOUT reducing output resolution.
+
+Pipeline:
+  frame 0: SAM2 segmentation → MatAnyOne refine (seed) → composite → write
+  frames 1..N: MatAnyOne propagate (no mask) → composite → write
 """
 
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Optional, Dict, Any, Callable
+from typing import Optional, Dict, Any, Callable, List
 import time
 import threading
+import shutil
+import subprocess
+import os
+import sys
 
 import cv2
 import numpy as np
 import torch
 
-# Logger (fallback to std logging if your project logger isn't available)
+# Logger (fallback to std logging if project logger not available)
 try:
     from utils.logging_setup import make_logger
     _log = make_logger("processing.video.video_processor")
@@ -32,7 +35,7 @@
     logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
     _log = logging.getLogger(__name__)
 
-# New, hardened utils (device-safe, SAM2↔MatAnyOne interop)
+# Hardened utils (device-safe, SAM2↔MatAnyOne interop)
 from utils import (
     segment_person_hq,
     refine_mask_hq,
@@ -43,6 +46,9 @@
 )
 
 
+# -----------------------------
+# Config
+# -----------------------------
 @dataclass
 class ProcessorConfig:
     # Use a valid preset key from PROFESSIONAL_BACKGROUNDS (e.g., "office", "studio", …)
@@ -50,14 +56,149 @@ class ProcessorConfig:
     # None -> keep source fps (if available), else default to 25.0
     write_fps: Optional[float] = None
 
-
+    # ▼ Performance knobs that DO NOT hurt final quality
+    # Downscale only for model inference when min(H,W) > max_model_size; composite at full-res
+    max_model_size: Optional[int] = None    # e.g., 1280 or 1024; None = disabled
+
+    # NVENC hardware encoding (requires ffmpeg with nv-codec, e.g., on T4/L4/RTX)
+    use_nvenc: bool = True
+    nvenc_codec: str = "h264"               # "h264" | "hevc"
+    nvenc_preset: str = "p5"                # p1..p7; lower is faster, p5≈HQ
+    nvenc_cq: int = 19                      # 14–24 typical; lower = higher quality
+    nvenc_tune_hq: bool = True              # add -tune hq
+    # Set pix_fmt=yuv420p for compatibility
+    nvenc_pix_fmt: str = "yuv420p"
+
+
+# -----------------------------
+# ffmpeg NVENC writer (via pipe)
+# -----------------------------
+class _FFmpegNVENCWriter:
+    def __init__(self, path: str, width: int, height: int, fps: float,
+                 codec: str = "h264", preset: str = "p5", cq: int = 19, tune_hq: bool = True,
+                 pix_fmt: str = "yuv420p"):
+        self.path = path
+        self.width = int(width)
+        self.height = int(height)
+        self.fps = float(fps)
+        self.codec = codec
+        self.preset = preset
+        self.cq = int(cq)
+        self.tune_hq = bool(tune_hq)
+        self.pix_fmt = pix_fmt
+        self.proc: Optional[subprocess.Popen] = None
+
+        self._start()
+
+    @staticmethod
+    def _which_ffmpeg() -> Optional[str]:
+        return shutil.which("ffmpeg")
+
+    @staticmethod
+    def _encoder_name(codec: str) -> str:
+        return "h264_nvenc" if codec.lower() == "h264" else "hevc_nvenc"
+
+    @classmethod
+    def is_available(cls, codec: str) -> bool:
+        ffmpeg = cls._which_ffmpeg()
+        if not ffmpeg:
+            return False
+        # quick capability check: run a tiny encode to null
+        enc = cls._encoder_name(codec)
+        cmd = [
+            ffmpeg, "-hide_banner", "-loglevel", "error",
+            "-f", "lavfi", "-i", f"color=c=black:s=16x16:d=0.04:r=10",
+            "-c:v", enc, "-f", "null", "-"
+        ]
+        try:
+            subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
+            return True
+        except Exception:
+            return False
+
+    def _start(self):
+        ffmpeg = self._which_ffmpeg()
+        if not ffmpeg:
+            raise RuntimeError("ffmpeg not found on PATH")
+
+        enc = self._encoder_name(self.codec)
+        cmd = [
+            ffmpeg, "-hide_banner", "-loglevel", "error",
+            "-f", "rawvideo",
+            "-pix_fmt", "rgb24",
+            "-s", f"{self.width}x{self.height}",
+            "-r", f"{self.fps:.06f}",
+            "-i", "-",
+            "-an",
+            "-c:v", enc,
+            "-preset", self.preset,
+            "-rc", "vbr_hq",
+            "-cq", str(self.cq),
+            "-b:v", "0",
+            "-pix_fmt", self.nvenc_pix_fmt if hasattr(self, "nvenc_pix_fmt") else "yuv420p",
+        ]
+        if self.tune_hq:
+            cmd += ["-tune", "hq"]
+        cmd += ["-y", self.path]
+
+        _log.info("Starting NVENC writer via ffmpeg: %s", " ".join(cmd))
+        self.proc = subprocess.Popen(cmd, stdin=subprocess.PIPE)
+
+    def write_rgb(self, frame_rgb: np.ndarray) -> None:
+        if self.proc is None or self.proc.stdin is None:
+            raise RuntimeError("NVENC writer is not started")
+        assert frame_rgb.dtype == np.uint8 and frame_rgb.ndim == 3 and frame_rgb.shape[2] == 3, \
+            "write_rgb expects HxWx3 uint8 RGB frame"
+        self.proc.stdin.write(frame_rgb.tobytes())
+
+    def is_opened(self) -> bool:
+        return self.proc is not None and self.proc.poll() is None
+
+    def release(self) -> None:
+        if self.proc is not None:
+            try:
+                if self.proc.stdin:
+                    self.proc.stdin.flush()
+                    self.proc.stdin.close()
+            finally:
+                self.proc.wait(timeout=10)
+            self.proc = None
+
+
+# -----------------------------
+# Helpers: scaling for model
+# -----------------------------
+def _model_scale(width: int, height: int, max_model_size: Optional[int]) -> float:
+    if not max_model_size:
+        return 1.0
+    m = min(width, height)
+    if m <= max_model_size:
+        return 1.0
+    return max_model_size / float(m)
+
+
+def _resize_rgb(img_rgb: np.ndarray, scale: float, interp=cv2.INTER_AREA) -> np.ndarray:
+    if scale == 1.0:
+        return img_rgb
+    h, w = img_rgb.shape[:2]
+    new_w = max(1, int(round(w * scale)))
+    new_h = max(1, int(round(h * scale)))
+    return cv2.resize(img_rgb, (new_w, new_h), interpolation=interp)
+
+
+def _resize_mask(mask_hw: np.ndarray, shape_hw: tuple[int, int], interp=cv2.INTER_CUBIC) -> np.ndarray:
+    H, W = shape_hw
+    return cv2.resize(mask_hw, (W, H), interpolation=interp)
+
+
+# -----------------------------
+# CoreVideoProcessor
+# -----------------------------
 class CoreVideoProcessor:
     """
-    Minimal, safe implementation used by app entrypoint.
-    It relies on a models provider (e.g., ModelLoader) that implements:
+    Uses a models provider (ModelLoader) with:
         - get_sam2()
         - get_matanyone()
-
     Supports progress callback and cancellation via stop_event.
     """
 
@@ -103,10 +244,9 @@ def _prepare_background_from_config(
             self.log.warning("Unknown background preset '%s'; using 'office'.", choice)
             choice = "office"
 
-        bg_rgb = create_professional_background(choice, width, height)  # returns RGB
-        return bg_rgb
+        return create_professional_background(choice, width, height)  # RGB
 
-    # ---------- Full video pipeline (first-frame seed + propagate) ----------
+    # ---------- Full video pipeline ----------
     def process_video(
         self,
         input_path: str,
@@ -119,12 +259,8 @@ def process_video(
         Process a full video with live progress and optional cancel.
         progress_callback(current_frame, total_frames, fps_live)
 
-        Pipeline:
-          - Read video (OpenCV)
-          - Build background (once)
-          - Frame 0: SAM2 segmentation → MatAnyOne refine (seed)
-          - Frames 1..N: MatAnyOne propagate (no mask)
-          - Composite each frame and write to MP4
+        We optionally downscale frames ONLY for model inference (not for output),
+        and optionally use NVENC to encode the final RGB frames.
         """
         # Validate input video
         ok = validate_video_file(input_path)
@@ -142,16 +278,11 @@ def process_video(
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
 
         fps_out = self.config.write_fps or (src_fps if src_fps and src_fps > 0 else 25.0)
-        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-        writer = cv2.VideoWriter(output_path, fourcc, float(fps_out), (width, height))
-        if not writer.isOpened():
-            cap.release()
-            raise RuntimeError(f"Could not open writer for: {output_path}")
 
-        # Build background (RGB)
+        # Background (RGB)
         background_rgb = self._prepare_background_from_config(bg_config, width, height)
 
-        # Models (allow fallbacks provided by app)
+        # Models
         predictor = None
         mat_core = None
         try:
@@ -165,14 +296,55 @@ def process_video(
         except Exception as e:
             self.log.warning("MatAnyOne core unavailable: %s", e)
 
-        # Device (only used by helpers internally; we keep tensors on that device)
+        # Device
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.log.info("Starting processing on device=%s (size=%dx%d, fps_out=%.2f, frames=%s)",
+        self.log.info("Start processing on %s (%dx%d @ %.2ffps, frames=%s)",
                       device, width, height, float(fps_out), total_frames or "unknown")
 
+        # Decide writer: NVENC or OpenCV
+        writer_rgb = None
+        nvenc_ok = False
+        if self.config.use_nvenc:
+            nvenc_ok = _FFmpegNVENCWriter.is_available(self.config.nvenc_codec)
+        if nvenc_ok:
+            try:
+                writer_rgb = _FFmpegNVENCWriter(
+                    path=output_path,
+                    width=width,
+                    height=height,
+                    fps=float(fps_out),
+                    codec=self.config.nvenc_codec,
+                    preset=self.config.nvenc_preset,
+                    cq=self.config.nvenc_cq,
+                    tune_hq=self.config.nvenc_tune_hq,
+                    pix_fmt=self.config.nvenc_pix_fmt
+                )
+                self.log.info("Using NVENC (%s) via ffmpeg.", self.config.nvenc_codec)
+            except Exception as e:
+                self.log.warning("NVENC init failed (%s). Falling back to OpenCV writer.", e)
+                writer_rgb = None
+
+        if writer_rgb is None:
+            # Fallback to OpenCV VideoWriter (CPU encoding)
+            fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+            writer = cv2.VideoWriter(output_path, fourcc, float(fps_out), (width, height))
+            if not writer.isOpened():
+                cap.release()
+                raise RuntimeError(f"Could not open writer for: {output_path}")
+            self.log.info("Using OpenCV writer (mp4v).")
+        else:
+            writer = None  # mutually exclusive
+
+        # Model-only scale
+        scale = _model_scale(width, height, self.config.max_model_size)
+        if scale != 1.0:
+            self.log.info("Model-only downscale enabled: scale=%.4f (max_model_size=%d)", scale, self.config.max_model_size or -1)
+        else:
+            self.log.info("Model-only downscale disabled (processing at full resolution for model).")
+
         frame_count = 0
         start_time = time.time()
-        refined_mask_prev: Optional[np.ndarray] = None
+        refined_mask_prev_model: Optional[np.ndarray] = None  # mask at model resolution
 
         try:
             # -------- First frame (seed) --------
@@ -182,33 +354,43 @@ def process_video(
 
             f0_rgb = cv2.cvtColor(f0_bgr, cv2.COLOR_BGR2RGB)
 
-            # Segmentation (SAM2 preferred, else fallback)
-            m0_hw = segment_person_hq(
+            # Segmentation (SAM2 preferred, else fallback) at FULL resolution
+            m0_hw_full = segment_person_hq(
                 frame_rgb=f0_rgb,
                 use_sam2=True,
                 sam2_predictor=predictor
             )
-            if m0_hw is None:
-                # As an absolute last resort, use a solid foreground mask (keeps pipeline alive)
+            if m0_hw_full is None:
                 self.log.warning("First-frame segmentation failed; using full-foreground mask.")
-                m0_hw = np.ones((f0_rgb.shape[0], f0_rgb.shape[1]), dtype=np.float32)
+                m0_hw_full = np.ones((f0_rgb.shape[0], f0_rgb.shape[1]), dtype=np.float32)
 
-            # Refine / seed MatAnyOne (first_frame=True makes the helper pass the mask)
-            refined_mask_0 = refine_mask_hq(
-                mask_hw_float01=m0_hw,
-                frame_rgb=f0_rgb,
+            # Prepare model-resolution inputs (downscale if needed)
+            f0_rgb_model = _resize_rgb(f0_rgb, scale, interp=cv2.INTER_AREA)
+            m0_hw_model  = _resize_mask(m0_hw_full, f0_rgb_model.shape[:2], interp=cv2.INTER_AREA)
+
+            # Refine / seed MatAnyOne at model resolution
+            refined_mask_model_0 = refine_mask_hq(
+                mask_hw_float01=m0_hw_model,
+                frame_rgb=f0_rgb_model,
                 use_matanyone=True,
                 mat_core=mat_core,
                 first_frame=True,
                 device=device
             )
-            refined_mask_prev = refined_mask_0
+            refined_mask_prev_model = refined_mask_model_0
 
-            # Composite & write
-            comp0_rgb = replace_background_hq(f0_rgb, refined_mask_0, background_rgb)
-            writer.write(cv2.cvtColor(comp0_rgb, cv2.COLOR_RGB2BGR))
-            frame_count = 1
+            # Upscale matte to full resolution for compositing
+            refined_mask_full_0 = _resize_mask(refined_mask_model_0, f0_rgb.shape[:2], interp=cv2.INTER_CUBIC)
+
+            # Composite & write (RGB)
+            comp0_rgb = replace_background_hq(f0_rgb, refined_mask_full_0, background_rgb)
 
+            if writer_rgb is not None:
+                writer_rgb.write_rgb(comp0_rgb)
+            else:
+                writer.write(cv2.cvtColor(comp0_rgb, cv2.COLOR_RGB2BGR))
+
+            frame_count = 1
             if progress_callback:
                 elapsed = time.time() - start_time
                 fps_live = frame_count / elapsed if elapsed > 0 else 0.0
@@ -227,24 +409,31 @@ def process_video(
                 if not ret:
                     break
 
-                frgb = cv2.cvtColor(fbgr, cv2.COLOR_BGR2RGB)
+                frgb_full = cv2.cvtColor(fbgr, cv2.COLOR_BGR2RGB)
+                frgb_model = _resize_rgb(frgb_full, scale, interp=cv2.INTER_AREA)
 
-                # Propagate (first_frame=False -> mask ignored internally, MatAnyOne uses memory)
-                refined_mask_t = refine_mask_hq(
-                    mask_hw_float01=refined_mask_prev if refined_mask_prev is not None else m0_hw,
-                    frame_rgb=frgb,
+                # Propagate at model resolution (first_frame=False → internal memory used)
+                refined_mask_model_t = refine_mask_hq(
+                    mask_hw_float01=refined_mask_prev_model if refined_mask_prev_model is not None else m0_hw_model,
+                    frame_rgb=frgb_model,
                     use_matanyone=True,
                     mat_core=mat_core,
                     first_frame=False,
                     device=device
                 )
-                refined_mask_prev = refined_mask_t
+                refined_mask_prev_model = refined_mask_model_t
 
-                comp_rgb = replace_background_hq(frgb, refined_mask_t, background_rgb)
-                writer.write(cv2.cvtColor(comp_rgb, cv2.COLOR_RGB2BGR))
+                # Upscale matte to full-res for compositing
+                refined_mask_full_t = _resize_mask(refined_mask_model_t, frgb_full.shape[:2], interp=cv2.INTER_CUBIC)
 
-                frame_count += 1
+                comp_rgb = replace_background_hq(frgb_full, refined_mask_full_t, background_rgb)
+
+                if writer_rgb is not None:
+                    writer_rgb.write_rgb(comp_rgb)
+                else:
+                    writer.write(cv2.cvtColor(comp_rgb, cv2.COLOR_RGB2BGR))
 
+                frame_count += 1
                 if progress_callback:
                     elapsed = time.time() - start_time
                     fps_live = frame_count / elapsed if elapsed > 0 else 0.0
@@ -255,7 +444,10 @@ def process_video(
 
         finally:
             cap.release()
-            writer.release()
+            if writer_rgb is not None:
+                writer_rgb.release()
+            else:
+                writer.release()
 
         self.log.info("Processed %d frames → %s", frame_count, output_path)
         return {