73 lines
2.0 KiB
Python
73 lines
2.0 KiB
Python
import os
|
|
from typing import List
|
|
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
from PIL import Image
|
|
from PIL.Image import Image as PILImage
|
|
from rembg.sessions import BaseSession
|
|
|
|
|
|
class CustomBaseSession(BaseSession):
|
|
def __init__(self, model_name: str):
|
|
sess_opts = ort.SessionOptions()
|
|
|
|
if "OMP_NUM_THREADS" in os.environ:
|
|
sess_opts.inter_op_num_threads = int(os.environ["OMP_NUM_THREADS"])
|
|
|
|
super().__init__(model_name, sess_opts)
|
|
|
|
|
|
class CustomSessionContainer:
|
|
def __init__(self, mean_x, mean_y, mean_z, std_x, std_y, std_z, width, height) -> None:
|
|
self.mean_x = mean_x
|
|
self.mean_y = mean_y
|
|
self.mean_z = mean_z
|
|
self.std_x = std_x
|
|
self.std_y = std_y
|
|
self.std_z = std_z
|
|
self.width = width
|
|
self.height = height
|
|
|
|
|
|
class CustomAbstractSession(CustomBaseSession, CustomSessionContainer):
|
|
def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
|
|
ort_outs = self.inner_session.run(
|
|
None,
|
|
self.normalize(
|
|
img,
|
|
(self.mean_x, self.mean_y, self.mean_z),
|
|
(self.std_x, self.std_y, self.std_z),
|
|
(self.width, self.height)
|
|
),
|
|
)
|
|
|
|
pred = ort_outs[0][:, 0, :, :]
|
|
|
|
ma = np.max(pred)
|
|
mi = np.min(pred)
|
|
|
|
pred = (pred - mi) / (ma - mi)
|
|
pred = np.squeeze(pred)
|
|
|
|
mask = Image.fromarray((pred * 255).astype("uint8"), mode="L")
|
|
mask = mask.resize(img.size, Image.LANCZOS)
|
|
|
|
return [mask]
|
|
|
|
@classmethod
|
|
def download_models(cls, *args, **kwargs):
|
|
return os.path.join(cls.u2net_home(), f"{cls.name()}")
|
|
|
|
def from_container(self, container: CustomSessionContainer):
|
|
self.mean_x = container.mean_x
|
|
self.mean_y = container.mean_y
|
|
self.mean_z = container.mean_z
|
|
self.std_x = container.std_x
|
|
self.std_y = container.std_y
|
|
self.std_z = container.std_z
|
|
self.width = container.width
|
|
self.height = container.height
|
|
|
|
return self
|