def test_opencv_crop(): dummy_img = np.ones((600, 600, 3), dtype=np.uint8) dummy_box = [20, 20, 120, 20, 120, 40, 20, 40] cropped_img = warp_img(dummy_img, dummy_box) with pytest.raises(AssertionError): warp_img(dummy_img, []) with pytest.raises(AssertionError): warp_img(dummy_img, [20, 40, 40, 20]) assert math.isclose(cropped_img.shape[0], 20) assert math.isclose(cropped_img.shape[1], 100)
def __call__(self, results): if 'img_info' not in results: return results crop_flag = True box = [] for key in self.box_keys: if key not in results['img_info']: crop_flag = False break box.append(float(results['img_info'][key])) if not crop_flag: return results jitter_flag = np.random.random() > self.jitter_prob kwargs = dict(jitter_flag=jitter_flag, jitter_ratio_x=self.max_jitter_ratio_x, jitter_ratio_y=self.max_jitter_ratio_y) crop_img = warp_img(results['img'], box, **kwargs) results['img'] = crop_img results['img_shape'] = crop_img.shape return results
def parse_labelme_json(json_file, img_dir, out_dir, tasks, ignore_marker='###', recog_format='jsonl', warp_flag=False): invalid_res = [[], [], []] json_obj = mmcv.load(json_file) img_file = osp.basename(json_obj['imagePath']) img_full_path = osp.join(img_dir, img_file) img_width = json_obj['imageWidth'] img_height = json_obj['imageHeight'] if 'recog' in tasks: src_img = mmcv.imread(img_full_path) img_basename = osp.splitext(img_file)[0] sub_dir = osp.join(out_dir, 'crops', img_basename) mmcv.mkdir_or_exist(sub_dir) det_line_json_list = [] recog_crop_line_str_list = [] recog_warp_line_str_list = [] shape_info = json_obj['shapes'] idx = 0 annos = [] for box_info in shape_info: shape = box_info['shape_type'] if shape not in ['rectangle', 'polygon']: msg = 'Only \'rectangle\' and \'polygon\' boxes are supported. ' msg += f'Boxes with {shape} will be discarded.' warnings.warn(msg) return invalid_res poly = [] box_points = box_info['points'] for point in box_points: poly.extend([int(x) for x in point]) x_list = poly[0::2] y_list = poly[1::2] quad = [] if shape == 'rectangle': warp_flag = False quad = [ poly[0], poly[1], poly[2], poly[1], poly[2], poly[3], poly[0], poly[3] ] else: if len(poly) < 8 or len(poly) % 2 != 0: msg = f'Invalid polygon {poly}. ' msg += 'The polygon is expected to have 8 or more than 8 ' msg += 'even number of coordinates in MMOCR.' warnings.warn(msg) return invalid_res if len(poly) == 8: quad = poly else: warp_flag = False x_min, x_max, y_min, y_max = min(x_list), max(x_list), min( y_list), max(y_list) quad = [x_min, y_min, x_max, y_min, x_max, y_max, x_min, y_max] text_label = box_info['label'] # for textdet anno = {} anno['iscrowd'] = 0 if text_label != ignore_marker else 1 anno['category_id'] = 1 w = max(x_list) - min(x_list) h = max(y_list) - min(y_list) anno['bbox'] = [min(x_list), min(y_list), w, h] if shape == 'rectangle': anno['segmentation'] = [quad] else: anno['segmentation'] = [poly] anno['text'] = text_label annos.append(anno) # for textrecog if 'recog' in tasks: if text_label == ignore_marker or len(text_label) == 0: continue cropped_img = crop_img(src_img, quad) img_path_cropped_img = osp.join(sub_dir, f'crop_{idx}.jpg') mmcv.imwrite(cropped_img, img_path_cropped_img) if recog_format == 'txt': recog_crop_line_str_list.append( f'{img_path_cropped_img} {text_label}') elif recog_format == 'jsonl': recog_crop_line_str_list.append( json.dumps({ 'filename': img_path_cropped_img, 'text': text_label })) else: raise NotImplementedError if warp_flag: warpped_img = warp_img(src_img, quad) img_path_warpped_img = osp.join(sub_dir, f'warp_{idx}.jpg') mmcv.imwrite(warpped_img, img_path_warpped_img) if recog_format == 'txt': recog_warp_line_str_list.append( f'{img_path_warpped_img} {text_label}') elif recog_format == 'jsonl': recog_warp_line_str_list.append( json.dumps({ 'filename': img_path_cropped_img, 'text': text_label })) idx += 1 line_json = { 'file_name': img_file, 'height': img_height, 'width': img_width, 'annotations': annos } det_line_json_list.append(json.dumps(line_json, ensure_ascii=False)) return [ det_line_json_list, recog_crop_line_str_list, recog_warp_line_str_list ]