import cv2 import numpy as np class ResizeImage(object): def __init__(self, resize_short=None): self.resize_short = resize_short def __call__(self, img): img_h, img_w = img.shape[:2] percent = float(self.resize_short) / min(img_w, img_h) w = int(round(img_w * percent)) h = int(round(img_h * percent)) return cv2.resize(img, (w, h)) class CropImage(object): def __init__(self, size): if type(size) is int: self.size = (size, size) else: self.size = size def __call__(self, img): w, h = self.size img_h, img_w = img.shape[:2] w_start = (img_w - w) // 2 h_start = (img_h - h) // 2 w_end = w_start + w h_end = h_start + h return img[h_start:h_end, w_start:w_end, :] class ToTensor(object): def __init__(self): pass def __call__(self, img): img = img.transpose((2, 0, 1)) return img class NormalizeImage(object): def __init__(self, scale=None, mean=None, std=None): self.scale = np.float32(scale if scale is not None else 1.0 / 255.0) mean = mean if mean is not None else [0.485, 0.456, 0.406] std = std if std is not None else [0.229, 0.224, 0.225] shape = (1, 1, 3) self.mean = np.array(mean).reshape(shape).astype('float32') self.std = np.array(std).reshape(shape).astype('float32') def __call__(self, img): return (img.astype('float32') * self.scale - self.mean) / self.std