1、生成识别数据

获取PPOCRLabel格式的数据集中的目标的四个点，然后使用getPerspectiveTransform和warpPerspective获取图片，生成识别数据集。

import json
import os
import numpy as np
import cv2

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
def write_txt_img(src_path,label_txt,file_dir):
    with open(src_path, 'r', encoding='utf-8') as f:
        for line in f.readlines():
            print(line)
            content = line.split('\t')
            print(content[0])
            imag_name = content[0].split('/')[1]
            image_path = './train_data/icdar2015/text_localization/' + content[0]
            img = cv2.imread(image_path)
            content[1] = content[1].replace("'", "\"")
            list_dict = json.loads(content[1])
            nsize = len(list_dict)
            print(nsize)
            num = 0
            for i in range(nsize):
                print(list_dict[i])
                lin = list_dict[i]
                info = lin['transcription']
                info=info.replace(" ","")
                points = lin['points']
                points = [list(x) for x in points]
                points = np.float32([list(map(float, item)) for item in points])
                imag_name=str(num)+"_"+imag_name
                save_path = './train_data/rec/' +file_dir+ imag_name
                dst_img = get_rotate_crop_image(img, points)
                cv2.imwrite(save_path, dst_img)
                label_txt.write(file_dir+imag_name+'\t'+info+'\n')
                num=num+1
if not os.path.exists('train_data/rec/train/'):
    os.makedirs('train_data/rec/train/')
if not os.path.exists('train_data/rec/val/'):
    os.makedirs('train_data/rec/val/')
src_path = r"./train_data/icdar2015/text_localization/train.txt"
label_txt=r"./train_data/rec/train.txt"
src_test_path = r"./train_data/icdar2015/text_localization/val.txt"
label_test_txt=r"./train_data/rec/val.txt"
with open(label_txt, 'w') as w_label:
    write_txt_img(src_path,w_label,'train/')
with open(label_test_txt, 'w') as w_label:
    write_txt_img(src_test_path, w_label,'val/')

2、切分训练集和验证集

按照一定的比例，将数据集切分为训练集和验证集

# 制作数据集，将Label.txt切分为训练集和验证集
import os
import shutil
from sklearn.model_selection import train_test_split

os.makedirs('train',exist_ok=True)
os.makedirs('val',exist_ok=True)
label_txt='Label.txt'
with open(label_txt, 'r',encoding='gbk') as f:
   txt_List=f.readlines()
   trainval_files, val_files = train_test_split(txt_List, test_size=0.2, random_state=42)
   train_list=[]
   for file_Line in trainval_files:
       image_path = file_Line.split('\t')[0]
       image_new_path='train/'+image_path.split('/')[1]+'\t'+file_Line.split('\t')[1]
       train_list.append(image_new_path)
   f = open("train.txt", "w")
   f.writelines(train_list)
   f.close()
   val_list = []
   for file_Line in val_files:
       image_path = file_Line.split('\t')[0]
       image_new_path = 'val/' + image_path.split('/')[1] + '\t' + file_Line.split('\t')[1]
       val_list.append(image_new_path)
   f = open("val.txt", "w")
   f.writelines(val_list)
   f.close()
   for txt in trainval_files:
       image_name=txt.split('\t')[0]
       new_path="./train/"+image_name.split('/')[1]
       shutil.copy(image_name, new_path)
       print(image_name)
   for txt in val_files:
       image_name=txt.split('\t')[0]
       new_path="./val/"+image_name.split('/')[1]
       shutil.copy(image_name, new_path)
       print(image_name)

3、将数据集生成LabelImg格式

将PPOCRLabel格式的数据集转为LabelImg标注的xml格式的数据集。

import os
from collections import defaultdict
import cv2
# import misc_utils as utils  # pip3 install utils-misc==0.0.5 -i https://pypi.douban.com/simple/
import json

os.makedirs('./Annotations', exist_ok=True)
print('建立Annotations目录', 3)
# os.makedirs('./PaddleOCR/train_data/ImageSets/Main', exist_ok=True)
# print('建立ImageSets/Main目录', 3)

mem = defaultdict(list)

with open('Label.txt', 'r', encoding='utf8') as fp:
    s = [i.replace('\n','').split('\t') for i in fp.readlines()]
    for i in enumerate(s):
        path = i[1][0]
        anno = json.loads(i[1][1])
        print(anno)
        filename = path.split('/')[1]
        img = cv2.imread(path)
        cv2.imwrite('Annotations/'+filename.split('.')[0]+'.jpg',img)
        height, width = img.shape[:-1]
        for j in range(len(anno)):
            label = 'No'
            x1 = min(int(anno[j - 1]['points'][0][0]), int(anno[j - 1]['points'][1][0]),
                     int(anno[j - 1]['points'][2][0]), int(anno[j - 1]['points'][3][0]))
            x2 = max(int(anno[j - 1]['points'][0][0]), int(anno[j - 1]['points'][1][0]),
                     int(anno[j - 1]['points'][2][0]), int(anno[j - 1]['points'][3][0]))
            y1 = min(int(anno[j - 1]['points'][0][1]), int(anno[j - 1]['points'][1][1]),
                     int(anno[j - 1]['points'][2][1]), int(anno[j - 1]['points'][3][1]))
            y2 = max(int(anno[j - 1]['points'][0][1]), int(anno[j - 1]['points'][1][1]),
                     int(anno[j - 1]['points'][2][1]), int(anno[j - 1]['points'][3][1]))
            mem[filename].append([label, x1, y1, x2, y2])

            # for i, filename in enumerate(mem):
            #     img = cv2.imread(os.path.join('train', filename))
            # height, width, _ = img.shape

            with open(os.path.join('./Annotations', filename.split('.')[0]) + '.xml', 'w') as f:
                f.write(f"""<annotation>
            <folder>JPEGImages</folder>
            <filename>{filename.split('.')[0]}.jpg</filename>
            <size>
                <width>{width}</width>
                <height>{height}</height>
                <depth>3</depth>
            </size>
            <segmented>0</segmented>\n""")
                for label, x1, y1, x2, y2 in mem[filename]:
                    f.write(f"""    <object>
                <name>{label}</name>
                <pose>Unspecified</pose>
                <truncated>0</truncated>
                <difficult>0</difficult>
                <bndbox>
                    <xmin>{x1}</xmin>
                    <ymin>{y1}</ymin>
                    <xmax>{x2}</xmax>
                    <ymax>{y2}</ymax>
                </bndbox>
            </object>\n""")
                f.write("</annotation>")

4、将PPOCRLabel格式的数据集转为DBNet训练用的icdar2015格式的数据集

import os
import json

def json_2_icdar(js_path, ic_path):
    with open(js_path, 'r', encoding='utf-8') as f:
        for line in f.readlines():
            print(line)
            content = line.split('\t')
            print(content[0])
            txt_file = str(content[0].split('.')[0])+'.txt'
            dst_file = os.path.join(ic_path, txt_file)
            # write file
            file_lineinfo = open(txt_file, 'w', encoding='utf-8')
            list_dict = json.loads(content[1])
            nsize = len(list_dict)
            print(nsize)
            for i in range(nsize):
                print(list_dict[i])
                lin = list_dict[i]
                info = lin['transcription']
                points = lin['points']
                points = [int(y) for x in points for y in x]
                pts = ','.join(map(str, points))
                lineinfo = pts + ',' + info + '\n'
                file_lineinfo.write(lineinfo)
            file_lineinfo.close()


if __name__ == "__main__":
    src_path = r"train/Label.txt"
    dst_path = r""
    json_2_icdar(src_path, dst_path)

5、数据增强

对标注的数据集做旋转、高斯模糊、色彩饱和度、亮度等增强。

import json
import os

import cv2
import numpy as np
import torchvision.transforms as transforms
from torchtoolbox.transform import Cutout
from PIL import Image
from random import randint
# 数据预处理7
t=[
    transforms.ColorJitter(brightness=0.3, contrast=0.5, saturation=0.5),
    transforms.GaussianBlur(5,sigma=(0.1,0.5)),
    ]
transform = transforms.Compose([
    transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.2),
    transforms.GaussianBlur(5, sigma=(0.1, 3.0)),
    transforms.ToTensor(),
    transforms.ToPILImage(),
])
"""
旋转后图片返回
"""
def dumpRotateImage(img, degree): #图片，角度
    height, width = img.shape[:2]
    heightNew = height
    widthNew = width
    matRotation = cv2.getRotationMatrix2D((width//2,height//2), degree, 1)
    imgRotation = cv2.warpAffine(img, matRotation, (widthNew, heightNew), borderValue=(255, 255, 255))
    return imgRotation, matRotation

os.makedirs('train', exist_ok=True)
src_path = "Label_new.txt"
d_path='dd.txt'
radom_p=[-3,-2,-1,0,1,2,3,4,5]
with open(d_path, 'w') as w_label:
    with open(src_path, 'r', encoding='utf-8') as f:
        for line in f.readlines():
            content = line.split('\t')
            imag_name = content[0].split('/')[1]
            image_path = content[0]
            img = cv2.imread(image_path)
            list_dict = json.loads(content[1])
            nsize = len(list_dict)
            num = 0
            box = []
            info = ''
            for i in range(nsize):
                lin = list_dict[i]
                info = lin['transcription']
                info = info.replace(" ", "")
                points = lin['points']
                points = [list(x) for x in points]
                print(points)
                box = [b for a in points for b in a]
                break
            image = cv2.imread(image_path)
            for i in range(5):
                imgRotation, matRotation = dumpRotateImage(image,randint(-15,15))  # 旋转
                imgRotation = Image.fromarray(cv2.cvtColor(imgRotation, cv2.COLOR_BGR2RGB))
                imgRotation=transform(imgRotation).convert('RGB')
                imgRotation = cv2.cvtColor(np.asarray(imgRotation), cv2.COLOR_RGB2BGR)
                pt1 = np.dot(matRotation, np.array([[box[0]], [box[1]], [1]]))
                pt2 = np.dot(matRotation, np.array([[box[2]], [box[3]], [1]]))
                pt3 = np.dot(matRotation, np.array([[box[4]], [box[5]], [1]]))
                pt4 = np.dot(matRotation, np.array([[box[6]], [box[7]], [1]]))
                print(int(pt4[0]))
                if int(pt1[0])<5 or int(pt1[1])<5 or int(pt2[0])<5 or int(pt2[1])<5 or int(pt3[0])<5 or int(pt3[1])<5 or int(pt4[0])<5 or int(pt4[1])<5:
                    continue
                result_info = [{"transcription": info,
                                "points": [[int(pt1[0])+radom_p[randint(0, len(radom_p)-1)], int(pt1[1])+radom_p[randint(0, len(radom_p)-1)]], [int(pt2[0])+radom_p[randint(0, len(radom_p)-1)], int(pt2[1])+radom_p[randint(0, len(radom_p)-1)]],
                                           [int(pt3[0])+radom_p[randint(0, len(radom_p)-1)], int(pt3[1])+radom_p[randint(0, len(radom_p)-1)]],
                                           [int(pt4[0])+radom_p[randint(0, len(radom_p)-1)], int(pt4[1])+radom_p[randint(0, len(radom_p)-1)]]], "difficult": "false"}]

                imag_d_path = "train/" + imag_name.split('.')[0] + "_" + str(i) + "_0726." + imag_name.split('.')[
                    1] + "\t" + str(result_info)+'\n'
                print(imag_d_path)
                cv2.imwrite('./train/' + imag_name.split('.')[0] + "_" + str(i) + "_0726." + imag_name.split('.')[1],
                            imgRotation)
                w_label.write(imag_d_path.replace('\'','\"'))

6、删除没有标注的图片

将标注的数据和图片的列表做差，将多余的图片删除。

import os

image_list = os.listdir('train/')
label_txt = "Label_new.txt"
label_list=[]
with open(label_txt, 'r') as label_s:
    for line in label_s.readlines():
        label_list.append(line.split('\t')[0].split('/')[1])
cha_list=list(set(image_list) - set(label_list))
for img in cha_list:
    os.remove('train/'+img)