Spaces:
Runtime error
Runtime error
| 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, inner_session: ort.InferenceSession): | |
| self.model_name = model_name | |
| self.inner_session = inner_session | |
| def normalize( | |
| self, | |
| img: PILImage, | |
| mean: Tuple[float, float, float], | |
| std: Tuple[float, float, float], | |
| size: Tuple[int, int], | |
| ) -> 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) -> List[PILImage]: | |
| raise NotImplementedError | |