Spaces:
Sleeping
Sleeping
File size: 2,336 Bytes
3faa99b 5f57808 3faa99b 5f57808 3faa99b 5f57808 3faa99b 5f57808 3faa99b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import os
from typing import Dict, List, Tuple
import numpy as np
import onnxruntime as ort
from PIL import Image
from PIL.Image import Image as PILImage
class BaseSession:
def __init__(
self,
model_name: str,
sess_opts: ort.SessionOptions,
providers=None,
*args,
**kwargs
):
self.model_name = model_name
self.providers = []
_providers = ort.get_available_providers()
if providers:
for provider in providers:
if provider in _providers:
self.providers.append(provider)
else:
self.providers.extend(_providers)
self.inner_session = ort.InferenceSession(
str(self.__class__.download_models()),
providers=self.providers,
sess_options=sess_opts,
)
def normalize(
self,
img: PILImage,
mean: Tuple[float, float, float],
std: Tuple[float, float, float],
size: Tuple[int, int],
*args,
**kwargs
) -> Dict[str, np.ndarray]:
im = img.convert("RGB").resize(size, Image.LANCZOS)
im_ary = np.array(im)
im_ary = im_ary / np.max(im_ary)
tmpImg = np.zeros((im_ary.shape[0], im_ary.shape[1], 3))
tmpImg[:, :, 0] = (im_ary[:, :, 0] - mean[0]) / std[0]
tmpImg[:, :, 1] = (im_ary[:, :, 1] - mean[1]) / std[1]
tmpImg[:, :, 2] = (im_ary[:, :, 2] - mean[2]) / std[2]
tmpImg = tmpImg.transpose((2, 0, 1))
return {
self.inner_session.get_inputs()[0]
.name: np.expand_dims(tmpImg, 0)
.astype(np.float32)
}
def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
raise NotImplementedError
@classmethod
def checksum_disabled(cls, *args, **kwargs):
return os.getenv("MODEL_CHECKSUM_DISABLED", None) is not None
@classmethod
def u2net_home(cls, *args, **kwargs):
return os.path.expanduser(
os.getenv(
"U2NET_HOME", os.path.join(os.getenv("XDG_DATA_HOME", "~"), ".u2net")
)
)
@classmethod
def download_models(cls, *args, **kwargs):
raise NotImplementedError
@classmethod
def name(cls, *args, **kwargs):
raise NotImplementedError
|