File size: 2,091 Bytes
1e3b872 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
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
|