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.
224 lines
8.1 KiB
Python
224 lines
8.1 KiB
Python
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)
|