EOD

rembg_video.py

@@ -1,4 +1,5 @@
 import argparse
+import io
 import os
 import ffmpeg
 import pathlib
@@ -24,14 +25,14 @@ parser.add_argument(
 parser.add_argument(
     "--model",
     type=str,
-    default="u2net-human-seg",
+    default="u2net_human_seg",
     help="rembg model to use (default: birefnet-general-lite)",
 )
 parser.add_argument(
     "--workers",
     type=int,
-    default=1,
-    help="Number of concurrent processing workers (default: 1)",
+    default=os.cpu_count() or 4,
+    help="Number of concurrent processing workers (default: cpu_count)",
 )
 parser.add_argument(
     "--smooth",
@@ -52,15 +53,108 @@ parser.add_argument(
     help="Number of processed frames to buffer before writing (default: 8)",
 )
 parser.add_argument(
-    "--skip-smooth", action="store_true", help="Skips temporal mask smoothing"
+    "--smooth-workers",
+    type=int,
+    default=os.cpu_count() or 4,
+    help="Number of threads to use for temporal mask smoothing (default: cpu count)",
 )
 args = parser.parse_args()
 
+
+def is_oom_error(exc):
+    text = str(exc).lower()
+    return any(
+        phrase in text
+        for phrase in (
+            "out of memory",
+            "cuda out of memory",
+            "failed to allocate",
+            "oom",
+            "memory error",
+        )
+    )
+
+
+def image_to_png_bytes(image):
+    with io.BytesIO() as buffer:
+        image.save(buffer, format="PNG")
+        return buffer.getvalue()
+
+
+def remove_with_mask_fallback(image_bytes, session, scales=(1.0, 0.8, 0.6, 0.4)):
+    original = Image.open(io.BytesIO(image_bytes)).convert("RGBA")
+    width, height = original.size
+    last_exc = None
+
+    for scale in scales:
+        try:
+            if scale == 1.0:
+                return remove(image_bytes, session=session)
+
+            resized = original.resize(
+                (
+                    max(1, int(width * scale)),
+                    max(1, int(height * scale)),
+                ),
+                Image.Resampling.LANCZOS,
+            )
+            with io.BytesIO() as buffer:
+                resized.save(buffer, format="PNG")
+                scaled_bytes = buffer.getvalue()
+
+            scaled_output = remove(scaled_bytes, session=session)
+            if isinstance(scaled_output, (bytes, bytearray)):
+                scaled_output_bytes = bytes(scaled_output)
+            elif isinstance(scaled_output, np.ndarray):
+                with io.BytesIO() as buffer:
+                    Image.fromarray(scaled_output).save(buffer, format="PNG")
+                    scaled_output_bytes = buffer.getvalue()
+            elif isinstance(scaled_output, Image.Image):
+                with io.BytesIO() as buffer:
+                    scaled_output.save(buffer, format="PNG")
+                    scaled_output_bytes = buffer.getvalue()
+            else:
+                raise RuntimeError(
+                    "Unexpected rembg remove() result type during mask fallback."
+                )
+            alpha = (
+                Image.open(io.BytesIO(scaled_output_bytes))
+                .convert("RGBA")
+                .getchannel("A")
+            )
+            alpha = alpha.resize((width, height), Image.Resampling.LANCZOS)
+
+            output_full = original.copy()
+            output_full.putalpha(alpha)
+            return image_to_png_bytes(output_full)
+
+        except Exception as exc:
+            last_exc = exc
+            if not is_oom_error(exc):
+                raise
+            if scale == scales[-1]:
+                raise RuntimeError(
+                    "Insufficient GPU memory: background removal failed even after fallback downscales."
+                ) from exc
+            print(
+                f"OOM detected during removal at scale {scale:.2f}; retrying with lower-resolution mask...",
+                flush=True,
+            )
+
+    if last_exc is not None:
+        raise RuntimeError(
+            "Background removal failed unexpectedly during fallback."
+        ) from last_exc
+    raise RuntimeError("Background removal failed unexpectedly.")
+
+
 # Extract video info
 probe = ffmpeg.probe(args.input)
 video_stream = next(
     (stream for stream in probe["streams"] if stream["codec_type"] == "video"), None
 )
+if video_stream is None:
+    raise ValueError(f"No video stream found in input file: {args.input}")
 width = int(video_stream["width"])
 height = int(video_stream["height"])
 whstr = str(width) + "x" + str(height)
@@ -73,7 +167,7 @@ processed_dir = os.path.join(str(pathlib.Path(__file__).parent.absolute()), "pro
 if not os.path.isdir(frames_dir):
     os.mkdir(frames_dir)
 stream = ffmpeg.input(args.input)
-stream = ffmpeg.output(stream, os.path.join(frames_dir, "%04d.png"))
+stream = ffmpeg.output(stream, os.path.join(frames_dir, "%04d.bmp"))
 ffmpeg.run(stream)
 
 _SENTINEL = object()
@@ -87,7 +181,9 @@ try:
     total_files = len(files)
 
     print(f"Loading rembg session (model={args.model})...", flush=True)
-    session = new_session(args.model, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
+    session = new_session(
+        args.model, providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
+    )
 
     read_queue = Queue(maxsize=args.read_ahead)
     write_queue = Queue(maxsize=args.write_buffer)
@@ -118,7 +214,7 @@ try:
                     break
                 idx, file, input_data = item
                 print(f"Processing frame {idx}/{total_files}: {file}", flush=True)
-                output_data = remove(input_data, session=session)
+                output_data = remove_with_mask_fallback(input_data, session=session)
                 write_queue.put((idx, file, output_data))
         except Exception as e:
             errors.append(e)
@@ -162,48 +258,94 @@ try:
         raise errors[0]
 
     # Temporal mask smoothing
-    if not args.skip_smooth and args.smooth > 0:
+    if args.smooth > 0:
         files = sorted(os.listdir(processed_dir))
         total = len(files)
         window = args.smooth
         half = window // 2
-        print(f"Applying temporal mask smoothing (window={window})...", flush=True)
+        print(
+            f"Applying temporal mask smoothing (window={window}, "
+            f"workers={args.smooth_workers})...",
+            flush=True,
+        )
 
-        buf = {}  # read_idx -> (filename, np.ndarray RGBA)
+        smoothing_errors = []
+        progress_lock = threading.Lock()
+        progress_count = [0]
 
-        for read_idx in range(total + half):
-            if read_idx < total:
-                file = files[read_idx]
-                img = Image.open(os.path.join(processed_dir, file)).convert("RGBA")
-                buf[read_idx] = (file, np.array(img))
+        # Cache decoded alpha channels so overlapping windows don't re-decode
+        # the same PNG repeatedly.
+        alpha_cache = {}
+        alpha_cache_lock = threading.Lock()
 
-            write_idx = read_idx - half
-            if 0 <= write_idx < total:
+        def get_alpha(idx):
+            with alpha_cache_lock:
+                cached = alpha_cache.get(idx)
+            if cached is not None:
+                return cached
+            file = files[idx]
+            img = Image.open(os.path.join(processed_dir, file)).convert("RGBA")
+            alpha = np.array(img)[:, :, 3].astype(np.float32)
+            with alpha_cache_lock:
+                alpha_cache[idx] = alpha
+            return alpha
+
+        def smooth_frame(write_idx):
+            try:
                 start = max(0, write_idx - half)
                 end = min(total - 1, write_idx + half)
-                alphas = np.stack(
-                    [
-                        buf[j][1][:, :, 3].astype(np.float32)
-                        for j in range(start, end + 1)
-                    ]
-                )
+                alphas = np.stack([get_alpha(j) for j in range(start, end + 1)])
                 smoothed_alpha = np.mean(alphas, axis=0).astype(np.uint8)
-                filename, arr = buf[write_idx]
-                out_arr = arr.copy()
+
+                filename = files[write_idx]
+                out_img = Image.open(os.path.join(processed_dir, filename)).convert(
+                    "RGBA"
+                )
+                out_arr = np.array(out_img)
                 out_arr[:, :, 3] = smoothed_alpha
                 Image.fromarray(out_arr).save(os.path.join(processed_dir, filename))
-                print(f"Smoothed frame {write_idx + 1}/{total}: {filename}", flush=True)
 
-                drop_idx = write_idx - half
-                if drop_idx in buf:
-                    del buf[drop_idx]
+                with progress_lock:
+                    progress_count[0] += 1
+                    print(
+                        f"Smoothed frame {progress_count[0]}/{total}: {filename}",
+                        flush=True,
+                    )
+            except Exception as e:
+                smoothing_errors.append(e)
+
+        smooth_queue = Queue()
+        for write_idx in range(total):
+            smooth_queue.put(write_idx)
+
+        def smoothing_worker():
+            while True:
+                try:
+                    write_idx = smooth_queue.get_nowait()
+                except Exception:
+                    return
+                if smoothing_errors:
+                    return
+                smooth_frame(write_idx)
+
+        smoothing_threads = [
+            threading.Thread(target=smoothing_worker, daemon=True)
+            for _ in range(max(1, args.smooth_workers))
+        ]
+        for t in smoothing_threads:
+            t.start()
+        for t in smoothing_threads:
+            t.join()
+
+        if smoothing_errors:
+            raise smoothing_errors[0]
 
     # Output video
     output_file = pathlib.Path(args.o)
     output_file.parent.mkdir(exist_ok=True, parents=True)
 
     stream = ffmpeg.input(
-        os.path.join(processed_dir, "%04d.png"),
+        os.path.join(processed_dir, "%04d.bmp"),
         r=framerate,
         f="image2",
         s=whstr,