Spaces:
Sleeping
Sleeping
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import argparse | |
| import math | |
| import os.path as osp | |
| import mmcv | |
| import mmengine | |
| import numpy as np | |
| from mmocr.utils import dump_ocr_data | |
| def parse_args(): | |
| parser = argparse.ArgumentParser( | |
| description='Generate training, validation and test set of IMGUR ') | |
| parser.add_argument('root_path', help='Root dir path of IMGUR') | |
| args = parser.parse_args() | |
| return args | |
| def collect_imgur_info(root_path, annotation_filename, print_every=1000): | |
| annotation_path = osp.join(root_path, 'annotations', annotation_filename) | |
| if not osp.exists(annotation_path): | |
| raise Exception( | |
| f'{annotation_path} not exists, please check and try again.') | |
| annotation = mmengine.load(annotation_path) | |
| images = annotation['index_to_ann_map'].keys() | |
| img_infos = [] | |
| for i, img_name in enumerate(images): | |
| if i >= 0 and i % print_every == 0: | |
| print(f'{i}/{len(images)}') | |
| img_path = osp.join(root_path, 'imgs', img_name + '.jpg') | |
| # Skip not exist images | |
| if not osp.exists(img_path): | |
| continue | |
| img = mmcv.imread(img_path, 'unchanged') | |
| # Skip broken images | |
| if img is None: | |
| continue | |
| img_info = dict( | |
| file_name=img_name + '.jpg', | |
| height=img.shape[0], | |
| width=img.shape[1]) | |
| anno_info = [] | |
| for ann_id in annotation['index_to_ann_map'][img_name]: | |
| ann = annotation['ann_id'][ann_id] | |
| # The original annotation is oriented rects [x, y, w, h, a] | |
| box = np.fromstring( | |
| ann['bounding_box'][1:-2], sep=',', dtype=float) | |
| quadrilateral = convert_oriented_box(box) | |
| xs, ys = quadrilateral[::2], quadrilateral[1::2] | |
| x = max(0, math.floor(min(xs))) | |
| y = max(0, math.floor(min(ys))) | |
| w = math.floor(max(xs)) - x | |
| h = math.floor(max(ys)) - y | |
| bbox = [x, y, w, h] | |
| segmentation = quadrilateral | |
| anno = dict( | |
| iscrowd=0, | |
| category_id=1, | |
| bbox=bbox, | |
| area=w * h, | |
| segmentation=[segmentation]) | |
| anno_info.append(anno) | |
| img_info.update(anno_info=anno_info) | |
| img_infos.append(img_info) | |
| return img_infos | |
| def convert_oriented_box(box): | |
| x_ctr, y_ctr, width, height, angle = box[:5] | |
| angle = -angle * math.pi / 180 | |
| tl_x, tl_y, br_x, br_y = -width / 2, -height / 2, width / 2, height / 2 | |
| rect = np.array([[tl_x, br_x, br_x, tl_x], [tl_y, tl_y, br_y, br_y]]) | |
| R = np.array([[np.cos(angle), -np.sin(angle)], | |
| [np.sin(angle), np.cos(angle)]]) | |
| poly = R.dot(rect) | |
| x0, x1, x2, x3 = poly[0, :4] + x_ctr | |
| y0, y1, y2, y3 = poly[1, :4] + y_ctr | |
| poly = np.array([x0, y0, x1, y1, x2, y2, x3, y3], dtype=np.float32) | |
| poly = get_best_begin_point_single(poly) | |
| return poly.tolist() | |
| def get_best_begin_point_single(coordinate): | |
| x1, y1, x2, y2, x3, y3, x4, y4 = coordinate | |
| xmin = min(x1, x2, x3, x4) | |
| ymin = min(y1, y2, y3, y4) | |
| xmax = max(x1, x2, x3, x4) | |
| ymax = max(y1, y2, y3, y4) | |
| combine = [[[x1, y1], [x2, y2], [x3, y3], [x4, y4]], | |
| [[x2, y2], [x3, y3], [x4, y4], [x1, y1]], | |
| [[x3, y3], [x4, y4], [x1, y1], [x2, y2]], | |
| [[x4, y4], [x1, y1], [x2, y2], [x3, y3]]] | |
| dst_coordinate = [[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]] | |
| force = 100000000.0 | |
| force_flag = 0 | |
| for i in range(4): | |
| temp_force = cal_line_length(combine[i][0], dst_coordinate[0]) \ | |
| + cal_line_length(combine[i][1], dst_coordinate[1]) \ | |
| + cal_line_length(combine[i][2], dst_coordinate[2]) \ | |
| + cal_line_length(combine[i][3], dst_coordinate[3]) | |
| if temp_force < force: | |
| force = temp_force | |
| force_flag = i | |
| if force_flag != 0: | |
| pass | |
| return np.array(combine[force_flag]).reshape(8) | |
| def cal_line_length(point1, point2): | |
| return math.sqrt( | |
| math.pow(point1[0] - point2[0], 2) + | |
| math.pow(point1[1] - point2[1], 2)) | |
| def main(): | |
| args = parse_args() | |
| root_path = args.root_path | |
| for split in ['train', 'val', 'test']: | |
| print(f'Processing {split} set...') | |
| with mmengine.Timer( | |
| print_tmpl='It takes {}s to convert IMGUR annotation'): | |
| anno_infos = collect_imgur_info( | |
| root_path, f'imgur5k_annotations_{split}.json') | |
| dump_ocr_data(anno_infos, | |
| osp.join(root_path, f'instances_{split}.json'), | |
| 'textdet') | |
| if __name__ == '__main__': | |
| main() | |