| import os |
| from typing import List |
|
|
| import numpy as np |
| import pooch |
| from PIL import Image |
| from PIL.Image import Image as PILImage |
|
|
| from .base import BaseSession |
|
|
|
|
| class DisSession(BaseSession): |
| def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]: |
| ort_outs = self.inner_session.run( |
| None, |
| self.normalize(img, (0.485, 0.456, 0.406), (1.0, 1.0, 1.0), (1024, 1024)), |
| ) |
|
|
| 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): |
| fname = f"{cls.name()}.onnx" |
| pooch.retrieve( |
| "https://github.com/danielgatis/rembg/releases/download/v0.0.0/isnet-general-use.onnx", |
| None |
| if cls.checksum_disabled(*args, **kwargs) |
| else "md5:fc16ebd8b0c10d971d3513d564d01e29", |
| fname=fname, |
| path=cls.u2net_home(*args, **kwargs), |
| progressbar=True, |
| ) |
|
|
| return os.path.join(cls.u2net_home(), fname) |
|
|
| @classmethod |
| def name(cls, *args, **kwargs): |
| return "isnet-general-use" |
|
|
