| from typing import List, Tuple, Union
|
|
|
| import cv2
|
| import numpy as np
|
| from config import ModelType
|
| from numpy import ndarray
|
|
|
|
|
| class Preprocess:
|
|
|
| def __init__(self, model_type: ModelType):
|
| if model_type in (ModelType.YOLOV5, ModelType.YOLOV6, ModelType.YOLOV7,
|
| ModelType.YOLOV8):
|
| mean = np.array([0, 0, 0], dtype=np.float32)
|
| std = np.array([255, 255, 255], dtype=np.float32)
|
| is_rgb = True
|
| elif model_type == ModelType.YOLOX:
|
| mean = np.array([0, 0, 0], dtype=np.float32)
|
| std = np.array([1, 1, 1], dtype=np.float32)
|
| is_rgb = False
|
| elif model_type == ModelType.PPYOLOE:
|
| mean = np.array([123.675, 116.28, 103.53], dtype=np.float32)
|
| std = np.array([58.395, 57.12, 57.375], dtype=np.float32)
|
| is_rgb = True
|
|
|
| elif model_type == ModelType.PPYOLOEP:
|
| mean = np.array([0, 0, 0], dtype=np.float32)
|
| std = np.array([255, 255, 255], dtype=np.float32)
|
| is_rgb = True
|
| elif model_type == ModelType.RTMDET:
|
| mean = np.array([103.53, 116.28, 123.675], dtype=np.float32)
|
| std = np.array([57.375, 57.12, 58.3955], dtype=np.float32)
|
| is_rgb = False
|
| else:
|
| raise NotImplementedError
|
|
|
| self.mean = mean.reshape((3, 1, 1))
|
| self.std = std.reshape((3, 1, 1))
|
| self.is_rgb = is_rgb
|
|
|
| def __call__(self,
|
| image: ndarray,
|
| new_size: Union[List[int], Tuple[int]] = (640, 640),
|
| **kwargs) -> Tuple[ndarray, Tuple[float, float]]:
|
|
|
| height, width = image.shape[:2]
|
| ratio_h, ratio_w = new_size[0] / height, new_size[1] / width
|
| image = cv2.resize(
|
| image, (0, 0),
|
| fx=ratio_w,
|
| fy=ratio_h,
|
| interpolation=cv2.INTER_LINEAR)
|
| image = np.ascontiguousarray(image.transpose(2, 0, 1))
|
| image = image.astype(np.float32)
|
| image -= self.mean
|
| image /= self.std
|
| return image[np.newaxis], (ratio_w, ratio_h)
|
|
|
