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 @@
|
||||
"""Core background removal and video processing package."""
|
||||
@@ -0,0 +1,239 @@
|
||||
import os
|
||||
import re
|
||||
import json
|
||||
import shutil
|
||||
import subprocess
|
||||
from typing import Dict, Any, Optional, Callable
|
||||
|
||||
def check_ffmpeg_installed() -> bool:
|
||||
"""Check if ffmpeg is available on the system PATH."""
|
||||
return shutil.which("ffmpeg") is not None
|
||||
|
||||
def check_ffprobe_installed() -> bool:
|
||||
"""Check if ffprobe is available on the system PATH."""
|
||||
return shutil.which("ffprobe") is not None
|
||||
|
||||
def get_video_info(video_path: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Get video metadata using ffprobe or ffmpeg.
|
||||
Returns dictionary with fps, width, height, duration, frame_count, has_audio.
|
||||
"""
|
||||
info = {
|
||||
"fps": 30.0,
|
||||
"width": 1920,
|
||||
"height": 1080,
|
||||
"duration": 0.0,
|
||||
"frame_count": 0,
|
||||
"has_audio": False
|
||||
}
|
||||
|
||||
if check_ffprobe_installed():
|
||||
cmd = [
|
||||
"ffprobe",
|
||||
"-v", "quiet",
|
||||
"-print_format", "json",
|
||||
"-show_format",
|
||||
"-show_streams",
|
||||
video_path
|
||||
]
|
||||
try:
|
||||
result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True)
|
||||
data = json.loads(result.stdout)
|
||||
|
||||
for stream in data.get("streams", []):
|
||||
if stream.get("codec_type") == "video" and "r_frame_rate" in stream:
|
||||
# Evaluate fps string like "30/1" or "30000/1001"
|
||||
r_fps = stream["r_frame_rate"]
|
||||
if "/" in r_fps:
|
||||
num, den = r_fps.split("/")
|
||||
if float(den) > 0:
|
||||
info["fps"] = float(num) / float(den)
|
||||
else:
|
||||
info["fps"] = float(r_fps)
|
||||
|
||||
info["width"] = int(stream.get("width", 1920))
|
||||
info["height"] = int(stream.get("height", 1080))
|
||||
|
||||
if "nb_frames" in stream and stream["nb_frames"].isdigit():
|
||||
info["frame_count"] = int(stream["nb_frames"])
|
||||
elif stream.get("codec_type") == "audio":
|
||||
info["has_audio"] = True
|
||||
|
||||
if "format" in data and "duration" in data["format"]:
|
||||
info["duration"] = float(data["format"]["duration"])
|
||||
if info["frame_count"] == 0 and info["fps"] > 0:
|
||||
info["frame_count"] = int(info["duration"] * info["fps"])
|
||||
return info
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Fallback parsing via ffmpeg stderr
|
||||
cmd = ["ffmpeg", "-i", video_path]
|
||||
res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||||
stderr = res.stderr
|
||||
|
||||
fps_match = re.search(r"(\d+(?:\.\d+)?)\s+fps", stderr)
|
||||
if fps_match:
|
||||
info["fps"] = float(fps_match.group(1))
|
||||
|
||||
res_match = re.search(r"(\d{2,5})x(\d{2,5})", stderr)
|
||||
if res_match:
|
||||
info["width"] = int(res_match.group(1))
|
||||
info["height"] = int(res_match.group(2))
|
||||
|
||||
dur_match = re.search(r"Duration:\s*(\d+):(\d+):(\d+(?:\.\d+)?)", stderr)
|
||||
if dur_match:
|
||||
h, m, s = float(dur_match.group(1)), float(dur_match.group(2)), float(dur_match.group(3))
|
||||
info["duration"] = h * 3600 + m * 60 + s
|
||||
info["frame_count"] = int(info["duration"] * info["fps"])
|
||||
|
||||
if "Audio:" in stderr:
|
||||
info["has_audio"] = True
|
||||
|
||||
return info
|
||||
|
||||
def extract_frames_lossless(
|
||||
video_path: str,
|
||||
output_dir: str,
|
||||
target_fps: Optional[float] = None,
|
||||
progress_callback: Optional[Callable[[str], None]] = None
|
||||
) -> int:
|
||||
"""
|
||||
Extract frames from a video into PNG images (100% Lossless) inside output_dir.
|
||||
Returns total number of extracted frames.
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
out_pattern = os.path.join(output_dir, "frame_%06d.png")
|
||||
|
||||
cmd = ["ffmpeg", "-y", "-i", video_path]
|
||||
|
||||
if target_fps and target_fps > 0:
|
||||
cmd.extend(["-vf", f"fps={target_fps}"])
|
||||
|
||||
# Lossless PNG output codec (-c:v png -pred mixed)
|
||||
cmd.extend(["-c:v", "png", "-pred", "mixed", out_pattern])
|
||||
|
||||
if progress_callback:
|
||||
progress_callback("Extracting video frames in lossless PNG format...")
|
||||
|
||||
process = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||||
if process.returncode != 0:
|
||||
raise RuntimeError(f"FFmpeg frame extraction failed: {process.stderr}")
|
||||
|
||||
frames = [f for f in os.listdir(output_dir) if f.startswith("frame_") and f.endswith(".png")]
|
||||
frames.sort()
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(f"Successfully extracted {len(frames)} lossless frames.")
|
||||
|
||||
return len(frames)
|
||||
|
||||
def extract_single_frame(video_path: str, output_path: str, time_offset: float = 0.5) -> str:
|
||||
"""Grab one representative frame from a video (for preview purposes)."""
|
||||
cmd = ["ffmpeg", "-y", "-ss", str(time_offset), "-i", video_path, "-frames:v", "1", output_path]
|
||||
res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||||
if res.returncode != 0 or not os.path.exists(output_path):
|
||||
raise RuntimeError(f"FFmpeg single-frame extraction failed: {res.stderr}")
|
||||
return output_path
|
||||
|
||||
def extract_audio(video_path: str, output_audio_path: str) -> bool:
|
||||
"""Extract audio stream from video file if present."""
|
||||
cmd = ["ffmpeg", "-y", "-i", video_path, "-vn", "-acodec", "aac", "-b:a", "192k", output_audio_path]
|
||||
res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||||
return res.returncode == 0 and os.path.exists(output_audio_path) and os.path.getsize(output_audio_path) > 0
|
||||
|
||||
def reassemble_video(
|
||||
frames_dir: str,
|
||||
output_video_path: str,
|
||||
fps: float = 30.0,
|
||||
audio_path: Optional[str] = None,
|
||||
export_format: str = "mov_transparent",
|
||||
bg_color: Optional[str] = None,
|
||||
progress_callback: Optional[Callable[[str], None]] = None
|
||||
) -> str:
|
||||
"""
|
||||
Reassemble processed frame sequence back into a video file or frame folder.
|
||||
|
||||
export_format options:
|
||||
- 'mov_transparent': QuickTime MOV with ProRes 4444 + Alpha Transparency
|
||||
- 'webm_transparent': WebM with VP9 + Alpha Transparency
|
||||
- 'mp4_solid': MP4 (H.264) with solid background color (e.g. Green screen #00FF00)
|
||||
- 'png_sequence': Copies/Leaves processed PNG frames in output directory
|
||||
"""
|
||||
if export_format == "png_sequence":
|
||||
out_dir = output_video_path
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
for fname in os.listdir(frames_dir):
|
||||
if fname.endswith(".png"):
|
||||
shutil.copy2(os.path.join(frames_dir, fname), os.path.join(out_dir, fname))
|
||||
if progress_callback:
|
||||
progress_callback(f"PNG sequence exported to {out_dir}")
|
||||
return out_dir
|
||||
|
||||
input_pattern = os.path.join(frames_dir, "frame_%06d.png")
|
||||
|
||||
# Ensure parent output directory exists
|
||||
os.makedirs(os.path.dirname(os.path.abspath(output_video_path)), exist_ok=True)
|
||||
|
||||
cmd = ["ffmpeg", "-y", "-framerate", str(fps), "-i", input_pattern]
|
||||
|
||||
if audio_path and os.path.exists(audio_path):
|
||||
cmd.extend(["-i", audio_path, "-c:a", "aac", "-shortest"])
|
||||
|
||||
if export_format == "mov_transparent":
|
||||
# ProRes 4444 with Alpha Channel (Lossless/visually lossless transparent video)
|
||||
cmd.extend([
|
||||
"-c:v", "prores_ks",
|
||||
"-profile:v", "4",
|
||||
"-pix_fmt", "yuva444p10le",
|
||||
output_video_path
|
||||
])
|
||||
elif export_format == "webm_transparent":
|
||||
# VP9 WebM with Alpha Channel
|
||||
cmd.extend([
|
||||
"-c:v", "libvpx-vp9",
|
||||
"-pix_fmt", "yuva420p",
|
||||
"-b:v", "4M",
|
||||
output_video_path
|
||||
])
|
||||
elif export_format == "mp4_solid":
|
||||
# Standard MP4 H.264
|
||||
cmd.extend([
|
||||
"-c:v", "libx264",
|
||||
"-pix_fmt", "yuv420p",
|
||||
"-crf", "18",
|
||||
output_video_path
|
||||
])
|
||||
else:
|
||||
# Default MP4 fallback
|
||||
cmd.extend([
|
||||
"-c:v", "libx264",
|
||||
"-pix_fmt", "yuv420p",
|
||||
output_video_path
|
||||
])
|
||||
|
||||
if progress_callback:
|
||||
progress_callback("Reassembling frames into output video format...")
|
||||
|
||||
res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||||
if res.returncode != 0:
|
||||
# Fallback to standard libx264 if prores_ks or libvpx-vp9 isn't available
|
||||
if export_format in ["mov_transparent", "webm_transparent"]:
|
||||
if progress_callback:
|
||||
progress_callback(f"Codec warning: trying fallback MP4 export for {output_video_path}...")
|
||||
fallback_cmd = [
|
||||
"ffmpeg", "-y", "-framerate", str(fps), "-i", input_pattern
|
||||
]
|
||||
if audio_path and os.path.exists(audio_path):
|
||||
fallback_cmd.extend(["-i", audio_path, "-c:a", "aac", "-shortest"])
|
||||
fallback_cmd.extend(["-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "18", output_video_path])
|
||||
res_fb = subprocess.run(fallback_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||||
if res_fb.returncode != 0:
|
||||
raise RuntimeError(f"FFmpeg video reassembly failed: {res_fb.stderr}")
|
||||
else:
|
||||
raise RuntimeError(f"FFmpeg video reassembly failed: {res.stderr}")
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(f"Video created successfully: {output_video_path}")
|
||||
|
||||
return output_video_path
|
||||
@@ -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)
|
||||
@@ -0,0 +1,220 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
RemoveBG Sidecar - JSON-over-stdio worker process for the native UI.
|
||||
|
||||
Protocol
|
||||
--------
|
||||
stdin (one JSON object per line): {"id": <int>, "cmd": <str>, "params": {...}}
|
||||
stdout (one JSON object per line): {"id": <int>, "event": "progress"|"log"|"done"|"error", ...}
|
||||
|
||||
A single {"event": "ready"} line (no "id") is emitted once on startup once the
|
||||
process is alive and reading. Model loading happens lazily on first use so
|
||||
startup is instant.
|
||||
|
||||
Anything a dependency (rembg / onnxruntime / tqdm) writes via print() or to
|
||||
sys.stdout is redirected to stderr so it can never corrupt the JSON stream.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import io
|
||||
import json
|
||||
import base64
|
||||
import tempfile
|
||||
import shutil
|
||||
import traceback
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
# Protect the protocol stream: keep the real stdout for our JSON lines only,
|
||||
# and reroute anything the rest of the process prints to stderr.
|
||||
_protocol_out = sys.stdout
|
||||
sys.stdout = sys.stderr
|
||||
|
||||
from PIL import Image
|
||||
from core.processor import BackgroundRemoverEngine, AVAILABLE_MODELS
|
||||
from core import ffmpeg_utils
|
||||
|
||||
PREVIEW_MAX_DIM = 1400
|
||||
|
||||
|
||||
def _encode_preview_pair(original: "Image.Image", processed: "Image.Image") -> Dict[str, Any]:
|
||||
proc_buf = io.BytesIO()
|
||||
processed.save(proc_buf, format="PNG")
|
||||
|
||||
orig_buf = io.BytesIO()
|
||||
original.convert("RGB").save(orig_buf, format="JPEG", quality=88)
|
||||
|
||||
return {
|
||||
"original_b64": base64.b64encode(orig_buf.getvalue()).decode("ascii"),
|
||||
"preview_b64": base64.b64encode(proc_buf.getvalue()).decode("ascii"),
|
||||
"width": processed.size[0],
|
||||
"height": processed.size[1],
|
||||
}
|
||||
|
||||
|
||||
def _downscale_for_preview(img: "Image.Image") -> "Image.Image":
|
||||
w, h = img.size
|
||||
scale = min(1.0, PREVIEW_MAX_DIM / max(w, h))
|
||||
if scale < 1.0:
|
||||
img = img.resize((max(1, int(w * scale)), max(1, int(h * scale))), Image.LANCZOS)
|
||||
return img
|
||||
|
||||
|
||||
def emit(payload: Dict[str, Any]) -> None:
|
||||
_protocol_out.write(json.dumps(payload, ensure_ascii=False) + "\n")
|
||||
_protocol_out.flush()
|
||||
|
||||
|
||||
class Sidecar:
|
||||
def __init__(self):
|
||||
self.engine: Optional[BackgroundRemoverEngine] = None
|
||||
self.current_model: Optional[str] = None
|
||||
self.cancel_requested = False
|
||||
|
||||
def get_engine(self, model_name: str) -> BackgroundRemoverEngine:
|
||||
if self.engine is None or self.current_model != model_name:
|
||||
emit({"event": "log", "message": f"Loading model '{model_name}'..."})
|
||||
self.engine = BackgroundRemoverEngine(model_name=model_name)
|
||||
self.current_model = model_name
|
||||
return self.engine
|
||||
|
||||
# ---- command handlers -------------------------------------------------
|
||||
|
||||
def cmd_check_ffmpeg(self, req_id: int, params: Dict[str, Any]) -> Dict[str, Any]:
|
||||
return {"installed": ffmpeg_utils.check_ffmpeg_installed()}
|
||||
|
||||
def cmd_list_models(self, req_id: int, params: Dict[str, Any]) -> Dict[str, Any]:
|
||||
return {"models": AVAILABLE_MODELS}
|
||||
|
||||
def cmd_get_video_info(self, req_id: int, params: Dict[str, Any]) -> Dict[str, Any]:
|
||||
path = params["path"]
|
||||
if not os.path.exists(path):
|
||||
raise FileNotFoundError(f"Video file not found: {path}")
|
||||
return ffmpeg_utils.get_video_info(path)
|
||||
|
||||
def cmd_preview_image(self, req_id: int, params: Dict[str, Any]) -> Dict[str, Any]:
|
||||
engine = self.get_engine(params.get("model", "u2net"))
|
||||
img = _downscale_for_preview(Image.open(params["input"]))
|
||||
processed = engine.remove_background(
|
||||
img,
|
||||
alpha_matting=bool(params.get("alpha_matting", False)),
|
||||
bg_color=params.get("bg_color"),
|
||||
)
|
||||
return _encode_preview_pair(img, processed)
|
||||
|
||||
def cmd_preview_video_frame(self, req_id: int, params: Dict[str, Any]) -> Dict[str, Any]:
|
||||
if not ffmpeg_utils.check_ffmpeg_installed():
|
||||
raise RuntimeError("FFmpeg is not installed or not found on system PATH.")
|
||||
|
||||
engine = self.get_engine(params.get("model", "u2net"))
|
||||
tmp_dir = tempfile.mkdtemp(prefix="rembg_preview_")
|
||||
try:
|
||||
frame_path = os.path.join(tmp_dir, "frame.png")
|
||||
ffmpeg_utils.extract_single_frame(
|
||||
params["video"], frame_path, time_offset=float(params.get("time", 0.5))
|
||||
)
|
||||
img = _downscale_for_preview(Image.open(frame_path))
|
||||
processed = engine.remove_background(
|
||||
img,
|
||||
alpha_matting=bool(params.get("alpha_matting", False)),
|
||||
bg_color=params.get("bg_color"),
|
||||
)
|
||||
return _encode_preview_pair(img, processed)
|
||||
finally:
|
||||
shutil.rmtree(tmp_dir, ignore_errors=True)
|
||||
|
||||
def cmd_process_image(self, req_id: int, params: Dict[str, Any]) -> Dict[str, Any]:
|
||||
engine = self.get_engine(params.get("model", "u2net"))
|
||||
out = engine.process_single_image(
|
||||
params["input"],
|
||||
params["output"],
|
||||
alpha_matting=bool(params.get("alpha_matting", False)),
|
||||
bg_color=params.get("bg_color"),
|
||||
)
|
||||
return {"output": out}
|
||||
|
||||
def cmd_process_batch(self, req_id: int, params: Dict[str, Any]) -> Dict[str, Any]:
|
||||
engine = self.get_engine(params.get("model", "u2net"))
|
||||
inputs = params["inputs"]
|
||||
total = len(inputs)
|
||||
|
||||
def pcb(current, total_count, fname):
|
||||
emit({"id": req_id, "event": "progress", "current": current, "total": total_count, "message": fname})
|
||||
|
||||
results = engine.process_batch_images(
|
||||
input_paths=inputs,
|
||||
output_dir=params["output_dir"],
|
||||
alpha_matting=bool(params.get("alpha_matting", False)),
|
||||
bg_color=params.get("bg_color"),
|
||||
output_format=params.get("format", "png"),
|
||||
progress_callback=pcb,
|
||||
)
|
||||
return {"results": results, "count": len(results)}
|
||||
|
||||
def cmd_process_video(self, req_id: int, params: Dict[str, Any]) -> Dict[str, Any]:
|
||||
if not ffmpeg_utils.check_ffmpeg_installed():
|
||||
raise RuntimeError("FFmpeg is not installed or not found on system PATH.")
|
||||
|
||||
engine = self.get_engine(params.get("model", "u2net"))
|
||||
fps = params.get("fps") or None
|
||||
|
||||
def pcb(current, total_count, message):
|
||||
emit({"id": req_id, "event": "progress", "current": current, "total": total_count, "message": message})
|
||||
|
||||
out = engine.process_video(
|
||||
video_path=params["video"],
|
||||
output_path=params["output"],
|
||||
fps=fps,
|
||||
export_format=params.get("export_format", "mov_transparent"),
|
||||
bg_color=params.get("bg_color"),
|
||||
alpha_matting=bool(params.get("alpha_matting", False)),
|
||||
progress_callback=pcb,
|
||||
)
|
||||
return {"output": out}
|
||||
|
||||
def dispatch(self, req: Dict[str, Any]) -> None:
|
||||
req_id = req.get("id")
|
||||
cmd = req.get("cmd")
|
||||
params = req.get("params") or {}
|
||||
|
||||
handler = getattr(self, f"cmd_{cmd}", None)
|
||||
if handler is None:
|
||||
emit({"id": req_id, "event": "error", "message": f"Unknown command: {cmd}"})
|
||||
return
|
||||
|
||||
try:
|
||||
result = handler(req_id, params)
|
||||
emit({"id": req_id, "event": "done", "result": result})
|
||||
except Exception as e:
|
||||
emit({
|
||||
"id": req_id,
|
||||
"event": "error",
|
||||
"message": str(e),
|
||||
"detail": traceback.format_exc(),
|
||||
})
|
||||
|
||||
|
||||
def main() -> None:
|
||||
sidecar = Sidecar()
|
||||
emit({"event": "ready"})
|
||||
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
req = json.loads(line)
|
||||
except json.JSONDecodeError as e:
|
||||
emit({"event": "error", "message": f"Malformed JSON request: {e}"})
|
||||
continue
|
||||
|
||||
if req.get("cmd") == "shutdown":
|
||||
break
|
||||
|
||||
sidecar.dispatch(req)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user