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,223 @@
|
||||
import os
|
||||
import tempfile
|
||||
import shutil
|
||||
from typing import List, Optional, Callable, Dict, Tuple
|
||||
from PIL import Image, ImageColor
|
||||
import rembg
|
||||
from core import ffmpeg_utils
|
||||
|
||||
# Available model descriptions
|
||||
AVAILABLE_MODELS = {
|
||||
"u2net": "U2-Net (Default - General Purpose)",
|
||||
"u2netp": "U2-Net Lite (Fast, Lightweight)",
|
||||
"u2net_human_seg": "U2-Net Human Seg (Optimized for People & Portraits)",
|
||||
"u2net_cloth_seg": "U2-Net Cloth Seg (Optimized for Clothes/Fashion)",
|
||||
"isnet-general-use": "IS-Net General Use (High Accuracy & Crisp Edges)",
|
||||
"birefnet-general": "BiRefNet General (State-of-the-Art High Precision)"
|
||||
}
|
||||
|
||||
class BackgroundRemoverEngine:
|
||||
def __init__(self, model_name: str = "u2net"):
|
||||
self.model_name = model_name
|
||||
self.session = None
|
||||
self._init_session(model_name)
|
||||
|
||||
def _init_session(self, model_name: str):
|
||||
self.model_name = model_name
|
||||
try:
|
||||
self.session = rembg.new_session(model_name)
|
||||
except Exception as e:
|
||||
# Fallback to standard u2net if specific model fails
|
||||
print(f"Warning: Failed to load model {model_name}: {e}. Falling back to u2net.")
|
||||
self.model_name = "u2net"
|
||||
self.session = rembg.new_session("u2net")
|
||||
|
||||
def change_model(self, model_name: str):
|
||||
if model_name != self.model_name:
|
||||
self._init_session(model_name)
|
||||
|
||||
def remove_background(
|
||||
self,
|
||||
image_input: Image.Image,
|
||||
alpha_matting: bool = False,
|
||||
af: int = 240,
|
||||
ab: int = 10,
|
||||
ae: int = 10,
|
||||
bg_color: Optional[str] = None
|
||||
) -> Image.Image:
|
||||
"""
|
||||
Process PIL Image and return PIL Image with background removed or replaced.
|
||||
"""
|
||||
kwargs = {}
|
||||
if alpha_matting:
|
||||
kwargs.update({
|
||||
"alpha_matting": True,
|
||||
"alpha_matting_foreground_threshold": af,
|
||||
"alpha_matting_background_threshold": ab,
|
||||
"alpha_matting_erode_size": ae
|
||||
})
|
||||
|
||||
output_img = rembg.remove(image_input, session=self.session, **kwargs)
|
||||
|
||||
# Ensure RGBA mode
|
||||
if output_img.mode != "RGBA":
|
||||
output_img = output_img.convert("RGBA")
|
||||
|
||||
# If solid background color requested (e.g. Green screen "#00FF00")
|
||||
if bg_color and bg_color.strip() and bg_color.lower() != "transparent":
|
||||
try:
|
||||
rgb = ImageColor.getrgb(bg_color)
|
||||
background = Image.new("RGBA", output_img.size, rgb + (255,))
|
||||
background.paste(output_img, (0, 0), output_img)
|
||||
return background.convert("RGB")
|
||||
except Exception as e:
|
||||
print(f"Warning: Invalid background color '{bg_color}': {e}")
|
||||
|
||||
return output_img
|
||||
|
||||
def process_single_image(
|
||||
self,
|
||||
input_path: str,
|
||||
output_path: str,
|
||||
alpha_matting: bool = False,
|
||||
bg_color: Optional[str] = None
|
||||
) -> str:
|
||||
"""Process a single image file."""
|
||||
if not os.path.exists(input_path):
|
||||
raise FileNotFoundError(f"Input file not found: {input_path}")
|
||||
|
||||
img = Image.open(input_path)
|
||||
processed = self.remove_background(
|
||||
img,
|
||||
alpha_matting=alpha_matting,
|
||||
bg_color=bg_color
|
||||
)
|
||||
|
||||
os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
|
||||
|
||||
# Save based on file extension
|
||||
ext = os.path.splitext(output_path)[1].lower()
|
||||
if ext in [".jpg", ".jpeg"] and processed.mode == "RGBA":
|
||||
# Convert RGBA to RGB with white background for JPEG
|
||||
bg = Image.new("RGB", processed.size, (255, 255, 255))
|
||||
bg.paste(processed, (0, 0), processed)
|
||||
bg.save(output_path, quality=95)
|
||||
else:
|
||||
processed.save(output_path)
|
||||
|
||||
return output_path
|
||||
|
||||
def process_batch_images(
|
||||
self,
|
||||
input_paths: List[str],
|
||||
output_dir: str,
|
||||
alpha_matting: bool = False,
|
||||
bg_color: Optional[str] = None,
|
||||
output_format: str = "png",
|
||||
progress_callback: Optional[Callable[[int, int, str], None]] = None
|
||||
) -> List[str]:
|
||||
"""Process a list of image paths or directory."""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
results = []
|
||||
total = len(input_paths)
|
||||
|
||||
for i, in_path in enumerate(input_paths, 1):
|
||||
fname = os.path.splitext(os.path.basename(in_path))[0]
|
||||
out_ext = ".png" if output_format.lower() == "png" else f".{output_format.lower()}"
|
||||
out_path = os.path.join(output_dir, f"{fname}_nobg{out_ext}")
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(i, total, os.path.basename(in_path))
|
||||
|
||||
res = self.process_single_image(
|
||||
in_path,
|
||||
out_path,
|
||||
alpha_matting=alpha_matting,
|
||||
bg_color=bg_color
|
||||
)
|
||||
results.append(res)
|
||||
|
||||
return results
|
||||
|
||||
def process_video(
|
||||
self,
|
||||
video_path: str,
|
||||
output_path: str,
|
||||
fps: Optional[float] = None,
|
||||
export_format: str = "mov_transparent",
|
||||
bg_color: Optional[str] = None,
|
||||
alpha_matting: bool = False,
|
||||
progress_callback: Optional[Callable[[int, int, str], None]] = None
|
||||
) -> str:
|
||||
"""
|
||||
Extract video frames in lossless PNG format via ffmpeg, process frames with rembg,
|
||||
and reassemble into output video or PNG frame sequence.
|
||||
"""
|
||||
if not ffmpeg_utils.check_ffmpeg_installed():
|
||||
raise RuntimeError("FFmpeg is not installed or not found on system PATH.")
|
||||
|
||||
if not os.path.exists(video_path):
|
||||
raise FileNotFoundError(f"Video file not found: {video_path}")
|
||||
|
||||
# Get video metadata
|
||||
vinfo = ffmpeg_utils.get_video_info(video_path)
|
||||
actual_fps = fps if (fps and fps > 0) else vinfo.get("fps", 30.0)
|
||||
|
||||
temp_dir = tempfile.mkdtemp(prefix="rembg_video_")
|
||||
extracted_frames_dir = os.path.join(temp_dir, "extracted")
|
||||
processed_frames_dir = os.path.join(temp_dir, "processed")
|
||||
audio_file = os.path.join(temp_dir, "audio.aac")
|
||||
|
||||
try:
|
||||
# 1. Extract frames losslessly
|
||||
if progress_callback:
|
||||
progress_callback(0, 100, "Extracting frames losslessly with FFmpeg...")
|
||||
num_frames = ffmpeg_utils.extract_frames_lossless(
|
||||
video_path,
|
||||
extracted_frames_dir,
|
||||
target_fps=actual_fps
|
||||
)
|
||||
|
||||
# 2. Extract audio track if present
|
||||
has_audio = False
|
||||
if vinfo.get("has_audio", False):
|
||||
if progress_callback:
|
||||
progress_callback(0, 100, "Extracting audio track...")
|
||||
has_audio = ffmpeg_utils.extract_audio(video_path, audio_file)
|
||||
|
||||
# 3. Process frames with rembg
|
||||
os.makedirs(processed_frames_dir, exist_ok=True)
|
||||
frame_files = sorted([f for f in os.listdir(extracted_frames_dir) if f.endswith(".png")])
|
||||
|
||||
for idx, frame_name in enumerate(frame_files, 1):
|
||||
if progress_callback:
|
||||
progress_callback(idx, len(frame_files), f"Removing background from frame {idx}/{len(frame_files)}")
|
||||
|
||||
in_f = os.path.join(extracted_frames_dir, frame_name)
|
||||
out_f = os.path.join(processed_frames_dir, frame_name)
|
||||
|
||||
self.process_single_image(
|
||||
in_f,
|
||||
out_f,
|
||||
alpha_matting=alpha_matting,
|
||||
bg_color=bg_color
|
||||
)
|
||||
|
||||
# 4. Reassemble output video / PNG sequence
|
||||
if progress_callback:
|
||||
progress_callback(len(frame_files), len(frame_files), "Reassembling video...")
|
||||
|
||||
final_output = ffmpeg_utils.reassemble_video(
|
||||
frames_dir=processed_frames_dir,
|
||||
output_video_path=output_path,
|
||||
fps=actual_fps,
|
||||
audio_path=audio_file if has_audio else None,
|
||||
export_format=export_format,
|
||||
bg_color=bg_color
|
||||
)
|
||||
|
||||
return final_output
|
||||
|
||||
finally:
|
||||
# Cleanup temp directory
|
||||
shutil.rmtree(temp_dir, ignore_errors=True)
|
||||
Reference in New Issue
Block a user