|
import os |
|
import shutil |
|
import subprocess |
|
import sys |
|
from typing import List |
|
|
|
import numpy as np |
|
import pooch |
|
from PIL import Image |
|
from PIL.Image import Image as PILImage |
|
|
|
from .CustomSession import CustomBaseSession |
|
|
|
|
|
class ModnetPhotographicSession(CustomBaseSession): |
|
def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]: |
|
ort_outs = self.inner_session.run( |
|
None, |
|
self.normalize(img, (0.5, 0.5, 0.5), (1.0, 1.0, 1.0), (512, 512)), |
|
) |
|
|
|
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" |
|
|
|
if not os.path.exists(os.path.join(cls.u2net_home(), fname)): |
|
pooch.retrieve( |
|
"https://raw.githubusercontent.com/ZHKKKe/MODNet/master/onnx/export_onnx.py", |
|
"SHA256:647990c98c409fbf6a72cd2a2db5fe19d2e4b15a3a436ef0302be0582458b63e", |
|
fname=f"export_onnx.py", |
|
path=os.path.join(cls.u2net_home(), "modnet-p/"), |
|
progressbar=True, |
|
) |
|
|
|
pooch.retrieve( |
|
"https://raw.githubusercontent.com/ZHKKKe/MODNet/master/onnx/modnet_onnx.py", |
|
"SHA256:0502cad1b7ab0bf2f866179960454c1e63df096390db05e93cb40145dbc26e1f", |
|
fname=f"modnet_onnx.py", |
|
path=os.path.join(cls.u2net_home(), "modnet-p/"), |
|
progressbar=True, |
|
) |
|
|
|
pooch.retrieve( |
|
"https://raw.githubusercontent.com/ZHKKKe/MODNet/master/src/models/backbones/__init__.py", |
|
"SHA256:28a5fb95f7dcf9e365edbf42c6d2e8ea0ca4839e51fd7f11bd0547d2359fcd96", |
|
fname=f"__init__.py", |
|
path=os.path.join(cls.u2net_home(), "modnet-p/src/models/backbones"), |
|
progressbar=True, |
|
) |
|
|
|
pooch.retrieve( |
|
"https://raw.githubusercontent.com/ZHKKKe/MODNet/master/src/models/backbones/mobilenetv2.py", |
|
"SHA256:e3cc8ad6a9933ba18a17a62d5f887c64e0721240871ea8b48742fb9a8a2c3199", |
|
fname=f"mobilenetv2.py", |
|
path=os.path.join(cls.u2net_home(), "modnet-p/src/models/backbones"), |
|
progressbar=True, |
|
) |
|
|
|
pooch.retrieve( |
|
"https://raw.githubusercontent.com/ZHKKKe/MODNet/master/src/models/backbones/wrapper.py", |
|
"SHA256:41197be7eb96b8a60dc034b55d8c9340dd682a41441dcf2ce67238955dfa5607", |
|
fname=f"wrapper.py", |
|
path=os.path.join(cls.u2net_home(), "modnet-p/src/models/backbones"), |
|
progressbar=True, |
|
) |
|
|
|
pooch.retrieve( |
|
"https://storage.openvinotoolkit.org/repositories/open_model_zoo/public/2022.2/modnet-photographic-portrait-matting/modnet_photographic_portrait_matting.ckpt", |
|
"SHA256:7c22235f0925deba15d4d63e53afcb654c47055bbcd98f56e393ab2584007ed8", |
|
fname=f"modnet_photographic_portrait_matting.ckpt", |
|
path=os.path.join(cls.u2net_home(), "modnet-p/"), |
|
progressbar=True, |
|
) |
|
|
|
replace_line( |
|
os.path.join(cls.u2net_home(), "modnet-p/export_onnx.py"), |
|
"from . import modnet_onnx", |
|
"import modnet_onnx" |
|
) |
|
|
|
subprocess.run([ |
|
sys.executable, |
|
os.path.join(cls.u2net_home(), "modnet-p/export_onnx.py"), |
|
"--ckpt-path=" + os.path.join(cls.u2net_home(), "modnet-p/modnet_photographic_portrait_matting.ckpt"), |
|
"--output-path=" + os.path.join(cls.u2net_home(), "modnet-p/../modnet-p.onnx"), |
|
]) |
|
|
|
shutil.rmtree(os.path.join(cls.u2net_home(), "modnet-p/")) |
|
|
|
return os.path.join(cls.u2net_home(), fname) |
|
|
|
@classmethod |
|
def name(cls, *args, **kwargs): |
|
return "modnet-p" |
|
|
|
|
|
def replace_line(path: str, old: str, new: str): |
|
with open(path, "r", encoding="utf-8") as file: |
|
data = file.readlines() |
|
|
|
for i in range(len(data)): |
|
if data[i].__contains__(old): |
|
data[i] = data[i].replace(old, new) |
|
|
|
with open(path, "w", encoding="utf-8") as file: |
|
file.writelines(data) |
|
|
