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,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