LyAhn
fa0170daa9
C++ WebView2 host app driving a Python/rembg engine as a JSON-stdio sidecar, with live before/after preview for images and video frames. Also includes the original CLI (removebg.py) built on the same core/ engine. The WebView2 SDK is fetched on demand via native/third_party/fetch-webview2.ps1 (wired into the Makefile) rather than vendored, since it's ~15MB of prebuilt/generated Microsoft SDK content. A Makefile provides make sync/build/run/clean as the standard entry points.
255 lines
8.9 KiB
Python
255 lines
8.9 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
RemoveBG Tool - Command Line & Graphical Interface for Image & Video Background Removal
|
|
Powered by rembg and FFmpeg.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import argparse
|
|
from typing import List
|
|
|
|
# Add parent directory to path if running directly
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
from core.processor import BackgroundRemoverEngine, AVAILABLE_MODELS
|
|
from core.ffmpeg_utils import check_ffmpeg_installed
|
|
|
|
def run_cli(args):
|
|
"""Execute command-line processing mode."""
|
|
print("=" * 60)
|
|
print("✨ RemoveBG CLI Processor")
|
|
print("=" * 60)
|
|
|
|
model_name = args.model
|
|
if model_name not in AVAILABLE_MODELS:
|
|
print(f"Warning: Unknown model '{model_name}'. Available options: {list(AVAILABLE_MODELS.keys())}")
|
|
model_name = "u2net"
|
|
|
|
print(f"Loading Model: {model_name} ({AVAILABLE_MODELS.get(model_name, '')})")
|
|
engine = BackgroundRemoverEngine(model_name=model_name)
|
|
|
|
bg_color = args.bg_color
|
|
if bg_color:
|
|
if bg_color.lower() == "green":
|
|
bg_color = "#00FF00"
|
|
elif bg_color.lower() == "white":
|
|
bg_color = "#FFFFFF"
|
|
elif bg_color.lower() == "black":
|
|
bg_color = "#000000"
|
|
|
|
# 1. Video Processing Mode
|
|
if args.video:
|
|
if not check_ffmpeg_installed():
|
|
print("❌ Error: FFmpeg is not installed or not found in system PATH. Video processing requires FFmpeg.")
|
|
sys.exit(1)
|
|
|
|
video_path = args.video
|
|
output_path = args.output
|
|
if not output_path:
|
|
dir_name, base_name = os.path.split(video_path)
|
|
fname, _ = os.path.splitext(base_name)
|
|
output_path = os.path.join(dir_name, f"{fname}_nobg.mov")
|
|
|
|
print(f"\n🎬 Input Video: {video_path}")
|
|
print(f"🎯 Output Destination: {output_path}")
|
|
print(f"⚙️ Export Format: {args.export_format}")
|
|
if bg_color:
|
|
print(f"🎨 Background Color: {bg_color}")
|
|
print(f"✂️ Alpha Matting: {'Enabled' if args.alpha_matting else 'Disabled'}")
|
|
|
|
from tqdm import tqdm
|
|
pbar = None
|
|
|
|
def pcb(current, total_count, msg):
|
|
nonlocal pbar
|
|
if total_count > 0:
|
|
if pbar is None or pbar.total != total_count:
|
|
if pbar:
|
|
pbar.close()
|
|
pbar = tqdm(total=total_count, desc="Processing Video Frames", unit="frame")
|
|
pbar.n = current
|
|
pbar.refresh()
|
|
else:
|
|
print(f"[Status] {msg}")
|
|
|
|
try:
|
|
result = engine.process_video(
|
|
video_path=video_path,
|
|
output_path=output_path,
|
|
fps=args.fps if args.fps > 0 else None,
|
|
export_format=args.export_format,
|
|
bg_color=bg_color,
|
|
alpha_matting=args.alpha_matting,
|
|
progress_callback=pcb
|
|
)
|
|
if pbar:
|
|
pbar.close()
|
|
print(f"\n✅ Video background removal complete! File saved to: {result}")
|
|
except Exception as e:
|
|
if pbar:
|
|
pbar.close()
|
|
print(f"\n❌ Error processing video: {e}")
|
|
sys.exit(1)
|
|
|
|
# 2. Batch Directory Mode
|
|
elif args.batch or (args.image and os.path.isdir(args.image[0])):
|
|
batch_dir = args.batch or args.image[0]
|
|
output_dir = args.output or os.path.join(batch_dir, "output_nobg")
|
|
|
|
image_files = []
|
|
for root, _, files in os.walk(batch_dir):
|
|
for f in files:
|
|
if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff")):
|
|
image_files.append(os.path.join(root, f))
|
|
|
|
if not image_files:
|
|
print(f"❌ No valid image files found in directory: {batch_dir}")
|
|
sys.exit(1)
|
|
|
|
print(f"\n📁 Batch Input Directory: {batch_dir} ({len(image_files)} images found)")
|
|
print(f"🎯 Output Directory: {output_dir}")
|
|
|
|
from tqdm import tqdm
|
|
pbar = tqdm(total=len(image_files), desc="Removing Backgrounds", unit="img")
|
|
|
|
def pcb(current, total_count, fname):
|
|
pbar.n = current
|
|
pbar.set_postfix_str(fname)
|
|
pbar.refresh()
|
|
|
|
results = engine.process_batch_images(
|
|
input_paths=image_files,
|
|
output_dir=output_dir,
|
|
alpha_matting=args.alpha_matting,
|
|
bg_color=bg_color,
|
|
output_format=args.format,
|
|
progress_callback=pcb
|
|
)
|
|
pbar.close()
|
|
print(f"\n✅ Batch processing complete! Saved {len(results)} images to: {output_dir}")
|
|
|
|
# 3. Single / Multiple Specific Images
|
|
elif args.image:
|
|
input_paths = args.image
|
|
if len(input_paths) == 1:
|
|
in_file = input_paths[0]
|
|
out_file = args.output
|
|
if not out_file:
|
|
dir_name, base_name = os.path.split(in_file)
|
|
fname, _ = os.path.splitext(base_name)
|
|
out_file = os.path.join(dir_name, f"{fname}_nobg.{args.format.lower()}")
|
|
|
|
print(f"\n🖼️ Input Image: {in_file}")
|
|
print(f"🎯 Output File: {out_file}")
|
|
|
|
try:
|
|
res = engine.process_single_image(
|
|
in_file,
|
|
out_file,
|
|
alpha_matting=args.alpha_matting,
|
|
bg_color=bg_color
|
|
)
|
|
print(f"✅ Success! Image saved to: {res}")
|
|
except Exception as e:
|
|
print(f"❌ Error: {e}")
|
|
sys.exit(1)
|
|
else:
|
|
out_dir = args.output or os.path.dirname(input_paths[0])
|
|
print(f"\n🖼️ Processing {len(input_paths)} images...")
|
|
from tqdm import tqdm
|
|
pbar = tqdm(total=len(input_paths), desc="Removing Backgrounds", unit="img")
|
|
|
|
def pcb(current, total_count, fname):
|
|
pbar.n = current
|
|
pbar.set_postfix_str(fname)
|
|
pbar.refresh()
|
|
|
|
results = engine.process_batch_images(
|
|
input_paths=input_paths,
|
|
output_dir=out_dir,
|
|
alpha_matting=args.alpha_matting,
|
|
bg_color=bg_color,
|
|
output_format=args.format,
|
|
progress_callback=pcb
|
|
)
|
|
pbar.close()
|
|
print(f"✅ Success! Processed {len(results)} images.")
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="RemoveBG Studio - AI Background Removal for Images & Video Files (FFmpeg)",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
Examples:
|
|
# For a graphical UI, build and run native/RemoveBGStudio.exe instead (see README.md).
|
|
|
|
# Single Image Background Removal:
|
|
python removebg.py -i input.jpg -o output.png
|
|
|
|
# Batch Image Directory Background Removal:
|
|
python removebg.py -b ./my_images -o ./no_bg_images --model birefnet-general
|
|
|
|
# Video Background Removal with FFmpeg:
|
|
python removebg.py -v input.mp4 -o transparent.mov --export-format mov_transparent
|
|
|
|
# Video Background Removal with Green Screen Replacement:
|
|
python removebg.py -v input.mp4 -o greenscreen.mp4 --export-format mp4_solid --bg-color green
|
|
"""
|
|
)
|
|
|
|
parser.add_argument("-i", "--image", nargs="+", help="Input image file path(s)")
|
|
parser.add_argument("-b", "--batch", help="Input directory containing images for batch processing")
|
|
parser.add_argument("-v", "--video", help="Input video file path to extract and process frames via FFmpeg")
|
|
parser.add_argument("-o", "--output", help="Output file path (single image/video) or output directory (batch)")
|
|
|
|
parser.add_argument(
|
|
"-m", "--model",
|
|
default="u2net",
|
|
choices=list(AVAILABLE_MODELS.keys()),
|
|
help="AI Model to use for background removal (default: u2net)"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"-bg", "--bg-color",
|
|
help="Background replacement color (e.g. 'green', 'white', 'black', or Hex '#00FF00'). Default: transparent"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"-am", "--alpha-matting",
|
|
action="store_true",
|
|
help="Enable alpha matting for finer edge detail (hair, fur)"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--format",
|
|
default="png",
|
|
choices=["png", "webp", "jpg"],
|
|
help="Output image format for single/batch processing (default: png)"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--fps",
|
|
type=float,
|
|
default=0.0,
|
|
help="Target frame rate for video frame extraction (0 = original video FPS)"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--export-format",
|
|
default="mov_transparent",
|
|
choices=["mov_transparent", "webm_transparent", "mp4_solid", "png_sequence"],
|
|
help="Export format for video reassembly (default: mov_transparent)"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
if not args.image and not args.batch and not args.video:
|
|
parser.print_help()
|
|
sys.exit(1)
|
|
|
|
run_cli(args)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|