Spaces:
Running
Running
| import math | |
| import sys | |
| import cv2 | |
| import numpy as np | |
| class DetResizeForTest(object): | |
| def __init__(self, **kwargs): | |
| super(DetResizeForTest, self).__init__() | |
| self.resize_type = 0 | |
| self.keep_ratio = False | |
| if 'image_shape' in kwargs: | |
| self.image_shape = kwargs['image_shape'] | |
| self.resize_type = 1 | |
| if 'keep_ratio' in kwargs: | |
| self.keep_ratio = kwargs['keep_ratio'] | |
| elif 'limit_side_len' in kwargs: | |
| self.limit_side_len = kwargs['limit_side_len'] | |
| self.limit_type = kwargs.get('limit_type', 'min') | |
| elif 'resize_long' in kwargs: | |
| self.resize_type = 2 | |
| self.resize_long = kwargs.get('resize_long', 960) | |
| else: | |
| self.limit_side_len = 736 | |
| self.limit_type = 'min' | |
| def __call__(self, data): | |
| img = data['image'] | |
| if 'max_sile_len' in data: | |
| self.limit_side_len = data['max_sile_len'] | |
| src_h, src_w, _ = img.shape | |
| if sum([src_h, src_w]) < 64: | |
| img = self.image_padding(img) | |
| if self.resize_type == 0: | |
| # img, shape = self.resize_image_type0(img) | |
| img, [ratio_h, ratio_w] = self.resize_image_type0(img) | |
| elif self.resize_type == 2: | |
| img, [ratio_h, ratio_w] = self.resize_image_type2(img) | |
| else: | |
| # img, shape = self.resize_image_type1(img) | |
| img, [ratio_h, ratio_w] = self.resize_image_type1(img) | |
| data['image'] = img | |
| data['shape'] = np.array([src_h, src_w, ratio_h, ratio_w]) | |
| return data | |
| def image_padding(self, im, value=0): | |
| h, w, c = im.shape | |
| im_pad = np.zeros((max(32, h), max(32, w), c), np.uint8) + value | |
| im_pad[:h, :w, :] = im | |
| return im_pad | |
| def resize_image_type1(self, img): | |
| resize_h, resize_w = self.image_shape | |
| ori_h, ori_w = img.shape[:2] # (h, w, c) | |
| if self.keep_ratio is True: | |
| resize_w = ori_w * resize_h / ori_h | |
| N = math.ceil(resize_w / 32) | |
| resize_w = N * 32 | |
| ratio_h = float(resize_h) / ori_h | |
| ratio_w = float(resize_w) / ori_w | |
| img = cv2.resize(img, (int(resize_w), int(resize_h))) | |
| # return img, np.array([ori_h, ori_w]) | |
| return img, [ratio_h, ratio_w] | |
| def resize_image_type0(self, img): | |
| """ | |
| resize image to a size multiple of 32 which is required by the network | |
| args: | |
| img(array): array with shape [h, w, c] | |
| return(tuple): | |
| img, (ratio_h, ratio_w) | |
| """ | |
| limit_side_len = self.limit_side_len | |
| h, w, c = img.shape | |
| # limit the max side | |
| if self.limit_type == 'max': | |
| if max(h, w) > limit_side_len: | |
| if h > w: | |
| ratio = float(limit_side_len) / h | |
| else: | |
| ratio = float(limit_side_len) / w | |
| else: | |
| ratio = 1.0 | |
| elif self.limit_type == 'min': | |
| if min(h, w) < limit_side_len: | |
| if h < w: | |
| ratio = float(limit_side_len) / h | |
| else: | |
| ratio = float(limit_side_len) / w | |
| else: | |
| ratio = 1.0 | |
| elif self.limit_type == 'resize_long': | |
| ratio = float(limit_side_len) / max(h, w) | |
| else: | |
| raise Exception('not support limit type, image ') | |
| resize_h = int(h * ratio) | |
| resize_w = int(w * ratio) | |
| resize_h = max(int(round(resize_h / 32) * 32), 32) | |
| resize_w = max(int(round(resize_w / 32) * 32), 32) | |
| try: | |
| if int(resize_w) <= 0 or int(resize_h) <= 0: | |
| return None, (None, None) | |
| img = cv2.resize(img, (int(resize_w), int(resize_h))) | |
| except: | |
| print(img.shape, resize_w, resize_h) | |
| sys.exit(0) | |
| ratio_h = resize_h / float(h) | |
| ratio_w = resize_w / float(w) | |
| return img, [ratio_h, ratio_w] | |
| def resize_image_type2(self, img): | |
| h, w, _ = img.shape | |
| resize_w = w | |
| resize_h = h | |
| if resize_h > resize_w: | |
| ratio = float(self.resize_long) / resize_h | |
| else: | |
| ratio = float(self.resize_long) / resize_w | |
| resize_h = int(resize_h * ratio) | |
| resize_w = int(resize_w * ratio) | |
| max_stride = 128 | |
| resize_h = (resize_h + max_stride - 1) // max_stride * max_stride | |
| resize_w = (resize_w + max_stride - 1) // max_stride * max_stride | |
| img = cv2.resize(img, (int(resize_w), int(resize_h))) | |
| ratio_h = resize_h / float(h) | |
| ratio_w = resize_w / float(w) | |
| return img, [ratio_h, ratio_w] | |