feat(rembg): add RemoveBG Studio - native background removal app
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.
This commit is contained in:
@@ -0,0 +1,254 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user