EOD

pyrightconfig.json

@@ -0,0 +1,6 @@
+{
+    "venvPath": ".",
+    "venv": ".venv",
+    "reportUnusedCallResult": false,
+    "reportAny": false
+}

readme.md

@@ -1,32 +1,71 @@
-# rembg_from_video
+# rembg_for_video
 
 Uses [ffmpeg-python](https://github.com/kkroening/ffmpeg-python) and [rembg](https://github.com/danielgatis/rembg) to attempt removal of a background from a video file.
 
-Two directories will be created in the same directory as the script to hold the video frames (before and after rembg is applied).
+Based on project [rembg_from_video](https://github.com/seth-tribbey/rembg_from_video) by [seth-tribbey](https://github.com/seth-tribbey)
 
-### Installation:
+## Installation:
-[Install rembg by following their instructions](https://github.com/danielgatis/rembg)
 
-rembg specifically requires Python 3.9 as of the time of this writing. Note that you must choose the `rembg[gpu]` version and configure onnxruntime accordingly if you wish to use your GPU for the image processing.
+### For Newer Architecture (GeForce 1600+ Series)
-
+```bash
-Then:
+python -m venv .venv
+pip install -r requirements.txt
 ```
-pip install ffmpeg-python
-```
-### Usage:
-```
-python .\rembg_video.py [-h] [-a] [-af AF] [-ab AB] [-ae AE] [--skip-extract] [--skip-process] input
 
-positional arguments:
+### For Pascal Arhitecture (GeForce 1000 Series)
-  input           Input video
+```bash
+python3.12 -m venv .venv
+pip install -r requirements(3.12).txt
+```
+> **Script for fixing the cudnn path on Linux:** <br>
+> export LD_LIBRARY_PATH=/path/to/kakisalmi/.venv/lib/python3.12/site-packages/nvidia/cudnn/lib:$LD_LIBRARY_PATH
+
+&nbsp;
+
+## Usage:
+```
+python .\rembg_video.py [-h] [--help] [-o] [--model] [--workers] [--smooth] [--read-ahead] [--write-buffer] input
+```
+<style>
+table {
+  border-collapse: separate;
+  border-spacing: 0;
+  border-radius: 5px;
+  overflow: hidden;
+}
+th,
+td {
+  border: 1px solid #a0a0a0;
+  padding: 8px 10px;
+  border-radius: 8px;
+  text-align: center;
+}
+td:first-child,
+th:first-child {
+  text-align: center;
+}
+td:nth-child(2),
+th:nth-child(2) {
+  text-align: left;
+}
+</style>
+<div style="display: flex; gap: 20px; align-items: flex-start;">
+  <table>
+    <tr><th colspan="2">Optional Arguments:</th></tr>
+    <tr><td>-o</td><td>Output Path</td></tr>
+    <tr><td>-h --help</td><td>Show Help</td></tr>
+    <tr><td>--model</td><td>Choose model for rembg</td></tr>
+  </table>
+
+  <table>
+    <tr><th colspan="2">Positional Arguments:</th></tr>
+    <tr><td>input</td><td>Input Video</td></tr>
+  </table>
+</div>
 
 optional arguments:
   -h, --help      show this help message and exit
-  -a              Turns on alpha matting during background removal
+  -o              Set output path (Default: export/output.mov)
-  -af AF          Alpha matting foreground threshold
+  -
-  -ab AB          Alpha matting background threshold
+  
-  -ae AE          Alpha matting erode size
-  --skip-extract  Skips ffmpeg frame extraction
-  --skip-process  Skips rembg frame processing
-  ```
   Tip: [Alpha matting can be used to refine the results](https://github.com/danielgatis/rembg#advance-usage)

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,
+    default=os.cpu_count() or 4,
-    help="Number of concurrent processing workers (default: 1)",
+    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):
+        # Cache decoded alpha channels so overlapping windows don't re-decode
-            if read_idx < total:
+        # the same PNG repeatedly.
-                file = files[read_idx]
+        alpha_cache = {}
-                img = Image.open(os.path.join(processed_dir, file)).convert("RGBA")
+        alpha_cache_lock = threading.Lock()
-                buf[read_idx] = (file, np.array(img))
 
-            write_idx = read_idx - half
+        def get_alpha(idx):
-            if 0 <= write_idx < total:
+            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(
+                alphas = np.stack([get_alpha(j) for j in range(start, end + 1)])
-                    [
-                        buf[j][1][:, :, 3].astype(np.float32)
-                        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
+                with progress_lock:
-                if drop_idx in buf:
+                    progress_count[0] += 1
-                    del buf[drop_idx]
+                    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,

rembg_video_claude.py

@@ -0,0 +1,375 @@
+import argparse
+import io
+import os
+import ffmpeg
+import pathlib
+import threading
+import numpy as np
+from queue import Queue
+from shutil import rmtree
+from PIL import Image
+from rembg import new_session, remove
+
+os.environ.setdefault(
+    "PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True,managed_memory:True"
+)
+
+# Parse args
+parser = argparse.ArgumentParser(
+    description="Applies rembg background removal to the frames of a video"
+)
+parser.add_argument("input", type=str, help="Input video")
+parser.add_argument(
+    "-o", type=str, default="export/output.mov", help="Define output path"
+)
+parser.add_argument(
+    "--model",
+    type=str,
+    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)",
+)
+parser.add_argument(
+    "--smooth",
+    type=int,
+    default=3,
+    help="Temporal mask smoothing window size in frames (default: 3, 0 to disable)",
+)
+parser.add_argument(
+    "--read-ahead",
+    type=int,
+    default=8,
+    help="Number of frames to read ahead into buffer (default: 8)",
+)
+parser.add_argument(
+    "--write-buffer",
+    type=int,
+    default=8,
+    help="Number of processed frames to buffer before writing (default: 8)",
+)
+parser.add_argument(
+    "--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)
+framerate = video_stream["avg_frame_rate"]
+
+frames_dir = os.path.join(str(pathlib.Path(__file__).parent.absolute()), "frames")
+processed_dir = os.path.join(str(pathlib.Path(__file__).parent.absolute()), "processed")
+
+# Extract input video frames
+# Always start from a clean slate: a prior crashed/interrupted run may have
+# left frames behind (possibly from a different input video, e.g. different
+# resolution), which would silently corrupt this run.
+rmtree(frames_dir, ignore_errors=True)
+os.mkdir(frames_dir)
+stream = ffmpeg.input(args.input)
+stream = ffmpeg.output(stream, os.path.join(frames_dir, "%04d.bmp"))
+ffmpeg.run(stream)
+
+_SENTINEL = object()
+
+try:
+    # Process frames with pipelined reader -> processors -> writer
+    if not os.path.isdir(processed_dir):
+        os.mkdir(processed_dir)
+
+    files = sorted(os.listdir(frames_dir))
+    total_files = len(files)
+
+    print(f"Loading rembg session (model={args.model})...", flush=True)
+    session = new_session(
+        args.model, providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
+    )
+
+    read_queue = Queue(maxsize=args.read_ahead)
+    write_queue = Queue(maxsize=args.write_buffer)
+    errors = []
+    active_processors = [
+        args.workers
+    ]  # list so processor() can mutate without nonlocal
+    active_processors_lock = threading.Lock()
+
+    def reader():
+        try:
+            for idx, file in enumerate(files, 1):
+                frame_path = os.path.join(frames_dir, file)
+                with open(frame_path, "rb") as f:
+                    data = f.read()
+                read_queue.put((idx, file, data))
+                os.remove(frame_path)
+        except Exception as e:
+            errors.append(e)
+        finally:
+            # One sentinel per worker so each one knows when to stop
+            for _ in range(args.workers):
+                read_queue.put(_SENTINEL)
+
+    def processor():
+        try:
+            while True:
+                item = read_queue.get()
+                if item is _SENTINEL:
+                    break
+                idx, file, input_data = item
+                print(f"Processing frame {idx}/{total_files}: {file}", flush=True)
+                output_data = remove_with_mask_fallback(input_data, session=session)
+                write_queue.put((idx, file, output_data))
+        except Exception as e:
+            errors.append(e)
+        finally:
+            # Signal writer only when the last processor finishes
+            with active_processors_lock:
+                active_processors[0] -= 1
+                if active_processors[0] == 0:
+                    write_queue.put(_SENTINEL)
+
+    def writer():
+        try:
+            while True:
+                item = write_queue.get()
+                if item is _SENTINEL:
+                    break
+                idx, file, output_data = item
+                with open(os.path.join(processed_dir, file), "wb") as f:
+                    f.write(output_data)
+                print(f"Written frame {idx}/{total_files}: {file}", flush=True)
+        except Exception as e:
+            errors.append(e)
+
+    reader_thread = threading.Thread(target=reader, daemon=True)
+    processor_threads = [
+        threading.Thread(target=processor, daemon=True) for _ in range(args.workers)
+    ]
+    writer_thread = threading.Thread(target=writer, daemon=True)
+
+    reader_thread.start()
+    for t in processor_threads:
+        t.start()
+    writer_thread.start()
+
+    reader_thread.join()
+    for t in processor_threads:
+        t.join()
+    writer_thread.join()
+
+    if errors:
+        raise errors[0]
+
+    # Temporal mask smoothing
+    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}, "
+            f"workers={args.smooth_workers})...",
+            flush=True,
+        )
+
+        smoothing_errors = []
+        progress_lock = threading.Lock()
+        progress_count = [0]
+        n_workers = max(1, args.smooth_workers)
+
+        # Write smoothed frames to a separate directory rather than
+        # overwriting processed_dir in place. Overlapping windows mean a
+        # frame can be a *read* dependency for several write_idx tasks;
+        # writing in place risked one thread reading a file while another
+        # was mid-save on it (truncated/corrupt PNG -> shape errors).
+        # Reading on demand (no full-clip RAM cache) avoids the OOM/disk
+        # blowup from caching every frame's alpha channel at once.
+        smoothed_dir = processed_dir + "_smoothed"
+        if not os.path.isdir(smoothed_dir):
+            os.mkdir(smoothed_dir)
+
+        def get_alpha(idx):
+            file = files[idx]
+            img = Image.open(os.path.join(processed_dir, file)).convert("RGBA")
+            return np.array(img)[:, :, 3].astype(np.float32)
+
+        def smooth_frame(write_idx):
+            try:
+                start = max(0, write_idx - half)
+                end = min(total - 1, write_idx + half)
+                alphas = np.stack([get_alpha(j) for j in range(start, end + 1)])
+                smoothed_alpha = np.mean(alphas, axis=0).astype(np.uint8)
+
+                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(smoothed_dir, filename))
+
+                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(n_workers)
+        ]
+        for t in smoothing_threads:
+            t.start()
+        for t in smoothing_threads:
+            t.join()
+
+        if smoothing_errors:
+            rmtree(smoothed_dir, ignore_errors=True)
+            raise smoothing_errors[0]
+
+        # Swap the smoothed frames in as the new processed_dir contents.
+        rmtree(processed_dir)
+        os.rename(smoothed_dir, processed_dir)
+
+    # 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.bmp"),
+        r=framerate,
+        f="image2",
+        s=whstr,
+        pix_fmt="yuva444p10le",
+    )
+    stream = ffmpeg.output(
+        stream, args.o, vcodec="prores_ks", **{"profile:v": "4", "bits_per_mb": "5000"}
+    )
+    ffmpeg.run(stream)
+
+except KeyboardInterrupt:
+    print("\nInterrupted by user")
+
+finally:
+    print("Removing temporary files...")
+    rmtree(processed_dir, ignore_errors=True)
+    rmtree(frames_dir, ignore_errors=True)

working_temp.py

@@ -0,0 +1,311 @@
+import argparse
+import io
+import os
+import ffmpeg
+import pathlib
+import threading
+import numpy as np
+from queue import Queue
+from shutil import rmtree
+from PIL import Image
+from rembg import new_session, remove
+
+os.environ.setdefault(
+    "PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True,managed_memory:True"
+)
+
+# Parse args
+parser = argparse.ArgumentParser(
+    description="Applies rembg background removal to the frames of a video"
+)
+parser.add_argument("input", type=str, help="Input video")
+parser.add_argument(
+    "-o", type=str, default="export/output.mov", help="Define output path"
+)
+parser.add_argument(
+    "--model",
+    type=str,
+    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)",
+)
+parser.add_argument(
+    "--smooth",
+    type=int,
+    default=3,
+    help="Temporal mask smoothing window size in frames (default: 3, 0 to disable)",
+)
+parser.add_argument(
+    "--read-ahead",
+    type=int,
+    default=8,
+    help="Number of frames to read ahead into buffer (default: 8)",
+)
+parser.add_argument(
+    "--write-buffer",
+    type=int,
+    default=8,
+    help="Number of processed frames to buffer before writing (default: 8)",
+)
+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)
+framerate = video_stream["avg_frame_rate"]
+
+frames_dir = os.path.join(str(pathlib.Path(__file__).parent.absolute()), "frames")
+processed_dir = os.path.join(str(pathlib.Path(__file__).parent.absolute()), "processed")
+
+# Extract input video frames
+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.bmp"))
+ffmpeg.run(stream)
+
+_SENTINEL = object()
+
+try:
+    # Process frames with pipelined reader -> processors -> writer
+    if not os.path.isdir(processed_dir):
+        os.mkdir(processed_dir)
+
+    files = sorted(os.listdir(frames_dir))
+    total_files = len(files)
+
+    print(f"Loading rembg session (model={args.model})...", flush=True)
+    session = new_session(
+        args.model, providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
+    )
+
+    read_queue = Queue(maxsize=args.read_ahead)
+    write_queue = Queue(maxsize=args.write_buffer)
+    errors = []
+    active_processors = [
+        args.workers
+    ]  # list so processor() can mutate without nonlocal
+    active_processors_lock = threading.Lock()
+
+    def reader():
+        try:
+            for idx, file in enumerate(files, 1):
+                with open(os.path.join(frames_dir, file), "rb") as f:
+                    data = f.read()
+                read_queue.put((idx, file, data))
+        except Exception as e:
+            errors.append(e)
+        finally:
+            # One sentinel per worker so each one knows when to stop
+            for _ in range(args.workers):
+                read_queue.put(_SENTINEL)
+
+    def processor():
+        try:
+            while True:
+                item = read_queue.get()
+                if item is _SENTINEL:
+                    break
+                idx, file, input_data = item
+                print(f"Processing frame {idx}/{total_files}: {file}", flush=True)
+                output_data = remove_with_mask_fallback(input_data, session=session)
+                write_queue.put((idx, file, output_data))
+        except Exception as e:
+            errors.append(e)
+        finally:
+            # Signal writer only when the last processor finishes
+            with active_processors_lock:
+                active_processors[0] -= 1
+                if active_processors[0] == 0:
+                    write_queue.put(_SENTINEL)
+
+    def writer():
+        try:
+            while True:
+                item = write_queue.get()
+                if item is _SENTINEL:
+                    break
+                idx, file, output_data = item
+                with open(os.path.join(processed_dir, file), "wb") as f:
+                    f.write(output_data)
+                print(f"Written frame {idx}/{total_files}: {file}", flush=True)
+        except Exception as e:
+            errors.append(e)
+
+    reader_thread = threading.Thread(target=reader, daemon=True)
+    processor_threads = [
+        threading.Thread(target=processor, daemon=True) for _ in range(args.workers)
+    ]
+    writer_thread = threading.Thread(target=writer, daemon=True)
+
+    reader_thread.start()
+    for t in processor_threads:
+        t.start()
+    writer_thread.start()
+
+    reader_thread.join()
+    for t in processor_threads:
+        t.join()
+    writer_thread.join()
+
+    if errors:
+        raise errors[0]
+
+    # Temporal mask smoothing
+    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)
+
+        alpha_buf = {}
+
+        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")
+                alpha_buf[read_idx] = (file, np.array(img)[:, :, 3].astype(np.float32))
+
+            write_idx = read_idx - half
+            if 0 <= write_idx < total:
+                start = max(0, write_idx - half)
+                end = min(total - 1, write_idx + half)
+                alphas = np.stack([alpha_buf[j][1] for j in range(start, end + 1)])
+                smoothed_alpha = np.mean(alphas, axis=0).astype(np.uint8)
+                filename = alpha_buf[write_idx][0]
+                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 alpha_buf:
+                    del alpha_buf[drop_idx]
+
+    # 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.bmp"),
+        r=framerate,
+        f="image2",
+        s=whstr,
+        pix_fmt="yuva444p10le",
+    )
+    stream = ffmpeg.output(
+        stream, args.o, vcodec="prores_ks", **{"profile:v": "4", "bits_per_mb": "5000"}
+    )
+    ffmpeg.run(stream)
+
+except KeyboardInterrupt:
+    print("\nInterrupted by user")
+
+finally:
+    print("Removing temporary files...")
+    rmtree(processed_dir, ignore_errors=True)
+    rmtree(frames_dir, ignore_errors=True)