120 lines
4.4 KiB
Python
120 lines
4.4 KiB
Python
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)
|