168 lines
6.2 KiB
Python
168 lines
6.2 KiB
Python
import cv2
|
|
import os
|
|
import json
|
|
from tqdm import tqdm
|
|
import numpy as np
|
|
|
|
provinces = [
|
|
"皖", "沪", "津", "渝", "冀", "晋", "蒙", "辽", "吉", "黑", "苏", "浙", "京", "闽", "赣",
|
|
"鲁", "豫", "鄂", "湘", "粤", "桂", "琼", "川", "贵", "云", "藏", "陕", "甘", "青", "宁",
|
|
"新", "警", "学", "O"
|
|
]
|
|
alphabets = [
|
|
'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q',
|
|
'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'O'
|
|
]
|
|
ads = [
|
|
'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q',
|
|
'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '0', '1', '2', '3', '4', '5',
|
|
'6', '7', '8', '9', 'O'
|
|
]
|
|
|
|
|
|
def make_label_2020(img_dir, save_gt_folder, phase):
|
|
crop_img_save_dir = os.path.join(save_gt_folder, phase, 'crop_imgs')
|
|
os.makedirs(crop_img_save_dir, exist_ok=True)
|
|
|
|
f_det = open(
|
|
os.path.join(save_gt_folder, phase, 'det.txt'), 'w', encoding='utf-8')
|
|
f_rec = open(
|
|
os.path.join(save_gt_folder, phase, 'rec.txt'), 'w', encoding='utf-8')
|
|
|
|
i = 0
|
|
for filename in tqdm(os.listdir(os.path.join(img_dir, phase))):
|
|
str_list = filename.split('-')
|
|
if len(str_list) < 5:
|
|
continue
|
|
coord_list = str_list[3].split('_')
|
|
txt_list = str_list[4].split('_')
|
|
boxes = []
|
|
for coord in coord_list:
|
|
boxes.append([int(x) for x in coord.split("&")])
|
|
boxes = [boxes[2], boxes[3], boxes[0], boxes[1]]
|
|
lp_number = provinces[int(txt_list[0])] + alphabets[int(txt_list[
|
|
1])] + ''.join([ads[int(x)] for x in txt_list[2:]])
|
|
|
|
# det
|
|
det_info = [{'points': boxes, 'transcription': lp_number}]
|
|
f_det.write('{}\t{}\n'.format(
|
|
os.path.join("CCPD2020/ccpd_green", phase, filename),
|
|
json.dumps(
|
|
det_info, ensure_ascii=False)))
|
|
|
|
# rec
|
|
boxes = np.float32(boxes)
|
|
img = cv2.imread(os.path.join(img_dir, phase, filename))
|
|
# crop_img = img[int(boxes[:,1].min()):int(boxes[:,1].max()),int(boxes[:,0].min()):int(boxes[:,0].max())]
|
|
crop_img = get_rotate_crop_image(img, boxes)
|
|
crop_img_save_filename = '{}_{}.jpg'.format(i, '_'.join(txt_list))
|
|
crop_img_save_path = os.path.join(crop_img_save_dir,
|
|
crop_img_save_filename)
|
|
cv2.imwrite(crop_img_save_path, crop_img)
|
|
f_rec.write('{}/{}/crop_imgs/{}\t{}\n'.format(
|
|
"CCPD2020/PPOCR", phase, crop_img_save_filename, lp_number))
|
|
i += 1
|
|
f_det.close()
|
|
f_rec.close()
|
|
|
|
|
|
def make_label_2019(list_dir, save_gt_folder, phase):
|
|
crop_img_save_dir = os.path.join(save_gt_folder, phase, 'crop_imgs')
|
|
os.makedirs(crop_img_save_dir, exist_ok=True)
|
|
|
|
f_det = open(
|
|
os.path.join(save_gt_folder, phase, 'det.txt'), 'w', encoding='utf-8')
|
|
f_rec = open(
|
|
os.path.join(save_gt_folder, phase, 'rec.txt'), 'w', encoding='utf-8')
|
|
|
|
with open(os.path.join(list_dir, phase + ".txt"), 'r') as rf:
|
|
imglist = rf.readlines()
|
|
|
|
i = 0
|
|
for idx, filename in enumerate(imglist):
|
|
if idx % 1000 == 0:
|
|
print("{}/{}".format(idx, len(imglist)))
|
|
filename = filename.strip()
|
|
str_list = filename.split('-')
|
|
if len(str_list) < 5:
|
|
continue
|
|
coord_list = str_list[3].split('_')
|
|
txt_list = str_list[4].split('_')
|
|
boxes = []
|
|
for coord in coord_list:
|
|
boxes.append([int(x) for x in coord.split("&")])
|
|
boxes = [boxes[2], boxes[3], boxes[0], boxes[1]]
|
|
lp_number = provinces[int(txt_list[0])] + alphabets[int(txt_list[
|
|
1])] + ''.join([ads[int(x)] for x in txt_list[2:]])
|
|
|
|
# det
|
|
det_info = [{'points': boxes, 'transcription': lp_number}]
|
|
f_det.write('{}\t{}\n'.format(
|
|
os.path.join("CCPD2019", filename),
|
|
json.dumps(
|
|
det_info, ensure_ascii=False)))
|
|
|
|
# rec
|
|
boxes = np.float32(boxes)
|
|
imgpath = os.path.join(list_dir[:-7], filename)
|
|
img = cv2.imread(imgpath)
|
|
# crop_img = img[int(boxes[:,1].min()):int(boxes[:,1].max()),int(boxes[:,0].min()):int(boxes[:,0].max())]
|
|
crop_img = get_rotate_crop_image(img, boxes)
|
|
crop_img_save_filename = '{}_{}.jpg'.format(i, '_'.join(txt_list))
|
|
crop_img_save_path = os.path.join(crop_img_save_dir,
|
|
crop_img_save_filename)
|
|
cv2.imwrite(crop_img_save_path, crop_img)
|
|
f_rec.write('{}/{}/crop_imgs/{}\t{}\n'.format(
|
|
"CCPD2019/PPOCR", phase, crop_img_save_filename, lp_number))
|
|
i += 1
|
|
f_det.close()
|
|
f_rec.close()
|
|
|
|
|
|
def get_rotate_crop_image(img, points):
|
|
'''
|
|
img_height, img_width = img.shape[0:2]
|
|
left = int(np.min(points[:, 0]))
|
|
right = int(np.max(points[:, 0]))
|
|
top = int(np.min(points[:, 1]))
|
|
bottom = int(np.max(points[:, 1]))
|
|
img_crop = img[top:bottom, left:right, :].copy()
|
|
points[:, 0] = points[:, 0] - left
|
|
points[:, 1] = points[:, 1] - top
|
|
'''
|
|
assert len(points) == 4, "shape of points must be 4*2"
|
|
img_crop_width = int(
|
|
max(
|
|
np.linalg.norm(points[0] - points[1]),
|
|
np.linalg.norm(points[2] - points[3])))
|
|
img_crop_height = int(
|
|
max(
|
|
np.linalg.norm(points[0] - points[3]),
|
|
np.linalg.norm(points[1] - points[2])))
|
|
pts_std = np.float32([[0, 0], [img_crop_width, 0],
|
|
[img_crop_width, img_crop_height],
|
|
[0, img_crop_height]])
|
|
M = cv2.getPerspectiveTransform(points, pts_std)
|
|
dst_img = cv2.warpPerspective(
|
|
img,
|
|
M, (img_crop_width, img_crop_height),
|
|
borderMode=cv2.BORDER_REPLICATE,
|
|
flags=cv2.INTER_CUBIC)
|
|
dst_img_height, dst_img_width = dst_img.shape[0:2]
|
|
if dst_img_height * 1.0 / dst_img_width >= 1.5:
|
|
dst_img = np.rot90(dst_img)
|
|
return dst_img
|
|
|
|
|
|
img_dir = './CCPD2020/ccpd_green'
|
|
save_gt_folder = './CCPD2020/PPOCR'
|
|
# phase = 'train' # change to val and test to make val dataset and test dataset
|
|
for phase in ['train', 'val', 'test']:
|
|
make_label_2020(img_dir, save_gt_folder, phase)
|
|
|
|
list_dir = './CCPD2019/splits/'
|
|
save_gt_folder = './CCPD2019/PPOCR'
|
|
|
|
for phase in ['train', 'val', 'test']:
|
|
make_label_2019(list_dir, save_gt_folder, phase)
|