Shortcuts

Source code for mmocr.datasets.pipelines.transforms

# Copyright (c) OpenMMLab. All rights reserved.
import math

import cv2
import mmcv
import numpy as np
import torchvision.transforms as transforms
from mmdet.core import BitmapMasks, PolygonMasks
from mmdet.datasets.builder import PIPELINES
from mmdet.datasets.pipelines.transforms import Resize
from PIL import Image
from shapely.geometry import Polygon as plg

import mmocr.core.evaluation.utils as eval_utils
from mmocr.utils import check_argument


[docs]@PIPELINES.register_module() class RandomCropInstances: """Randomly crop images and make sure to contain text instances. Args: target_size (tuple or int): (height, width) positive_sample_ratio (float): The probability of sampling regions that go through positive regions. """ def __init__( self, target_size, instance_key, mask_type='inx0', # 'inx0' or 'union_all' positive_sample_ratio=5.0 / 8.0): assert mask_type in ['inx0', 'union_all'] self.mask_type = mask_type self.instance_key = instance_key self.positive_sample_ratio = positive_sample_ratio self.target_size = target_size if (target_size is None or isinstance( target_size, tuple)) else (target_size, target_size) def sample_offset(self, img_gt, img_size): h, w = img_size t_h, t_w = self.target_size # target size is bigger than origin size t_h = t_h if t_h < h else h t_w = t_w if t_w < w else w if (img_gt is not None and np.random.random_sample() < self.positive_sample_ratio and np.max(img_gt) > 0): # make sure to crop the positive region # the minimum top left to crop positive region (h,w) tl = np.min(np.where(img_gt > 0), axis=1) - (t_h, t_w) tl[tl < 0] = 0 # the maximum top left to crop positive region br = np.max(np.where(img_gt > 0), axis=1) - (t_h, t_w) br[br < 0] = 0 # if br is too big so that crop the outside region of img br[0] = min(br[0], h - t_h) br[1] = min(br[1], w - t_w) # h = np.random.randint(tl[0], br[0]) if tl[0] < br[0] else 0 w = np.random.randint(tl[1], br[1]) if tl[1] < br[1] else 0 else: # make sure not to crop outside of img h = np.random.randint(0, h - t_h) if h - t_h > 0 else 0 w = np.random.randint(0, w - t_w) if w - t_w > 0 else 0 return (h, w) @staticmethod def crop_img(img, offset, target_size): h, w = img.shape[:2] br = np.min( np.stack((np.array(offset) + np.array(target_size), np.array( (h, w)))), axis=0) return img[offset[0]:br[0], offset[1]:br[1]], np.array( [offset[1], offset[0], br[1], br[0]]) def crop_bboxes(self, bboxes, canvas_bbox): kept_bboxes = [] kept_inx = [] canvas_poly = eval_utils.box2polygon(canvas_bbox) tl = canvas_bbox[0:2] for idx, bbox in enumerate(bboxes): poly = eval_utils.box2polygon(bbox) area, inters = eval_utils.poly_intersection( poly, canvas_poly, return_poly=True) if area == 0: continue xmin, ymin, xmax, ymax = inters.bounds kept_bboxes += [ np.array( [xmin - tl[0], ymin - tl[1], xmax - tl[0], ymax - tl[1]], dtype=np.float32) ] kept_inx += [idx] if len(kept_inx) == 0: return np.array([]).astype(np.float32).reshape(0, 4), kept_inx return np.stack(kept_bboxes), kept_inx @staticmethod def generate_mask(gt_mask, type): if type == 'inx0': return gt_mask.masks[0] if type == 'union_all': mask = gt_mask.masks[0].copy() for idx in range(1, len(gt_mask.masks)): mask = np.logical_or(mask, gt_mask.masks[idx]) return mask raise NotImplementedError def __call__(self, results): gt_mask = results[self.instance_key] mask = None if len(gt_mask.masks) > 0: mask = self.generate_mask(gt_mask, self.mask_type) results['crop_offset'] = self.sample_offset(mask, results['img'].shape[:2]) # crop img. bbox = [x1,y1,x2,y2] img, bbox = self.crop_img(results['img'], results['crop_offset'], self.target_size) results['img'] = img img_shape = img.shape results['img_shape'] = img_shape # crop masks for key in results.get('mask_fields', []): results[key] = results[key].crop(bbox) # for mask rcnn for key in results.get('bbox_fields', []): results[key], kept_inx = self.crop_bboxes(results[key], bbox) if key == 'gt_bboxes': # ignore gt_labels accordingly if 'gt_labels' in results: ori_labels = results['gt_labels'] ori_inst_num = len(ori_labels) results['gt_labels'] = [ ori_labels[idx] for idx in range(ori_inst_num) if idx in kept_inx ] # ignore g_masks accordingly if 'gt_masks' in results: ori_mask = results['gt_masks'].masks kept_mask = [ ori_mask[idx] for idx in range(ori_inst_num) if idx in kept_inx ] target_h, target_w = bbox[3] - bbox[1], bbox[2] - bbox[0] if len(kept_inx) > 0: kept_mask = np.stack(kept_mask) else: kept_mask = np.empty((0, target_h, target_w), dtype=np.float32) results['gt_masks'] = BitmapMasks(kept_mask, target_h, target_w) return results def __repr__(self): repr_str = self.__class__.__name__ return repr_str
[docs]@PIPELINES.register_module() class RandomRotateTextDet: """Randomly rotate images.""" def __init__(self, rotate_ratio=1.0, max_angle=10): self.rotate_ratio = rotate_ratio self.max_angle = max_angle @staticmethod def sample_angle(max_angle): angle = np.random.random_sample() * 2 * max_angle - max_angle return angle @staticmethod def rotate_img(img, angle): h, w = img.shape[:2] rotation_matrix = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1) img_target = cv2.warpAffine( img, rotation_matrix, (w, h), flags=cv2.INTER_NEAREST) assert img_target.shape == img.shape return img_target def __call__(self, results): if np.random.random_sample() < self.rotate_ratio: # rotate imgs results['rotated_angle'] = self.sample_angle(self.max_angle) img = self.rotate_img(results['img'], results['rotated_angle']) results['img'] = img img_shape = img.shape results['img_shape'] = img_shape # rotate masks for key in results.get('mask_fields', []): masks = results[key].masks mask_list = [] for m in masks: rotated_m = self.rotate_img(m, results['rotated_angle']) mask_list.append(rotated_m) results[key] = BitmapMasks(mask_list, *(img_shape[:2])) return results def __repr__(self): repr_str = self.__class__.__name__ return repr_str
[docs]@PIPELINES.register_module() class ColorJitter: """An interface for torch color jitter so that it can be invoked in mmdetection pipeline.""" def __init__(self, **kwargs): self.transform = transforms.ColorJitter(**kwargs) def __call__(self, results): # img is bgr img = results['img'][..., ::-1] img = Image.fromarray(img) img = self.transform(img) img = np.asarray(img) img = img[..., ::-1] results['img'] = img return results def __repr__(self): repr_str = self.__class__.__name__ return repr_str
[docs]@PIPELINES.register_module() class ScaleAspectJitter(Resize): """Resize image and segmentation mask encoded by coordinates. Allowed resize types are `around_min_img_scale`, `long_short_bound`, and `indep_sample_in_range`. """ def __init__(self, img_scale=None, multiscale_mode='range', ratio_range=None, keep_ratio=False, resize_type='around_min_img_scale', aspect_ratio_range=None, long_size_bound=None, short_size_bound=None, scale_range=None): super().__init__( img_scale=img_scale, multiscale_mode=multiscale_mode, ratio_range=ratio_range, keep_ratio=keep_ratio) assert not keep_ratio assert resize_type in [ 'around_min_img_scale', 'long_short_bound', 'indep_sample_in_range' ] self.resize_type = resize_type if resize_type == 'indep_sample_in_range': assert ratio_range is None assert aspect_ratio_range is None assert short_size_bound is None assert long_size_bound is None assert scale_range is not None else: assert scale_range is None assert isinstance(ratio_range, tuple) assert isinstance(aspect_ratio_range, tuple) assert check_argument.equal_len(ratio_range, aspect_ratio_range) if resize_type in ['long_short_bound']: assert short_size_bound is not None assert long_size_bound is not None self.aspect_ratio_range = aspect_ratio_range self.long_size_bound = long_size_bound self.short_size_bound = short_size_bound self.scale_range = scale_range @staticmethod def sample_from_range(range): assert len(range) == 2 min_value, max_value = min(range), max(range) value = np.random.random_sample() * (max_value - min_value) + min_value return value def _random_scale(self, results): if self.resize_type == 'indep_sample_in_range': w = self.sample_from_range(self.scale_range) h = self.sample_from_range(self.scale_range) results['scale'] = (int(w), int(h)) # (w,h) results['scale_idx'] = None return h, w = results['img'].shape[0:2] if self.resize_type == 'long_short_bound': scale1 = 1 if max(h, w) > self.long_size_bound: scale1 = self.long_size_bound / max(h, w) scale2 = self.sample_from_range(self.ratio_range) scale = scale1 * scale2 if min(h, w) * scale <= self.short_size_bound: scale = (self.short_size_bound + 10) * 1.0 / min(h, w) elif self.resize_type == 'around_min_img_scale': short_size = min(self.img_scale[0]) ratio = self.sample_from_range(self.ratio_range) scale = (ratio * short_size) / min(h, w) else: raise NotImplementedError aspect = self.sample_from_range(self.aspect_ratio_range) h_scale = scale * math.sqrt(aspect) w_scale = scale / math.sqrt(aspect) results['scale'] = (int(w * w_scale), int(h * h_scale)) # (w,h) results['scale_idx'] = None
@PIPELINES.register_module() class AffineJitter: """An interface for torchvision random affine so that it can be invoked in mmdet pipeline.""" def __init__(self, degrees=4, translate=(0.02, 0.04), scale=(0.9, 1.1), shear=None, resample=False, fillcolor=0): self.transform = transforms.RandomAffine( degrees=degrees, translate=translate, scale=scale, shear=shear, resample=resample, fillcolor=fillcolor) def __call__(self, results): # img is bgr img = results['img'][..., ::-1] img = Image.fromarray(img) img = self.transform(img) img = np.asarray(img) img = img[..., ::-1] results['img'] = img return results def __repr__(self): repr_str = self.__class__.__name__ return repr_str
[docs]@PIPELINES.register_module() class RandomCropPolyInstances: """Randomly crop images and make sure to contain at least one intact instance.""" def __init__(self, instance_key='gt_masks', crop_ratio=5.0 / 8.0, min_side_ratio=0.4): super().__init__() self.instance_key = instance_key self.crop_ratio = crop_ratio self.min_side_ratio = min_side_ratio def sample_valid_start_end(self, valid_array, min_len, max_start, min_end): assert isinstance(min_len, int) assert len(valid_array) > min_len start_array = valid_array.copy() max_start = min(len(start_array) - min_len, max_start) start_array[max_start:] = 0 start_array[0] = 1 diff_array = np.hstack([0, start_array]) - np.hstack([start_array, 0]) region_starts = np.where(diff_array < 0)[0] region_ends = np.where(diff_array > 0)[0] region_ind = np.random.randint(0, len(region_starts)) start = np.random.randint(region_starts[region_ind], region_ends[region_ind]) end_array = valid_array.copy() min_end = max(start + min_len, min_end) end_array[:min_end] = 0 end_array[-1] = 1 diff_array = np.hstack([0, end_array]) - np.hstack([end_array, 0]) region_starts = np.where(diff_array < 0)[0] region_ends = np.where(diff_array > 0)[0] region_ind = np.random.randint(0, len(region_starts)) end = np.random.randint(region_starts[region_ind], region_ends[region_ind]) return start, end
[docs] def sample_crop_box(self, img_size, results): """Generate crop box and make sure not to crop the polygon instances. Args: img_size (tuple(int)): The image size (h, w). results (dict): The results dict. """ assert isinstance(img_size, tuple) h, w = img_size[:2] key_masks = results[self.instance_key].masks x_valid_array = np.ones(w, dtype=np.int32) y_valid_array = np.ones(h, dtype=np.int32) selected_mask = key_masks[np.random.randint(0, len(key_masks))] selected_mask = selected_mask[0].reshape((-1, 2)).astype(np.int32) max_x_start = max(np.min(selected_mask[:, 0]) - 2, 0) min_x_end = min(np.max(selected_mask[:, 0]) + 3, w - 1) max_y_start = max(np.min(selected_mask[:, 1]) - 2, 0) min_y_end = min(np.max(selected_mask[:, 1]) + 3, h - 1) for key in results.get('mask_fields', []): if len(results[key].masks) == 0: continue masks = results[key].masks for mask in masks: assert len(mask) == 1 mask = mask[0].reshape((-1, 2)).astype(np.int32) clip_x = np.clip(mask[:, 0], 0, w - 1) clip_y = np.clip(mask[:, 1], 0, h - 1) min_x, max_x = np.min(clip_x), np.max(clip_x) min_y, max_y = np.min(clip_y), np.max(clip_y) x_valid_array[min_x - 2:max_x + 3] = 0 y_valid_array[min_y - 2:max_y + 3] = 0 min_w = int(w * self.min_side_ratio) min_h = int(h * self.min_side_ratio) x1, x2 = self.sample_valid_start_end(x_valid_array, min_w, max_x_start, min_x_end) y1, y2 = self.sample_valid_start_end(y_valid_array, min_h, max_y_start, min_y_end) return np.array([x1, y1, x2, y2])
def crop_img(self, img, bbox): assert img.ndim == 3 h, w, _ = img.shape assert 0 <= bbox[1] < bbox[3] <= h assert 0 <= bbox[0] < bbox[2] <= w return img[bbox[1]:bbox[3], bbox[0]:bbox[2]] def __call__(self, results): if len(results[self.instance_key].masks) < 1: return results if np.random.random_sample() < self.crop_ratio: crop_box = self.sample_crop_box(results['img'].shape, results) results['crop_region'] = crop_box img = self.crop_img(results['img'], crop_box) results['img'] = img results['img_shape'] = img.shape # crop and filter masks x1, y1, x2, y2 = crop_box w = max(x2 - x1, 1) h = max(y2 - y1, 1) labels = results['gt_labels'] valid_labels = [] for key in results.get('mask_fields', []): if len(results[key].masks) == 0: continue results[key] = results[key].crop(crop_box) # filter out polygons beyond crop box. masks = results[key].masks valid_masks_list = [] for ind, mask in enumerate(masks): assert len(mask) == 1 polygon = mask[0].reshape((-1, 2)) if (polygon[:, 0] > -4).all() and (polygon[:, 0] < w + 4).all() and ( polygon[:, 1] > -4).all() and (polygon[:, 1] < h + 4).all(): mask[0][::2] = np.clip(mask[0][::2], 0, w) mask[0][1::2] = np.clip(mask[0][1::2], 0, h) if key == self.instance_key: valid_labels.append(labels[ind]) valid_masks_list.append(mask) results[key] = PolygonMasks(valid_masks_list, h, w) results['gt_labels'] = np.array(valid_labels) return results def __repr__(self): repr_str = self.__class__.__name__ return repr_str
@PIPELINES.register_module() class RandomRotatePolyInstances: def __init__(self, rotate_ratio=0.5, max_angle=10, pad_with_fixed_color=False, pad_value=(0, 0, 0)): """Randomly rotate images and polygon masks. Args: rotate_ratio (float): The ratio of samples to operate rotation. max_angle (int): The maximum rotation angle. pad_with_fixed_color (bool): The flag for whether to pad rotated image with fixed value. If set to False, the rotated image will be padded onto cropped image. pad_value (tuple(int)): The color value for padding rotated image. """ self.rotate_ratio = rotate_ratio self.max_angle = max_angle self.pad_with_fixed_color = pad_with_fixed_color self.pad_value = pad_value def rotate(self, center, points, theta, center_shift=(0, 0)): # rotate points. (center_x, center_y) = center center_y = -center_y x, y = points[::2], points[1::2] y = -y theta = theta / 180 * math.pi cos = math.cos(theta) sin = math.sin(theta) x = (x - center_x) y = (y - center_y) _x = center_x + x * cos - y * sin + center_shift[0] _y = -(center_y + x * sin + y * cos) + center_shift[1] points[::2], points[1::2] = _x, _y return points def cal_canvas_size(self, ori_size, degree): assert isinstance(ori_size, tuple) angle = degree * math.pi / 180.0 h, w = ori_size[:2] cos = math.cos(angle) sin = math.sin(angle) canvas_h = int(w * math.fabs(sin) + h * math.fabs(cos)) canvas_w = int(w * math.fabs(cos) + h * math.fabs(sin)) canvas_size = (canvas_h, canvas_w) return canvas_size def sample_angle(self, max_angle): angle = np.random.random_sample() * 2 * max_angle - max_angle return angle def rotate_img(self, img, angle, canvas_size): h, w = img.shape[:2] rotation_matrix = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1) rotation_matrix[0, 2] += int((canvas_size[1] - w) / 2) rotation_matrix[1, 2] += int((canvas_size[0] - h) / 2) if self.pad_with_fixed_color: target_img = cv2.warpAffine( img, rotation_matrix, (canvas_size[1], canvas_size[0]), flags=cv2.INTER_NEAREST, borderValue=self.pad_value) else: mask = np.zeros_like(img) (h_ind, w_ind) = (np.random.randint(0, h * 7 // 8), np.random.randint(0, w * 7 // 8)) img_cut = img[h_ind:(h_ind + h // 9), w_ind:(w_ind + w // 9)] img_cut = mmcv.imresize(img_cut, (canvas_size[1], canvas_size[0])) mask = cv2.warpAffine( mask, rotation_matrix, (canvas_size[1], canvas_size[0]), borderValue=[1, 1, 1]) target_img = cv2.warpAffine( img, rotation_matrix, (canvas_size[1], canvas_size[0]), borderValue=[0, 0, 0]) target_img = target_img + img_cut * mask return target_img def __call__(self, results): if np.random.random_sample() < self.rotate_ratio: img = results['img'] h, w = img.shape[:2] angle = self.sample_angle(self.max_angle) canvas_size = self.cal_canvas_size((h, w), angle) center_shift = (int( (canvas_size[1] - w) / 2), int((canvas_size[0] - h) / 2)) # rotate image results['rotated_poly_angle'] = angle img = self.rotate_img(img, angle, canvas_size) results['img'] = img img_shape = img.shape results['img_shape'] = img_shape # rotate polygons for key in results.get('mask_fields', []): if len(results[key].masks) == 0: continue masks = results[key].masks rotated_masks = [] for mask in masks: rotated_mask = self.rotate((w / 2, h / 2), mask[0], angle, center_shift) rotated_masks.append([rotated_mask]) results[key] = PolygonMasks(rotated_masks, *(img_shape[:2])) return results def __repr__(self): repr_str = self.__class__.__name__ return repr_str @PIPELINES.register_module() class SquareResizePad: def __init__(self, target_size, pad_ratio=0.6, pad_with_fixed_color=False, pad_value=(0, 0, 0)): """Resize or pad images to be square shape. Args: target_size (int): The target size of square shaped image. pad_with_fixed_color (bool): The flag for whether to pad rotated image with fixed value. If set to False, the rescales image will be padded onto cropped image. pad_value (tuple(int)): The color value for padding rotated image. """ assert isinstance(target_size, int) assert isinstance(pad_ratio, float) assert isinstance(pad_with_fixed_color, bool) assert isinstance(pad_value, tuple) self.target_size = target_size self.pad_ratio = pad_ratio self.pad_with_fixed_color = pad_with_fixed_color self.pad_value = pad_value def resize_img(self, img, keep_ratio=True): h, w, _ = img.shape if keep_ratio: t_h = self.target_size if h >= w else int(h * self.target_size / w) t_w = self.target_size if h <= w else int(w * self.target_size / h) else: t_h = t_w = self.target_size img = mmcv.imresize(img, (t_w, t_h)) return img, (t_h, t_w) def square_pad(self, img): h, w = img.shape[:2] if h == w: return img, (0, 0) pad_size = max(h, w) if self.pad_with_fixed_color: expand_img = np.ones((pad_size, pad_size, 3), dtype=np.uint8) expand_img[:] = self.pad_value else: (h_ind, w_ind) = (np.random.randint(0, h * 7 // 8), np.random.randint(0, w * 7 // 8)) img_cut = img[h_ind:(h_ind + h // 9), w_ind:(w_ind + w // 9)] expand_img = mmcv.imresize(img_cut, (pad_size, pad_size)) if h > w: y0, x0 = 0, (h - w) // 2 else: y0, x0 = (w - h) // 2, 0 expand_img[y0:y0 + h, x0:x0 + w] = img offset = (x0, y0) return expand_img, offset def square_pad_mask(self, points, offset): x0, y0 = offset pad_points = points.copy() pad_points[::2] = pad_points[::2] + x0 pad_points[1::2] = pad_points[1::2] + y0 return pad_points def __call__(self, results): img = results['img'] if np.random.random_sample() < self.pad_ratio: img, out_size = self.resize_img(img, keep_ratio=True) img, offset = self.square_pad(img) else: img, out_size = self.resize_img(img, keep_ratio=False) offset = (0, 0) results['img'] = img results['img_shape'] = img.shape for key in results.get('mask_fields', []): if len(results[key].masks) == 0: continue results[key] = results[key].resize(out_size) masks = results[key].masks processed_masks = [] for mask in masks: square_pad_mask = self.square_pad_mask(mask[0], offset) processed_masks.append([square_pad_mask]) results[key] = PolygonMasks(processed_masks, *(img.shape[:2])) return results def __repr__(self): repr_str = self.__class__.__name__ return repr_str @PIPELINES.register_module() class RandomScaling: def __init__(self, size=800, scale=(3. / 4, 5. / 2)): """Random scale the image while keeping aspect. Args: size (int) : Base size before scaling. scale (tuple(float)) : The range of scaling. """ assert isinstance(size, int) assert isinstance(scale, float) or isinstance(scale, tuple) self.size = size self.scale = scale if isinstance(scale, tuple) \ else (1 - scale, 1 + scale) def __call__(self, results): image = results['img'] h, w, _ = results['img_shape'] aspect_ratio = np.random.uniform(min(self.scale), max(self.scale)) scales = self.size * 1.0 / max(h, w) * aspect_ratio scales = np.array([scales, scales]) out_size = (int(h * scales[1]), int(w * scales[0])) image = mmcv.imresize(image, out_size[::-1]) results['img'] = image results['img_shape'] = image.shape for key in results.get('mask_fields', []): if len(results[key].masks) == 0: continue results[key] = results[key].resize(out_size) return results @PIPELINES.register_module() class RandomCropFlip: def __init__(self, pad_ratio=0.1, crop_ratio=0.5, iter_num=1, min_area_ratio=0.2): """Random crop and flip a patch of the image. Args: crop_ratio (float): The ratio of cropping. iter_num (int): Number of operations. min_area_ratio (float): Minimal area ratio between cropped patch and original image. """ assert isinstance(crop_ratio, float) assert isinstance(iter_num, int) assert isinstance(min_area_ratio, float) self.pad_ratio = pad_ratio self.epsilon = 1e-2 self.crop_ratio = crop_ratio self.iter_num = iter_num self.min_area_ratio = min_area_ratio def __call__(self, results): for i in range(self.iter_num): results = self.random_crop_flip(results) return results def random_crop_flip(self, results): image = results['img'] polygons = results['gt_masks'].masks ignore_polygons = results['gt_masks_ignore'].masks all_polygons = polygons + ignore_polygons if len(polygons) == 0: return results if np.random.random() >= self.crop_ratio: return results h, w, _ = results['img_shape'] area = h * w pad_h = int(h * self.pad_ratio) pad_w = int(w * self.pad_ratio) h_axis, w_axis = self.generate_crop_target(image, all_polygons, pad_h, pad_w) if len(h_axis) == 0 or len(w_axis) == 0: return results attempt = 0 while attempt < 10: attempt += 1 polys_keep = [] polys_new = [] ign_polys_keep = [] ign_polys_new = [] xx = np.random.choice(w_axis, size=2) xmin = np.min(xx) - pad_w xmax = np.max(xx) - pad_w xmin = np.clip(xmin, 0, w - 1) xmax = np.clip(xmax, 0, w - 1) yy = np.random.choice(h_axis, size=2) ymin = np.min(yy) - pad_h ymax = np.max(yy) - pad_h ymin = np.clip(ymin, 0, h - 1) ymax = np.clip(ymax, 0, h - 1) if (xmax - xmin) * (ymax - ymin) < area * self.min_area_ratio: # area too small continue pts = np.stack([[xmin, xmax, xmax, xmin], [ymin, ymin, ymax, ymax]]).T.astype(np.int32) pp = plg(pts) fail_flag = False for polygon in polygons: ppi = plg(polygon[0].reshape(-1, 2)) ppiou = eval_utils.poly_intersection(ppi, pp) if np.abs(ppiou - float(ppi.area)) > self.epsilon and \ np.abs(ppiou) > self.epsilon: fail_flag = True break elif np.abs(ppiou - float(ppi.area)) < self.epsilon: polys_new.append(polygon) else: polys_keep.append(polygon) for polygon in ignore_polygons: ppi = plg(polygon[0].reshape(-1, 2)) ppiou = eval_utils.poly_intersection(ppi, pp) if np.abs(ppiou - float(ppi.area)) > self.epsilon and \ np.abs(ppiou) > self.epsilon: fail_flag = True break elif np.abs(ppiou - float(ppi.area)) < self.epsilon: ign_polys_new.append(polygon) else: ign_polys_keep.append(polygon) if fail_flag: continue else: break cropped = image[ymin:ymax, xmin:xmax, :] select_type = np.random.randint(3) if select_type == 0: img = np.ascontiguousarray(cropped[:, ::-1]) elif select_type == 1: img = np.ascontiguousarray(cropped[::-1, :]) else: img = np.ascontiguousarray(cropped[::-1, ::-1]) image[ymin:ymax, xmin:xmax, :] = img results['img'] = image if len(polys_new) + len(ign_polys_new) != 0: height, width, _ = cropped.shape if select_type == 0: for idx, polygon in enumerate(polys_new): poly = polygon[0].reshape(-1, 2) poly[:, 0] = width - poly[:, 0] + 2 * xmin polys_new[idx] = [poly.reshape(-1, )] for idx, polygon in enumerate(ign_polys_new): poly = polygon[0].reshape(-1, 2) poly[:, 0] = width - poly[:, 0] + 2 * xmin ign_polys_new[idx] = [poly.reshape(-1, )] elif select_type == 1: for idx, polygon in enumerate(polys_new): poly = polygon[0].reshape(-1, 2) poly[:, 1] = height - poly[:, 1] + 2 * ymin polys_new[idx] = [poly.reshape(-1, )] for idx, polygon in enumerate(ign_polys_new): poly = polygon[0].reshape(-1, 2) poly[:, 1] = height - poly[:, 1] + 2 * ymin ign_polys_new[idx] = [poly.reshape(-1, )] else: for idx, polygon in enumerate(polys_new): poly = polygon[0].reshape(-1, 2) poly[:, 0] = width - poly[:, 0] + 2 * xmin poly[:, 1] = height - poly[:, 1] + 2 * ymin polys_new[idx] = [poly.reshape(-1, )] for idx, polygon in enumerate(ign_polys_new): poly = polygon[0].reshape(-1, 2) poly[:, 0] = width - poly[:, 0] + 2 * xmin poly[:, 1] = height - poly[:, 1] + 2 * ymin ign_polys_new[idx] = [poly.reshape(-1, )] polygons = polys_keep + polys_new ignore_polygons = ign_polys_keep + ign_polys_new results['gt_masks'] = PolygonMasks(polygons, *(image.shape[:2])) results['gt_masks_ignore'] = PolygonMasks(ignore_polygons, *(image.shape[:2])) return results def generate_crop_target(self, image, all_polys, pad_h, pad_w): """Generate crop target and make sure not to crop the polygon instances. Args: image (ndarray): The image waited to be crop. all_polys (list[list[ndarray]]): All polygons including ground truth polygons and ground truth ignored polygons. pad_h (int): Padding length of height. pad_w (int): Padding length of width. Returns: h_axis (ndarray): Vertical cropping range. w_axis (ndarray): Horizontal cropping range. """ h, w, _ = image.shape h_array = np.zeros((h + pad_h * 2), dtype=np.int32) w_array = np.zeros((w + pad_w * 2), dtype=np.int32) text_polys = [] for polygon in all_polys: rect = cv2.minAreaRect(polygon[0].astype(np.int32).reshape(-1, 2)) box = cv2.boxPoints(rect) box = np.int0(box) text_polys.append([box[0], box[1], box[2], box[3]]) polys = np.array(text_polys, dtype=np.int32) for poly in polys: poly = np.round(poly, decimals=0).astype(np.int32) minx = np.min(poly[:, 0]) maxx = np.max(poly[:, 0]) w_array[minx + pad_w:maxx + pad_w] = 1 miny = np.min(poly[:, 1]) maxy = np.max(poly[:, 1]) h_array[miny + pad_h:maxy + pad_h] = 1 h_axis = np.where(h_array == 0)[0] w_axis = np.where(w_array == 0)[0] return h_axis, w_axis
[docs]@PIPELINES.register_module() class PyramidRescale: """Resize the image to the base shape, downsample it with gaussian pyramid, and rescale it back to original size. Adapted from https://github.com/FangShancheng/ABINet. Args: factor (int): The decay factor from base size, or the number of downsampling operations from the base layer. base_shape (tuple(int)): The shape of the base layer of the pyramid. randomize_factor (bool): If True, the final factor would be a random integer in [0, factor]. :Required Keys: - | ``img`` (ndarray): The input image. :Affected Keys: :Modified: - | ``img`` (ndarray): The modified image. """ def __init__(self, factor=4, base_shape=(128, 512), randomize_factor=True): assert isinstance(factor, int) assert isinstance(base_shape, list) or isinstance(base_shape, tuple) assert len(base_shape) == 2 assert isinstance(randomize_factor, bool) self.factor = factor if not randomize_factor else np.random.randint( 0, factor + 1) self.base_w, self.base_h = base_shape def __call__(self, results): assert 'img' in results if self.factor == 0: return results img = results['img'] src_h, src_w = img.shape[:2] scale_img = mmcv.imresize(img, (self.base_w, self.base_h)) for _ in range(self.factor): scale_img = cv2.pyrDown(scale_img) scale_img = mmcv.imresize(scale_img, (src_w, src_h)) results['img'] = scale_img return results def __repr__(self): repr_str = self.__class__.__name__ repr_str += f'(factor={self.factor}, ' repr_str += f'basew={self.basew}, baseh={self.baseh})' return repr_str
Read the Docs v: v0.4.0
Versions
latest
stable
v0.4.0
v0.3.0
v0.2.1
v0.2.0
v0.1.0
Downloads
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.