| 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 |
|
|
