de4dn0te/kakisalmi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Automatic cleanup

Browse Source
This commit is contained in:
de4dn0te
2026-05-29 10:10:28 +03:00
parent ceb3db92e6
commit a9d5f7b97c
7 changed files with 125 additions and 263 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+5 -3
View File
@@ -1,3 +1,5 @@
/build/**
/lib/yt-dlp
/testfootage/
.venv
/processed
/frames
/export
*.mp4
LICENSE
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2021 Seth Tribbey
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
lib/fonts/DejaVuSans-Bold.ttf
BIN
View File
Binary file not shown.
lib/fonts/DejaVuSans.ttf
BIN
View File
Binary file not shown.
readme.md
+32 -1
View File
@@ -1 +1,32 @@
# Automatic Background Removal
# rembg_from_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).
### 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.
Then:
```
pip install ffmpeg-python
```
### Usage:
```
python .\rembg_video.py [-h] [-a] [-af AF] [-ab AB] [-ae AE] [--skip-extract] [--skip-process] input
positional arguments:
input Input video
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)
rembg_video.py
+67
View File
@@ -0,0 +1,67 @@
import argparse
import os
import ffmpeg
import pathlib
import shutil
from rembg.bg import 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')
args = parser.parse_args()
#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']
#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)
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)
ffmpeg.run(stream)
#Cleanup
print("Removing temporary files...")
shutil.rmtree(processed_dir)
shutil.rmtree(frames_dir)
src/booth_rembg.py
-259
View File
@@ -1,259 +0,0 @@
#!/usr/bin/env python3
"""
Video background removal using rembg.
Outputs a transparent .webm/.mov, or composites onto a solid/image background.
Usage:
# Transparent output (WebM with alpha)
python remove_bg.py input.mp4 output.webm
# Composite onto a colour
python remove_bg.py input.mp4 output.mp4 --bg-color 0,255,0
# Composite onto an image
python remove_bg.py input.mp4 output.mp4 --bg-image background.jpg
# Use a specific rembg model (default: u2net)
python remove_bg.py input.mp4 output.webm --model birefnet-general
Options:
--model rembg model name (see MODEL NOTES below)
--bg-color R,G,B background colour (0-255)
--bg-image Path to a background image/video frame
--fps Override output FPS (default: match source)
--start Start time in seconds (default: 0)
--end End time in seconds (default: end of video)
--workers Parallel worker threads (default: 2)
--no-gpu Disable GPU/ONNX GPU provider
MODEL NOTES:
u2net Default. Good general-purpose, fast.
u2net_human_seg Tuned for people — better than u2net for humans.
birefnet-general Best quality overall. Slower, higher VRAM.
birefnet-portrait Best for portrait/bust shots of people.
isnet-general-use Strong edges, good alternative to birefnet.
silueta Lightweight, fast, less accurate.
"""
import argparse
import os
import sys
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
import cv2
import numpy as np
from rembg import new_session, remove
from tqdm import tqdm
# --- Helpers ---
def open_video(path: str) -> cv2.VideoCapture:
cap = cv2.VideoCapture(path)
if not cap.isOpened():
sys.exit(f"[error] Cannot open video: {path}")
return cap
def video_meta(cap: cv2.VideoCapture) -> dict:
return {
"fps": cap.get(cv2.CAP_PROP_FPS),
"width": int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
"height": int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
"total": int(cap.get(cv2.CAP_PROP_FRAME_COUNT)),
}
def make_writer(path: str, fps: float, width: int, height: int, alpha: bool) -> cv2.VideoWriter:
ext = Path(path).suffix.lower()
if alpha:
if ext == ".webm":
fourcc = cv2.VideoWriter_fourcc(*"VP90")
elif ext in (".mov", ".avi"):
fourcc = cv2.VideoWriter_fourcc(*"png ") # PNG codec for lossless alpha
else:
print(f"[warn] Alpha channel requested but '{ext}' may not support it. "
"Use .webm or .mov for transparency.")
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
return cv2.VideoWriter(path, fourcc, fps, (width, height), isColor=True)
else:
if ext == ".mp4":
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
elif ext == ".avi":
fourcc = cv2.VideoWriter_fourcc(*"XVID")
elif ext == ".webm":
fourcc = cv2.VideoWriter_fourcc(*"VP90")
else:
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
return cv2.VideoWriter(path, fourcc, fps, (width, height))
def load_bg_image(path: str, width: int, height: int) -> np.ndarray:
img = cv2.imread(path)
if img is None:
sys.exit(f"[error] Cannot read background image: {path}")
return cv2.resize(img, (width, height))
def composite_on_color(rgba: np.ndarray, bg_color: tuple[int, int, int]) -> np.ndarray:
"""Blend RGBA frame onto a solid colour. Returns BGR."""
alpha = rgba[:, :, 3:4].astype(np.float32) / 255.0
fg = rgba[:, :, :3].astype(np.float32)
bg = np.full_like(fg, bg_color[::-1], dtype=np.float32) # RGB→BGR
out = (fg * alpha + bg * (1.0 - alpha)).astype(np.uint8)
return out
def composite_on_image(rgba: np.ndarray, bg: np.ndarray) -> np.ndarray:
"""Blend RGBA frame onto a BGR background image. Returns BGR."""
alpha = rgba[:, :, 3:4].astype(np.float32) / 255.0
fg = rgba[:, :, :3].astype(np.float32)
bg_f = bg.astype(np.float32)
out = (fg * alpha + bg_f * (1.0 - alpha)).astype(np.uint8)
return out
# ---------------------------------------------------------------------------
# Per-frame processing
# ---------------------------------------------------------------------------
_session_local = threading.local()
def process_frame(
frame_bgr: np.ndarray,
model_name: str,
providers: list[str],
) -> np.ndarray:
"""Remove background from a single BGR frame. Returns RGBA numpy array."""
# Each thread gets its own rembg session (not thread-safe to share)
if not hasattr(_session_local, "session"):
_session_local.session = new_session(model_name, providers=providers)
# rembg expects PIL or bytes; convert BGR→RGB bytes via PNG
rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
success, buf = cv2.imencode(".png", rgb)
if not success:
raise RuntimeError("Failed to encode frame as PNG")
result_bytes = remove(buf.tobytes(), session=_session_local.session)
result_arr = np.frombuffer(result_bytes, dtype=np.uint8)
rgba = cv2.imdecode(result_arr, cv2.IMREAD_UNCHANGED) # RGBA
if rgba is None or rgba.shape[2] != 4:
raise RuntimeError("rembg did not return an RGBA image")
return rgba
# ---------------------------------------------------------------------------
# Main pipeline
# ---------------------------------------------------------------------------
def run(args: argparse.Namespace) -> None:
cap = open_video(args.input)
meta = video_meta(cap)
fps = args.fps or meta["fps"]
W, H = meta["width"], meta["height"]
total = meta["total"]
# Seek to start frame
start_frame = int((args.start or 0) * meta["fps"])
end_frame = int(args.end * meta["fps"]) if args.end else total
if start_frame > 0:
cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
n_frames = end_frame - start_frame
# Determine output mode
ext = Path(args.output).suffix.lower()
alpha_mode = (args.bg_color is None and args.bg_image is None)
# ONNX providers
providers = ["CPUExecutionProvider"] if args.no_gpu else \
["CUDAExecutionProvider", "CPUExecutionProvider"]
# Background image (loaded once)
bg_img = None
if args.bg_image:
bg_img = load_bg_image(args.bg_image, W, H)
# Output writer
writer = make_writer(args.output, fps, W, H, alpha=alpha_mode)
print(f"[info] Input : {args.input} ({W}×{H} @ {meta['fps']:.2f} fps, {total} frames)")
print(f"[info] Output : {args.output} ({'transparent' if alpha_mode else 'composited'})")
print(f"[info] Model : {args.model}")
print(f"[info] Frames : {start_frame}–{end_frame} ({n_frames} frames)")
print(f"[info] Workers: {args.workers}")
# Read all frames into memory (batched for thread safety)
# For very long videos you may want to chunk this
print("[info] Reading frames...")
frames = []
for _ in range(n_frames):
ok, frame = cap.read()
if not ok:
break
frames.append(frame)
cap.release()
print(f"[info] Processing {len(frames)} frames with rembg ({args.model})...")
# Process frames in parallel
results = [None] * len(frames)
with ThreadPoolExecutor(max_workers=args.workers) as pool:
future_to_idx = {
pool.submit(process_frame, f, args.model, providers): i
for i, f in enumerate(frames)
}
with tqdm(total=len(frames), unit="frame") as pbar:
for future in as_completed(future_to_idx):
idx = future_to_idx[future]
try:
rgba = future.result()
results[idx] = rgba
except Exception as e:
print(f"\n[warn] Frame {idx} failed: {e} — using blank frame")
results[idx] = np.zeros((H, W, 4), dtype=np.uint8)
pbar.update(1)
# Write output in order
print("[info] Writing output video...")
for rgba in tqdm(results, unit="frame"):
if alpha_mode:
# Write RGBA as BGRA
bgra = cv2.cvtColor(rgba, cv2.COLOR_RGBA2BGRA)
writer.write(bgra)
elif bg_img is not None:
bgr = composite_on_image(rgba, bg_img)
writer.write(bgr)
else:
bgr = composite_on_color(rgba, args.bg_color)
writer.write(bgr)
writer.release()
size_mb = os.path.getsize(args.output) / 1024 / 1024
print(f"[done] Saved → {args.output} ({size_mb:.1f} MB)")
def main():
p = argparse.ArgumentParser(
description="Remove background from a video using rembg.",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog=__doc__,
)
p.add_argument("input", help="Input video path")
p.add_argument("output", help="Output video path (.webm/.mov for alpha, .mp4 for composite)")
args = p.parse_args()
run(args)
if __name__ == "__main__":
main()