Automatic cleanup

2026-05-29 10:10:28 +03:00
parent ceb3db92e6
commit a9d5f7b97c
7 changed files with 125 additions and 263 deletions
@@ -1,3 +1,5 @@
-/build/**
+.venv
-/lib/yt-dlp
+/processed
-/testfootage/
+/frames
 /export
 *.mp4
@@ -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.
@@ -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)
@@ -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)
@@ -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()