Merge pull request 'v0.2.0' (#6) from rembg into main

Reviewed-on: #6

.gitignore

@@ -3,3 +3,7 @@
 /frames
 /export
 *.mp4
+*.webm
+*.mov
+*.mkv
+__pycache__

.gitmodules

@@ -1,6 +0,0 @@
-[submodule "lib/sys"]
-	path = lib/sys
-	url = https://go.googlesource.com/sys
-[submodule "lib/fsnotify"]
-	path = lib/fsnotify
-	url = https://github.com/fsnotify/fsnotify.git

LICENSE

@@ -1,6 +1,7 @@
 MIT License
 
 Copyright (c) 2021 Seth Tribbey
+Copyright (c) 2026 Uula Ilola
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

pyrightconfig.json

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

readme.md

@@ -1,32 +1,64 @@
-# 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:
-[Install rembg by following their instructions](https://github.com/danielgatis/rembg)
+## Installation:
 
-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.
-
-Then:
+Currently i've gotten rembg to work only with Python **3.12** because of onnxruntime's shenanigans X/
+```bash
+python3.12 -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
+> **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
 
-positional arguments:
-  input           Input video
+&nbsp;
 
-optional arguments:
-  -h, --help      show this help message and exit
-  -a              Turns on alpha matting during background removal
-  -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)
+## Usage:
+```
+python .\rembg_video.py [-h] [--help] [-o] [--model] [--workers] [--smooth] [--smooth-workers] [--buffer-size] [--output-type] 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>
+    <tr><td>--workers</td><td>Number of concurrent process workers</td></tr>
+    <tr><td>--smooth</td><td>Size of window for temporal smoothing</td></tr>
+    <tr><td>--smooth-workers</td><td>Number of CPU threads for temporal smoothing</td></tr>
+    <tr><td>--buffer-size</td><td>Set size of processing buffer</td></tr>
+    <tr><td>--output-type</td><td>Choose how video is exported 
+    <br>"complete" = Full Color .mov <br>"mask" = Only mask <br>"mask_seq" = Mask image sequence)
+  </table>
+
+  <table>
+    <tr><th colspan="2">Positional Arguments:</th></tr>
+    <tr><td>input</td><td>Input Video</td></tr>
+  </table>
+</div>

rembg_video.py

@@ -1,67 +1,346 @@
 import argparse
+import io
 import os
 import ffmpeg
 import pathlib
-import shutil
-from rembg.bg import remove
+import threading
+import numpy as np
+from queue import Queue
+from shutil import rmtree, move
+from PIL import Image
+from rembg import new_session, remove
 
-#Parse args
-parser = argparse.ArgumentParser(description='Applies rembg to the frames of a video')
-parser.add_argument('input', type=str, help='Input video')
-parser.add_argument('-o', type=str, default="export/output.mp4", help="Define output path")
-parser.add_argument('-a', action="store_true", help="Turns on alpha matting during background removal")
-parser.add_argument('-af', type=int, default=240, help="Alpha matting foreground threshold")
-parser.add_argument('-ab', type=int, default=10, help="Alpha matting background threshold")
-parser.add_argument('-ae', type=int, default=10, help="Alpha matting erode size")
-parser.add_argument('--skip-extract', action="store_true", help='Skips ffmpeg frame extraction')
-parser.add_argument('--skip-process', action="store_true", help='Skips rembg frame processing')
+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", help="Define output path")
+parser.add_argument(
+    "--model",
+    type=str,
+    default="u2net_human_seg",
+    help="rembg model to use (default: u2net_human_seg)",
+)
+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(
+    "--buffer-size",
+    type=int,
+    default=8,
+    help="Set maximum number of frames in buffer (default: 8)",
+)
+parser.add_argument(
+    "--smooth-workers",
+    type=int,
+    default=os.cpu_count() or 4,
+    help="Number of cpu threads to use for temporal mask smoothing (default: cpu count)",
+)
+parser.add_argument(
+    "--output-type",
+    type=str,
+    choices=["complete", "mask", "mask_seq"],
+    default="complete",
+    help="What way to output keyed video. (default: complete)",
+)
 args = parser.parse_args()
 
-#Extract video info
+
+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_bytes(image):
+    with io.BytesIO() as buffer:
+        image.save(buffer, format="TIFF")
+        return buffer.getvalue()
+
+
+def remove_with_fallback(image_bytes, session):
+    try:
+        rembg_out = remove(image_bytes, session=session)
+        if args.output_type == "complete":
+            return rembg_out
+        elif args.output_type in ("mask", "mask_seq"):
+            if isinstance(rembg_out, (bytes, bytearray)):
+                mask_bytes = bytes(rembg_out)
+            else:
+                raise TypeError('mask_bytes is not of type "bytes"')
+            alpha = Image.open(io.BytesIO(mask_bytes)).convert("RGBA").getchannel("A")
+            return image_to_bytes(alpha)
+        else:
+            raise Exception("Unknown output type.")
+
+    except Exception as exc:
+        if is_oom_error(exc):
+            raise RuntimeError("Insufficient GPU memory!") from exc
+        else:
+            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)
-width = int(video_stream['width'])
-height = int(video_stream['height'])
-whstr = str(width) + 'x' + str(height)
-framerate = video_stream['avg_frame_rate']
+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"]
 
-#Extract input video frames
-if not args.skip_extract:
-  frames_dir = os.path.join(str(pathlib.Path(__file__).parent.absolute()), "frames")
-  if not os.path.isdir(frames_dir):
-    os.mkdir(frames_dir)
+# Setup workspace folders
+frames_dir = os.path.join(str(pathlib.Path(__file__).parent.absolute()), "frames")
+processed_dir = os.path.join(str(pathlib.Path(__file__).parent.absolute()), "processed")
+smoothed_dir = processed_dir + "_smoothed"
+rmtree(frames_dir, ignore_errors=True)
+rmtree(processed_dir, ignore_errors=True)
+rmtree(smoothed_dir, ignore_errors=True)
 
-  stream = ffmpeg.input(args.input)
-  stream = ffmpeg.output(stream, os.path.join(frames_dir, "%04d.png"))
-  ffmpeg.run(stream)
-
-#Process frames with rembg
-if not args.skip_process:
-  files_dir = os.path.join(str(pathlib.Path(__file__).parent.absolute()), "frames")
-  processed_dir = os.path.join(str(pathlib.Path(__file__).parent.absolute()), "processed")
-  if not os.path.isdir(processed_dir):
-    os.mkdir(processed_dir)
-
-  files = sorted(os.listdir(files_dir))
-  total_files = len(files)
-  for idx, file in enumerate(files, 1):
-    print(f"Processing frame {idx}/{total_files}: {file}", flush=True)
-    with open(os.path.join(files_dir, file), "rb") as i:
-      with open(os.path.join(processed_dir, file), "wb") as o:
-          input = i.read()
-          output = remove(input, alpha_matting=args.a, alpha_matting_foreground_threshold=args.af, alpha_matting_background_threshold=args.ab, alpha_matting_erode_size=args.ae)
-          o.write(output)
-    print(f"Completed frame {idx}/{total_files}", flush=True)
-
-#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"), r=framerate, f='image2', s=whstr, pix_fmt='yuv420p')
-stream = ffmpeg.output(stream, args.o, vcodec='libx264', crf=25)
+# Extract input video frames
+os.mkdir(frames_dir)
+stream = ffmpeg.input(args.input)
+stream = ffmpeg.output(stream, os.path.join(frames_dir, "%04d.tiff"))
 ffmpeg.run(stream)
 
-#Cleanup
-print("Removing temporary files...")
-shutil.rmtree(processed_dir)
-shutil.rmtree(frames_dir)
+_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.buffer_size)
+    write_queue = Queue(maxsize=args.buffer_size)
+    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_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 img -> shape errors).
+        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)
+
+    if args.output_type != "mask_seq":
+        # Output video
+        stream = ffmpeg.input(
+            os.path.join(processed_dir, "%04d.tiff"),
+            r=framerate,
+            f="image2",
+            s=whstr,
+        )
+        if args.output_type == "mask":
+            output_file = pathlib.Path(args.o) / ("output.mp4")
+            output_file.parent.mkdir(exist_ok=True, parents=True)
+            stream = ffmpeg.output(
+                stream,
+                str(output_file),
+                vcodec="libx264",
+                pix_fmt="gray",
+                crf=0,
+                preset="veryslow",
+                tune="animation",
+            )
+        else:
+            output_file = pathlib.Path(args.o) / ("output.mov")
+            output_file.parent.mkdir(exist_ok=True, parents=True)
+            stream = ffmpeg.output(
+                stream, str(output_file), vcodec="prores_ks", **{"profile:v": "4"}
+            )
+
+        ffmpeg.run(stream, overwrite_output=True)
+
+    else:
+        img_seq_out_folder = os.path.join(args.o, "output_img_seq")
+        rmtree(img_seq_out_folder, ignore_errors=True)
+        move(processed_dir, img_seq_out_folder)
+
+except KeyboardInterrupt:
+    print("\nInterrupted by user")
+
+finally:
+    print("Removing temporary files...")
+    rmtree(processed_dir, ignore_errors=True)
+    rmtree(frames_dir, ignore_errors=True)
+    rmtree(smoothed_dir, ignore_errors=True)

requirements(3.12).txt

@@ -0,0 +1,42 @@
+attrs==26.1.0
+certifi==2026.6.17
+charset-normalizer==3.4.7
+coloredlogs==15.0.1
+ffmpeg-python==0.2.0
+flatbuffers==25.12.19
+future==1.0.0
+humanfriendly==10.0
+idna==3.18
+ImageIO==2.37.3
+jsonschema==4.26.0
+jsonschema-specifications==2025.9.1
+lazy-loader==0.5
+llvmlite==0.47.0
+mpmath==1.3.0
+networkx==3.6.1
+numba==0.65.1
+numpy==2.4.6
+nvidia-cublas-cu12==12.9.2.10
+nvidia-cuda-nvrtc-cu12==12.9.86
+nvidia-cuda-runtime-cu12==12.9.79
+nvidia-cudnn-cu12==9.10.2.21
+nvidia-cufft-cu12==11.4.1.4
+nvidia-nvjitlink-cu12==12.9.86
+onnxruntime-gpu==1.19.2
+packaging==26.2
+pillow==12.2.0
+platformdirs==4.10.0
+pooch==1.9.0
+protobuf==7.35.1
+PyMatting==1.1.15
+referencing==0.37.0
+rembg==2.0.76
+requests==2.34.2
+rpds-py==2026.5.1
+scikit-image==0.26.0
+scipy==1.18.0
+sympy==1.14.0
+tifffile==2026.6.1
+tqdm==4.68.3
+typing_extensions==4.15.0
+urllib3==2.7.0