#!/usr/bin/env python3 """ RemoveBG Sidecar - JSON-over-stdio worker process for the native UI. Protocol -------- stdin (one JSON object per line): {"id": , "cmd": , "params": {...}} stdout (one JSON object per line): {"id": , "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()