71 lines
2.7 KiB
Python
71 lines
2.7 KiB
Python
|
# /app/server/ormbg.py
|
||
|
|
from __future__ import annotations
|
||
|
|
import io
|
||
|
|
import os
|
||
|
|
from typing import Optional, Union
|
||
|
|
|
||
|
|
from rembg import remove as rembg_remove, new_session
|
||
|
|
from PIL import Image
|
||
|
|
import onnxruntime as ort
|
||
|
|
|
||
|
|
# Choose GPU if available, else CPU
|
||
|
|
def _default_providers() -> list[str]:
|
||
|
|
providers = ort.get_available_providers()
|
||
|
|
return ["CUDAExecutionProvider", "CPUExecutionProvider"] if "CUDAExecutionProvider" in providers else ["CPUExecutionProvider"]
|
||
|
|
|
||
|
|
class ORMBGProcessor:
|
||
|
|
"""
|
||
|
|
Minimal adapter used by server.py.
|
||
|
|
Provides:
|
||
|
|
- .warmup()
|
||
|
|
- .process_bytes(input_bytes) -> bytes (PNG with alpha)
|
||
|
|
- .process_pil(img: PIL.Image.Image) -> PIL.Image.Image
|
||
|
|
- .close() (no-op; for API symmetry)
|
||
|
|
"""
|
||
|
|
def __init__(self, model_name: str = "isnet-general-use", providers: Optional[list[str]] = None):
|
||
|
|
# Some rembg internals pick providers automatically; onnxruntime-gpu presence
|
||
|
|
# will make CUDA provider available. We still compute a list so we can log/verify.
|
||
|
|
self.providers = providers or _default_providers()
|
||
|
|
# rembg's API chooses the best available EP under the hood if ORT-GPU is present.
|
||
|
|
# Just ensure the model name exists; common: isnet-general-use, u2net, u2netp, u2net_human_seg
|
||
|
|
self.model_name = model_name
|
||
|
|
self.session = new_session(self.model_name)
|
||
|
|
|
||
|
|
def warmup(self) -> None:
|
||
|
|
# Tiny 1x1 transparent PNG to trigger model load
|
||
|
|
img = Image.new("RGBA", (1, 1), (0, 0, 0, 0))
|
||
|
|
_ = self.process_pil(img)
|
||
|
|
|
||
|
|
def process_bytes(self, data: bytes, return_png: bool = True) -> bytes:
|
||
|
|
"""
|
||
|
|
Accept raw image bytes (jpg/png/etc), return PNG bytes with alpha.
|
||
|
|
"""
|
||
|
|
out = rembg_remove(data, session=self.session)
|
||
|
|
if return_png:
|
||
|
|
# rembg already returns PNG bytes by default; if it's bytes, just pass through
|
||
|
|
if isinstance(out, (bytes, bytearray)):
|
||
|
|
return bytes(out)
|
||
|
|
# if it's a PIL image, convert to PNG bytes
|
||
|
|
if isinstance(out, Image.Image):
|
||
|
|
buf = io.BytesIO()
|
||
|
|
out.save(buf, format="PNG")
|
||
|
|
return buf.getvalue()
|
||
|
|
# fallback: ensure bytes
|
||
|
|
if isinstance(out, Image.Image):
|
||
|
|
buf = io.BytesIO()
|
||
|
|
out.save(buf, format="PNG")
|
||
|
|
return buf.getvalue()
|
||
|
|
return bytes(out)
|
||
|
|
|
||
|
|
def process_pil(self, img: Image.Image) -> Image.Image:
|
||
|
|
out = rembg_remove(img, session=self.session)
|
||
|
|
if isinstance(out, Image.Image):
|
||
|
|
return out
|
||
|
|
# if rembg returned bytes, decode back to PIL
|
||
|
|
return Image.open(io.BytesIO(out)).convert("RGBA")
|
||
|
|
|
||
|
|
def close(self) -> None:
|
||
|
|
# Nothing to dispose explicitly; keep for API compatibility
|
||
|
|
pass
|
||
|
|
|