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.
221 lines
7.8 KiB
Python
221 lines
7.8 KiB
Python
#!/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()
|