31 lines
741 B
Python
31 lines
741 B
Python
from typing import List
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
from PIL.Image import Image as PILImage
|
|
|
|
from .session_base import BaseSession
|
|
|
|
|
|
class SimpleSession(BaseSession):
|
|
def predict(self, img: PILImage) -> List[PILImage]:
|
|
ort_outs = self.inner_session.run(
|
|
None,
|
|
self.normalize(
|
|
img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (320, 320)
|
|
),
|
|
)
|
|
|
|
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]
|