|
|
|
|
|
import argparse |
|
|
import json |
|
|
import os.path as osp |
|
|
import time |
|
|
|
|
|
import lmdb |
|
|
import mmcv |
|
|
import numpy as np |
|
|
from scipy.io import loadmat |
|
|
from shapely.geometry import Polygon |
|
|
|
|
|
from mmocr.utils import check_argument |
|
|
|
|
|
|
|
|
def trace_boundary(char_boxes): |
|
|
"""Trace the boundary point of text. |
|
|
|
|
|
Args: |
|
|
char_boxes (list[ndarray]): The char boxes for one text. Each element |
|
|
is 4x2 ndarray. |
|
|
|
|
|
Returns: |
|
|
boundary (ndarray): The boundary point sets with size nx2. |
|
|
""" |
|
|
assert check_argument.is_type_list(char_boxes, np.ndarray) |
|
|
|
|
|
|
|
|
p_top = [box[0:2] for box in char_boxes] |
|
|
|
|
|
p_bottom = [ |
|
|
char_boxes[idx][[2, 3], :] |
|
|
for idx in range(len(char_boxes) - 1, -1, -1) |
|
|
] |
|
|
|
|
|
p = p_top + p_bottom |
|
|
|
|
|
boundary = np.concatenate(p).astype(int) |
|
|
|
|
|
return boundary |
|
|
|
|
|
|
|
|
def match_bbox_char_str(bboxes, char_bboxes, strs): |
|
|
"""match the bboxes, char bboxes, and strs. |
|
|
|
|
|
Args: |
|
|
bboxes (ndarray): The text boxes of size (2, 4, num_box). |
|
|
char_bboxes (ndarray): The char boxes of size (2, 4, num_char_box). |
|
|
strs (ndarray): The string of size (num_strs,) |
|
|
""" |
|
|
assert isinstance(bboxes, np.ndarray) |
|
|
assert isinstance(char_bboxes, np.ndarray) |
|
|
assert isinstance(strs, np.ndarray) |
|
|
bboxes = bboxes.astype(np.int32) |
|
|
char_bboxes = char_bboxes.astype(np.int32) |
|
|
|
|
|
if len(char_bboxes.shape) == 2: |
|
|
char_bboxes = np.expand_dims(char_bboxes, axis=2) |
|
|
char_bboxes = np.transpose(char_bboxes, (2, 1, 0)) |
|
|
if len(bboxes.shape) == 2: |
|
|
bboxes = np.expand_dims(bboxes, axis=2) |
|
|
bboxes = np.transpose(bboxes, (2, 1, 0)) |
|
|
chars = ''.join(strs).replace('\n', '').replace(' ', '') |
|
|
num_boxes = bboxes.shape[0] |
|
|
|
|
|
poly_list = [Polygon(bboxes[iter]) for iter in range(num_boxes)] |
|
|
poly_box_list = [bboxes[iter] for iter in range(num_boxes)] |
|
|
|
|
|
poly_char_list = [[] for iter in range(num_boxes)] |
|
|
poly_char_idx_list = [[] for iter in range(num_boxes)] |
|
|
poly_charbox_list = [[] for iter in range(num_boxes)] |
|
|
|
|
|
words = [] |
|
|
for s in strs: |
|
|
words += s.split() |
|
|
words_len = [len(w) for w in words] |
|
|
words_end_inx = np.cumsum(words_len) |
|
|
start_inx = 0 |
|
|
for word_inx, end_inx in enumerate(words_end_inx): |
|
|
for char_inx in range(start_inx, end_inx): |
|
|
poly_char_idx_list[word_inx].append(char_inx) |
|
|
poly_char_list[word_inx].append(chars[char_inx]) |
|
|
poly_charbox_list[word_inx].append(char_bboxes[char_inx]) |
|
|
start_inx = end_inx |
|
|
|
|
|
for box_inx in range(num_boxes): |
|
|
assert len(poly_charbox_list[box_inx]) > 0 |
|
|
|
|
|
poly_boundary_list = [] |
|
|
for item in poly_charbox_list: |
|
|
boundary = np.ndarray((0, 2)) |
|
|
if len(item) > 0: |
|
|
boundary = trace_boundary(item) |
|
|
poly_boundary_list.append(boundary) |
|
|
|
|
|
return (poly_list, poly_box_list, poly_boundary_list, poly_charbox_list, |
|
|
poly_char_idx_list, poly_char_list) |
|
|
|
|
|
|
|
|
def convert_annotations(root_path, gt_name, lmdb_name): |
|
|
"""Convert the annotation into lmdb dataset. |
|
|
|
|
|
Args: |
|
|
root_path (str): The root path of dataset. |
|
|
gt_name (str): The ground truth filename. |
|
|
lmdb_name (str): The output lmdb filename. |
|
|
""" |
|
|
assert isinstance(root_path, str) |
|
|
assert isinstance(gt_name, str) |
|
|
assert isinstance(lmdb_name, str) |
|
|
start_time = time.time() |
|
|
gt = loadmat(gt_name) |
|
|
img_num = len(gt['imnames'][0]) |
|
|
env = lmdb.open(lmdb_name, map_size=int(1e9 * 40)) |
|
|
with env.begin(write=True) as txn: |
|
|
for img_id in range(img_num): |
|
|
if img_id % 1000 == 0 and img_id > 0: |
|
|
total_time_sec = time.time() - start_time |
|
|
avg_time_sec = total_time_sec / img_id |
|
|
eta_mins = (avg_time_sec * (img_num - img_id)) / 60 |
|
|
print(f'\ncurrent_img/total_imgs {img_id}/{img_num} | ' |
|
|
f'eta: {eta_mins:.3f} mins') |
|
|
|
|
|
img_file = osp.join(root_path, 'imgs', gt['imnames'][0][img_id][0]) |
|
|
img = mmcv.imread(img_file, 'unchanged') |
|
|
height, width = img.shape[0:2] |
|
|
img_json = {} |
|
|
img_json['file_name'] = gt['imnames'][0][img_id][0] |
|
|
img_json['height'] = height |
|
|
img_json['width'] = width |
|
|
img_json['annotations'] = [] |
|
|
wordBB = gt['wordBB'][0][img_id] |
|
|
charBB = gt['charBB'][0][img_id] |
|
|
txt = gt['txt'][0][img_id] |
|
|
poly_list, _, poly_boundary_list, _, _, _ = match_bbox_char_str( |
|
|
wordBB, charBB, txt) |
|
|
for poly_inx in range(len(poly_list)): |
|
|
|
|
|
polygon = poly_list[poly_inx] |
|
|
min_x, min_y, max_x, max_y = polygon.bounds |
|
|
bbox = [min_x, min_y, max_x - min_x, max_y - min_y] |
|
|
anno_info = dict() |
|
|
anno_info['iscrowd'] = 0 |
|
|
anno_info['category_id'] = 1 |
|
|
anno_info['bbox'] = bbox |
|
|
anno_info['segmentation'] = [ |
|
|
poly_boundary_list[poly_inx].flatten().tolist() |
|
|
] |
|
|
|
|
|
img_json['annotations'].append(anno_info) |
|
|
string = json.dumps(img_json) |
|
|
txn.put(str(img_id).encode('utf8'), string.encode('utf8')) |
|
|
key = 'total_number'.encode('utf8') |
|
|
value = str(img_num).encode('utf8') |
|
|
txn.put(key, value) |
|
|
|
|
|
|
|
|
def parse_args(): |
|
|
parser = argparse.ArgumentParser( |
|
|
description='Convert synthtext to lmdb dataset') |
|
|
parser.add_argument('synthtext_path', help='synthetic root path') |
|
|
parser.add_argument('-o', '--out-dir', help='output path') |
|
|
args = parser.parse_args() |
|
|
return args |
|
|
|
|
|
|
|
|
def main(): |
|
|
args = parse_args() |
|
|
synthtext_path = args.synthtext_path |
|
|
out_dir = args.out_dir if args.out_dir else synthtext_path |
|
|
mmcv.mkdir_or_exist(out_dir) |
|
|
|
|
|
gt_name = osp.join(synthtext_path, 'gt.mat') |
|
|
lmdb_name = 'synthtext.lmdb' |
|
|
convert_annotations(synthtext_path, gt_name, osp.join(out_dir, lmdb_name)) |
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
main() |
|
|
|
