import copy
import json
import os
import random
import uuid
import cv2
import re
import pandas as pd
import numpy as np
import jieba
from shapely.geometry import Polygon, MultiPoint
from tools import get_file_paths, load_json
from word2vec import jwq_word2vec, simple_word2vec, jieba_and_tencent_word2vec
def bbox_iou(go_bbox, label_bbox, mode='iou'):
# 所有点的最小凸的表示形式，四边形对象，会自动计算四个点，最后顺序为：左上 左下 右下 右上 左上
go_poly = Polygon(go_bbox).convex_hull
label_poly = Polygon(label_bbox).convex_hull
if not go_poly.is_valid or not label_poly.is_valid:
print('formatting errors for boxes!!!! ')
return 0
if go_poly.area == 0 or label_poly.area == 0 :
return 0
inter = Polygon(go_poly).intersection(Polygon(label_poly)).area
go_area = Polygon(go_poly).area
return inter / go_area
# if mode == 'iou':
# union = go_poly.area + label_poly.area - inter
# elif mode =='tiou':
# union_poly = np.concatenate((go_bbox, label_bbox)) #合并两个box坐标，变为8*2
# union = MultiPoint(union_poly).convex_hull.area
# # coors = MultiPoint(union_poly).convex_hull.wkt
# elif mode == 'giou':
# union_poly = np.concatenate((go_bbox, label_bbox))
# union = MultiPoint(union_poly).envelope.area
# # coors = MultiPoint(union_poly).envelope.wkt
# elif mode == 'r_giou':
# union_poly = np.concatenate((go_bbox, label_bbox))
# union = MultiPoint(union_poly).minimum_rotated_rectangle.area
# # coors = MultiPoint(union_poly).minimum_rotated_rectangle.wkt
# else:
# raise Exception('incorrect mode!')
# if union == 0:
# return 0
# else:
# return inter / union
def clean_go_res(go_res_dir):
go_res_json_paths = get_file_paths(go_res_dir, ['.json', ])
for go_res_json_path in go_res_json_paths:
print('Info: start {0}'.format(go_res_json_path))
remove_idx_set = set()
src_go_res_list = load_json(go_res_json_path)
for idx, (_, text) in enumerate(src_go_res_list):
if text.strip() == '':
remove_idx_set.add(idx)
print(text)
if len(remove_idx_set) > 0:
for del_idx in remove_idx_set:
del src_go_res_list[del_idx]
go_res_list = sorted(src_go_res_list, key=lambda x: (x[0][1], x[0][0]), reverse=False)
with open(go_res_json_path, 'w') as fp:
json.dump(go_res_list, fp)
print('Rerewirte {0}'.format(go_res_json_path))
def char_length_statistics(go_res_dir):
max_char_length = None
target_file_name = None
go_res_json_paths = get_file_paths(go_res_dir, ['.json', ])
for go_res_json_path in go_res_json_paths:
print('Info: start {0}'.format(go_res_json_path))
src_go_res_list = load_json(go_res_json_path)
for _, text in src_go_res_list:
if max_char_length is None or len(text.strip()) > max_char_length:
max_char_length = len(text.strip())
target_file_name = go_res_json_path
return max_char_length, target_file_name
def char_length_statistics_jieba(go_res_dir):
max_char_length = None
target_file_name = None
max_char_length = None
statistics_dict = {}
go_res_json_paths = get_file_paths(go_res_dir, ['.json', ])
for go_res_json_path in go_res_json_paths:
print('Info: start {0}'.format(go_res_json_path))
src_go_res_list = load_json(go_res_json_path)
for _, text in src_go_res_list:
jieba_char_list = list(filter(lambda x:re.match(r'[\u4e00-\u9fa5]', x), jieba.lcut(text.strip())))
length = len(jieba_char_list)
if length in statistics_dict:
statistics_dict[length] += 1
else:
statistics_dict[length] = 1
if max_char_length is None or length > max_char_length:
target_file_name = go_res_json_path
target_jieba_char_list = jieba_char_list
max_char_length = length
return max_char_length, target_file_name, target_jieba_char_list, statistics_dict
def bbox_statistics(go_res_dir):
max_seq_count = None
seq_sum = 0
file_count = 0
go_res_json_paths = get_file_paths(go_res_dir, ['.json', ])
for go_res_json_path in go_res_json_paths:
print('Info: start {0}'.format(go_res_json_path))
go_res_list = load_json(go_res_json_path)
seq_sum += len(go_res_list)
file_count += 1
if max_seq_count is None or len(go_res_list) > max_seq_count:
max_seq_count = len(go_res_list)
max_seq_file_name = go_res_json_path
seq_lens_mean = seq_sum // file_count
return max_seq_count, seq_lens_mean, max_seq_file_name
def text_statistics(go_res_dir):
"""
Args:
go_res_dir: str 通用OCR的JSON文件夹
Returns: list 出现次数最多的文本及其次数
"""
json_count = 0
text_dict = {}
go_res_json_paths = get_file_paths(go_res_dir, ['.json', ])
for go_res_json_path in go_res_json_paths:
print('Info: start {0}'.format(go_res_json_path))
json_count += 1
go_res = load_json(go_res_json_path)
for _, text in go_res.values():
if text in text_dict:
text_dict[text] += 1
else:
text_dict[text] = 1
top_text_list = []
# 按照次数排序
for text, count in sorted(text_dict.items(), key=lambda x: x[1], reverse=True):
if text == '':
continue
# 丢弃：次数少于总数的2/3
if count <= json_count // 3:
break
top_text_list.append((text, count))
return top_text_list
def build_anno_file(dataset_dir, anno_file_path):
img_list = os.listdir(dataset_dir)
random.shuffle(img_list)
df = pd.DataFrame(columns=['name'])
df['name'] = img_list
df.to_csv(anno_file_path)
def build_dataset(img_dir, go_res_dir, label_dir, top_text_list, skip_list, save_dir, is_create_map=False):
"""
Args:
img_dir: str 图片目录
go_res_dir: str 通用OCR的JSON保存目录
label_dir: str 标注的JSON保存目录
top_text_list: list 出现次数最多的文本及其次数
skip_list: list 跳过的图片列表
save_dir: str 数据集保存目录
"""
if os.path.exists(save_dir):
return
else:
os.makedirs(save_dir, exist_ok=True)
# 开票日期 发票代码 机打号码 车辆类型 电话
# 发动机号码 车架号 帐号 开户银行 小写
group_cn_list = ['开票日期', '发票代码', '机打号码', '车辆类型', '电话', '发动机号码', '车架号', '帐号', '开户银行', '小写']
test_group_id = [1, 2, 5, 9, 20, 15, 16, 22, 24, 28]
create_map = {}
for img_name in sorted(os.listdir(img_dir)):
if img_name in skip_list:
print('Info: skip {0}'.format(img_name))
continue
print('Info: start {0}'.format(img_name))
image_path = os.path.join(img_dir, img_name)
img = cv2.imread(image_path)
h, w, _ = img.shape
base_image_name, _ = os.path.splitext(img_name)
go_res_json_path = os.path.join(go_res_dir, '{0}.json'.format(base_image_name))
go_res_list = load_json(go_res_json_path)
valid_lens = len(go_res_list)
top_text_idx_set = set()
for top_text, _ in top_text_list:
for go_idx, (_, text) in enumerate(go_res_list):
if text == top_text:
top_text_idx_set.add(go_idx)
break
label_json_path = os.path.join(label_dir, '{0}.json'.format(base_image_name))
label_res = load_json(label_json_path)
group_list = []
for group_id in test_group_id:
for item in label_res.get("shapes", []):
if item.get("group_id") == group_id:
label_bbox = list()
for point in item['points']:
label_bbox.extend(point)
group_list.append(label_bbox)
break
else:
group_list.append(None)
label_idx_dict = dict()
for label_idx, label_bbox in enumerate(group_list):
if isinstance(label_bbox, list):
for go_idx, (go_bbox, _) in enumerate(go_res_list):
if go_idx in top_text_idx_set or go_idx in label_idx_dict:
continue
go_bbox_rebuild = [
[go_bbox[0], go_bbox[1]],
[go_bbox[2], go_bbox[3]],
[go_bbox[4], go_bbox[5]],
[go_bbox[6], go_bbox[7]],
]
label_bbox_rebuild = [
[label_bbox[0], label_bbox[1]],
[label_bbox[2], label_bbox[1]],
[label_bbox[2], label_bbox[3]],
[label_bbox[0], label_bbox[3]],
]
iou = bbox_iou(go_bbox_rebuild, label_bbox_rebuild)
if iou >= 0.2:
label_idx_dict[go_idx] = label_idx
X = list()
y_true = list()
X_no_text = list()
# dim = 1 + 5 + 8
# text_vec_max_lens = 15 * 50
# dim = 1 + 5 + 8 + text_vec_max_lens
max_jieba_char = 4
text_vec_max_lens = max_jieba_char * 100
dim = 1 + 5 + 8 + text_vec_max_lens
num_classes = 10
for i in range(160):
if i >= valid_lens:
X.append([0. for _ in range(dim)])
y_true.append([0 for _ in range(num_classes)])
X_no_text.append([0. for _ in range(dim)])
elif i in top_text_idx_set:
(x0, y0, x1, y1, x2, y2, x3, y3), text = go_res_list[i]
feature_vec = [1.]
feature_vec.extend(simple_word2vec(text))
feature_vec.extend([(x0/w), (y0/h), (x1/w), (y1/h), (x2/w), (y2/h), (x3/w), (y3/h)])
# feature_vec.extend(jwq_word2vec(text, text_vec_max_lens))
feature_vec.extend(jieba_and_tencent_word2vec(text, max_jieba_char))
X.append(feature_vec)
y_true.append([0 for _ in range(num_classes)])
feature_vec_no_text = [1.]
feature_vec_no_text.extend([0. for _ in range(5)])
feature_vec_no_text.extend([(x0/w), (y0/h), (x1/w), (y1/h), (x2/w), (y2/h), (x3/w), (y3/h)])
# feature_vec.extend(jwq_word2vec(text, text_vec_max_lens))
feature_vec_no_text.extend([0. for _ in range(text_vec_max_lens)])
X_no_text.append(feature_vec_no_text)
elif i in label_idx_dict:
(x0, y0, x1, y1, x2, y2, x3, y3), text = go_res_list[i]
feature_vec = [0.]
feature_vec.extend(simple_word2vec(text))
feature_vec.extend([(x0/w), (y0/h), (x1/w), (y1/h), (x2/w), (y2/h), (x3/w), (y3/h)])
# feature_vec.extend(jwq_word2vec(text, text_vec_max_lens))
feature_vec.extend(jieba_and_tencent_word2vec(text, max_jieba_char))
X.append(feature_vec)
base_label_list = [0 for _ in range(num_classes)]
base_label_list[label_idx_dict[i]] = 1
y_true.append(base_label_list)
feature_vec_no_text = [0.]
feature_vec_no_text.extend([0. for _ in range(5)])
feature_vec_no_text.extend([(x0/w), (y0/h), (x1/w), (y1/h), (x2/w), (y2/h), (x3/w), (y3/h)])
# feature_vec.extend(jwq_word2vec(text, text_vec_max_lens))
feature_vec_no_text.extend([0. for _ in range(text_vec_max_lens)])
X_no_text.append(feature_vec_no_text)
else:
(x0, y0, x1, y1, x2, y2, x3, y3), text = go_res_list[i]
feature_vec = [0.]
feature_vec.extend(simple_word2vec(text))
feature_vec.extend([(x0/w), (y0/h), (x1/w), (y1/h), (x2/w), (y2/h), (x3/w), (y3/h)])
# feature_vec.extend(jwq_word2vec(text, text_vec_max_lens))
feature_vec.extend(jieba_and_tencent_word2vec(text, max_jieba_char))
X.append(feature_vec)
y_true.append([0 for _ in range(num_classes)])
feature_vec_no_text = [0.]
feature_vec_no_text.extend([0. for _ in range(5)])
feature_vec_no_text.extend([(x0/w), (y0/h), (x1/w), (y1/h), (x2/w), (y2/h), (x3/w), (y3/h)])
# feature_vec.extend(jwq_word2vec(text, text_vec_max_lens))
feature_vec_no_text.extend([0. for _ in range(text_vec_max_lens)])
X_no_text.append(feature_vec_no_text)
all_data = [X, y_true, valid_lens]
all_data_no_text = [X_no_text, y_true, valid_lens]
save_json_name = '{0}.json'.format(uuid.uuid3(uuid.NAMESPACE_DNS, img_name))
with open(os.path.join(save_dir, save_json_name), 'w') as fp:
json.dump(all_data, fp)
save_json_name_2 = '{0}.json'.format(uuid.uuid3(uuid.NAMESPACE_DNS, '{0}_no_text'.format(img_name)))
with open(os.path.join(save_dir, save_json_name_2), 'w') as fp:
json.dump(all_data_no_text, fp)
if is_create_map:
create_map[img_name] = {
'x_y_valid_lens': save_json_name,
'find_top_text': [go_res_list[i][-1] for i in top_text_idx_set],
'find_value': {go_res_list[k][-1]: group_cn_list[v] for k, v in label_idx_dict.items()}
}
# break
# print(create_map)
# print(is_create_map)
if create_map:
# print(create_map)
with open(os.path.join(os.path.dirname(save_dir), 'create_map.json'), 'w') as fp:
json.dump(create_map, fp)
# print('top text find:')
# for i in top_text_idx_set:
# _, text = go_res_list[i]
# print(text)
# print('-------------')
# print('label value find:')
# for k, v in label_idx_dict.items():
# _, text = go_res_list[k]
# print('{0}: {1}'.format(group_cn_list[v], text))
# break
if __name__ == '__main__':
base_dir = '/Users/zhouweiqi/Downloads/gcfp/data'
go_dir = os.path.join(base_dir, 'go_res')
dataset_save_dir = os.path.join(base_dir, 'dataset160x414x10-no-text')
label_dir = os.path.join(base_dir, 'labeled')
train_go_path = os.path.join(go_dir, 'train')
train_image_path = os.path.join(label_dir, 'train', 'image')
train_label_path = os.path.join(label_dir, 'train', 'label')
train_dataset_dir = os.path.join(dataset_save_dir, 'train')
train_anno_file_path = os.path.join(dataset_save_dir, 'train.csv')
valid_go_path = os.path.join(go_dir, 'valid')
valid_image_path = os.path.join(label_dir, 'valid', 'image')
valid_label_path = os.path.join(label_dir, 'valid', 'label')
valid_dataset_dir = os.path.join(dataset_save_dir, 'valid')
valid_anno_file_path = os.path.join(dataset_save_dir, 'valid.csv')
# max_seq_lens, seq_lens_mean, max_seq_file_name = bbox_statistics(go_dir)
# print(max_seq_lens) # 152
# print(max_seq_file_name) # train/CH-B101805176_page_2_img_0.json
# print(seq_lens_mean) # 92
# max_char_lens, target_file_name = char_length_statistics(go_dir)
# print(max_char_lens) # 72
# print(target_file_name) # train/CH-B103053828-4.json
# max_char_length, target_file_name, target_jieba_char_list, statistics_dict = char_length_statistics_jieba(go_dir)
# print(max_char_length) # 24
# print(target_file_name) # train/CH-B102551568-6.json
# print(target_jieba_char_list)
# print(statistics_dict) # {2: 12077, 1: 12751, 0: 13073, 3: 4423, 4: 1212, 5: 969, 6: 744, 7: 524, 8: 199, 10: 45, 12: 9, 18: 44, 9: 109, 11: 19, 13: 4, 16: 4, 21: 2, 19: 2, 15: 8, 17: 7, 14: 3, 20: 1, 24: 1}
# top_text_list = text_statistics(go_dir)
# for t in top_text_list:
# print(t)
filter_from_top_text_list = [
('机器编号', 496),
('购买方名称', 496),
('合格证号', 495),
('进口证明书号', 495),
('机打代码', 494),
('车辆类型', 492),
('完税凭证号码', 492),
('机打号码', 491),
('发动机号码', 491),
('主管税务', 491),
('价税合计', 489),
('机关及代码', 489),
('销货单位名称', 486),
('厂牌型号', 485),
('产地', 485),
('商检单号', 483),
('电话', 476),
('开户银行', 472),
('车辆识别代号/车架号码', 463),
('身份证号码', 454),
('吨位', 452),
('备注:一车一票', 439),
('地', 432),
('账号', 431),
('统一社会信用代码/', 424),
('限乘人数', 404),
('税额', 465),
('址', 392)
]
skip_list_train = [
'CH-B101910792-page-12.jpg',
'CH-B101655312-page-13.jpg',
'CH-B102278656.jpg',
'CH-B101846620_page_1_img_0.jpg',
'CH-B103062528-0.jpg',
'CH-B102613120-3.jpg',
'CH-B102997980-3.jpg',
'CH-B102680060-3.jpg',
# # 'CH-B102995500-2.jpg', # 没value
]
skip_list_valid = [
# 'CH-B102897920-2.jpg',
# 'CH-B102551284-0.jpg',
# 'CH-B102879376-2.jpg',
# 'CH-B101509488-page-16.jpg',
# 'CH-B102708352-2.jpg',
]
build_dataset(train_image_path, train_go_path, train_label_path, filter_from_top_text_list, skip_list_train, train_dataset_dir)
build_anno_file(train_dataset_dir, train_anno_file_path)
build_dataset(valid_image_path, valid_go_path, valid_label_path, filter_from_top_text_list, skip_list_valid, valid_dataset_dir, True)
build_anno_file(valid_dataset_dir, valid_anno_file_path)
# print(simple_word2vec(' fd2jk接口 额24；叁‘，。测ADF壹试！¥？ '))
# print(jwq_word2vec('发', 15*50))