mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
8d8ab4bc45
Three fixes to ensure zero network access after docker pull: 1. rembg model allowlist: validate model parameter against the 7 pre-downloaded models, preventing rembg from attempting to download unknown models via a raw API call. 2. GFPGAN/CodeFormer auxiliary models: pre-download facexlib's detection_Resnet50_Final.pth and parsing_parsenet.pth at build time. These were previously downloaded on first use via basicsr. Symlinks in /app/gfpgan/weights/ ensure codeformer-pip also finds them. 3. OpenCV colorize models: pre-download the prototxt, caffemodel, and pts_in_hull.npy so the lightweight OpenCV colorizer fallback works in addition to the primary DDColor method. Co-authored-by: stirling-image <stirling-image@users.noreply.github.com>
129 lines
3.8 KiB
Python
129 lines
3.8 KiB
Python
"""Background removal using rembg with state-of-the-art BiRefNet models."""
|
|
import sys
|
|
import json
|
|
import os
|
|
|
|
|
|
def emit_progress(percent, stage):
|
|
"""Emit structured progress to stderr for bridge.ts to capture."""
|
|
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
|
|
|
|
|
|
ALLOWED_MODELS = {
|
|
"u2net",
|
|
"isnet-general-use",
|
|
"bria-rmbg",
|
|
"birefnet-general-lite",
|
|
"birefnet-portrait",
|
|
"birefnet-general",
|
|
"birefnet-matting",
|
|
}
|
|
|
|
_matting_registered = False
|
|
|
|
def _register_matting_session(sessions_class):
|
|
"""Register the BiRefNet-matting ONNX session for Ultra quality mode."""
|
|
global _matting_registered
|
|
if _matting_registered:
|
|
return
|
|
_matting_registered = True
|
|
|
|
import os
|
|
import pooch
|
|
from rembg.sessions.birefnet_general import BiRefNetSessionGeneral
|
|
|
|
class BiRefNetMattingSession(BiRefNetSessionGeneral):
|
|
@classmethod
|
|
def download_models(cls, *args, **kwargs):
|
|
fname = f"{cls.name(*args, **kwargs)}.onnx"
|
|
pooch.retrieve(
|
|
"https://github.com/ZhengPeng7/BiRefNet/releases/download/v1/BiRefNet-matting-epoch_100.onnx",
|
|
None, # Skip checksum for GitHub release assets
|
|
fname=fname,
|
|
path=cls.u2net_home(*args, **kwargs),
|
|
progressbar=True,
|
|
)
|
|
return os.path.join(cls.u2net_home(*args, **kwargs), fname)
|
|
|
|
@classmethod
|
|
def name(cls, *args, **kwargs):
|
|
return "birefnet-matting"
|
|
|
|
sessions_class.append(BiRefNetMattingSession)
|
|
|
|
|
|
def main():
|
|
input_path = sys.argv[1]
|
|
output_path = sys.argv[2]
|
|
settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {}
|
|
|
|
model = settings.get("model", "birefnet-general-lite")
|
|
if model not in ALLOWED_MODELS:
|
|
model = "birefnet-general-lite"
|
|
|
|
# Redirect stdout to stderr so library download/progress output
|
|
# cannot contaminate our JSON result on stdout.
|
|
stdout_fd = os.dup(1)
|
|
os.dup2(2, 1)
|
|
|
|
try:
|
|
from rembg import remove, new_session
|
|
from rembg.sessions import sessions_class
|
|
from gpu import onnx_providers
|
|
|
|
# Register BiRefNet-matting (Ultra quality) if not already present
|
|
_register_matting_session(sessions_class)
|
|
|
|
emit_progress(10, "Loading model")
|
|
|
|
session = new_session(model, providers=onnx_providers())
|
|
|
|
emit_progress(25, "Model loaded")
|
|
|
|
with open(input_path, "rb") as f:
|
|
input_data = f.read()
|
|
|
|
# Try with alpha matting for better edges, fall back without
|
|
emit_progress(30, "Analyzing image")
|
|
try:
|
|
output_data = remove(
|
|
input_data,
|
|
session=session,
|
|
alpha_matting=True,
|
|
alpha_matting_foreground_threshold=240,
|
|
alpha_matting_background_threshold=10,
|
|
)
|
|
except Exception:
|
|
output_data = remove(input_data, session=session)
|
|
|
|
emit_progress(80, "Background removed")
|
|
|
|
# Always return transparent PNG. All background compositing
|
|
# (solid color, gradient, blur, shadow) is handled by Node.js/Sharp.
|
|
|
|
emit_progress(95, "Saving result")
|
|
with open(output_path, "wb") as f:
|
|
f.write(output_data)
|
|
|
|
result = json.dumps({"success": True, "model": model})
|
|
|
|
except ImportError:
|
|
result = json.dumps(
|
|
{
|
|
"success": False,
|
|
"error": "rembg is not installed. Install with: pip install rembg[cpu]",
|
|
}
|
|
)
|
|
except Exception as e:
|
|
result = json.dumps({"success": False, "error": str(e)})
|
|
|
|
# Restore original stdout and write only our JSON result
|
|
os.dup2(stdout_fd, 1)
|
|
os.close(stdout_fd)
|
|
sys.stdout.write(result + "\n")
|
|
sys.stdout.flush()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|