import cv2 import numpy as np def resize_short(img, target_size): """ resize_short """ percent = float(target_size) / min(img.shape[0], img.shape[1]) resized_width = int(round(img.shape[1] * percent)) resized_height = int(round(img.shape[0] * percent)) resized = cv2.resize(img, (resized_width, resized_height)) return resized def crop_image(img, target_size, center): """ crop_image """ height, width = img.shape[:2] size = target_size if center == True: w_start = (width - size) / 2 h_start = (height - size) / 2 else: w_start = np.random.randint(0, width - size + 1) h_start = np.random.randint(0, height - size + 1) w_end = w_start + size h_end = h_start + size img = img[int(h_start):int(h_end), int(w_start):int(w_end), :] return img def preprocess(img): mean = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] img = resize_short(img, 224) img = crop_image(img, 224, True) # bgr-> rgb && hwc->chw img = img[:, :, ::-1].astype('float32').transpose((2, 0, 1)) / 255 img_mean = np.array(mean).reshape((3, 1, 1)) img_std = np.array(std).reshape((3, 1, 1)) img -= img_mean img /= img_std return img[np.newaxis, :]