DocumentAIKit/idcardgenerator/idcard_det_rec_generate_forward.py

import argparse
import os.path
import cv2
from PIL import Image, ImageDraw, ImageFont
from tqdm import tqdm
from img_augmentation import img_augmentation_light as light
from img_augmentation import img_augmentation_noise as noise


def generate_idcard(info_dir, save_dir):
    if not os.path.exists(save_dir + '/crop_imgs'):
        os.makedirs(save_dir + '/crop_imgs')
    with open(info_dir, 'r', encoding='utf-8') as fp:
        text_lines = fp.readlines()
        cnt = 0
        for text_line in tqdm(text_lines):
            info_list = text_line.rstrip('\n').split(',')
            bg_path = "usedres/tw_idcard_bg_1.jpg"
            if 20 < cnt <= 60:
                bg_path = "usedres/tw_idcard_bg_2.png"
            if cnt > 60:
                bg_path = "usedres/tw_idcard_bg_color.jpg"
            pil_image = Image.open(bg_path)
            # pil_image 接收住这个图片对象
            # width 为图片的宽, height为图片的高

            width, height = pil_image.size
            # 生成一张尺寸为 width * height  背景色为白色的图片
            bg = Image.new('RGB', (width, height), color=(0, 0, 0))

            bg.paste(pil_image, (0, 0))  # 写入底图
            # 第一个参数为图片
            # 第二个参数为图片的要从哪里开始写入,

            name = info_list[0]  # 写入的文字
            year = info_list[1]
            month = info_list[2]
            day = info_list[3]
            birth = '民國   ' + year + '  年   ' + month + '  月   ' + day + '  日'
            addr = info_list[4]
            sex = info_list[5]
            id = info_list[6]
            icon_addr = info_list[7]
            name_font_size = 45
            temp_str = ''
            if len(name) <= 6:
                num_of_space = (4 * (6 - len(name)))/(len(name)-1)
                for i in range(0, len(name) - 1):
                    temp_str += name[i]
                    for j in range(0, int(num_of_space)):
                        temp_str += ' '
                temp_str += name[len(name)-1]
                name = temp_str
            elif 6 < len(name) < 9:
                name_font_size = 35
            elif 9 <= len(name) <= 12:
                name_font_size = 25

            font = "./usedres/R-PMingLiU-TW-2.ttf"  # 一个字体文件
            name_font = ImageFont.truetype(font, name_font_size)  # 设置名字字体和大小
            addr_font = ImageFont.truetype(font, 22)
            other_font = ImageFont.truetype(font, 26)  # 设置字体和大小
            id_font = ImageFont.truetype(font, 34)
            t_font = ImageFont.truetype(font, 25)
            # 计算出要写入的文字占用的像素
            # 创建一个可以在给定图像上绘图的对象
            draw = ImageDraw.Draw(bg)

            draw.text((110, 195), name, fill="#1E1E1E", font=name_font)
            draw.text((110, 290), birth, fill="#1E1E1E", font=other_font)
            draw.text((570, 295), sex, fill="#1E1E1E", font=t_font)
            draw.text((110, 365), addr, fill="#1E1E1E", font=addr_font)
            draw.text((465, 365), id, fill="#B4495C", font=id_font)

            path_list = []
            # 保存画布
            img_path = save_dir + '/idcard_' + str(cnt) + '_' + icon_addr.split('/')[-1].split('.')[0] + '.png'
            # img_path = save_dir + '/idcard_' + str(cnt) + '_mask.png'
            bg.save(img_path, "PNG")
            path_list.append(img_path)

            # 增亮、變暗
            light.work(img_path, '')
            path_list.append(str(img_path)[0:-4] + '_dark.png')
            path_list.append(str(img_path)[0:-4] + '_bright.png')

            # 高斯噪聲、椒鹽噪聲
            noise.add_noise(img_path, '')
            path_list.append(str(img_path)[0:-4] + "_guassian.png")
            path_list.append(str(img_path)[0:-4] + "_sp.png")


            cnt += 1
            boxes = []

            box_name_0 = [105, 195, 400, 245]
            box0 = []
            box0.append(box_name_0)
            box0.append(name.replace(' ', ''))
            box0.append('name')
            boxes.append(box0)

            box_birth_0 = [105, 290, 420, 320]
            box1 = []
            box1.append(box_birth_0)
            box1.append(birth.replace(' ', ''))
            box1.append('birth')
            boxes.append(box1)

            box_sex_0 = [570, 295, 595, 325]
            box2 = []
            box2.append(box_sex_0)
            box2.append(sex.replace(' ', ''))
            box2.append('sex')
            boxes.append(box2)

            box_addr_0 = [105, 365, 420, 390]
            box3 = []
            box3.append(box_addr_0)
            box3.append(addr.replace(' ', ''))
            box3.append('date')
            boxes.append(box3)

            box_id_0 = [465, 365, 640, 400]
            box4 = []
            box4.append(box_id_0)
            box4.append(id)
            box4.append('id')
            boxes.append(box4)

            num = 0
            fp = open(save_dir + '/rec_gt.txt', 'a', encoding='utf-8')
            fp1 = open(save_dir + '/Label.txt', 'a', encoding='utf-8')
            for path in path_list:
                # print(path)
                jishu = 1
                anno = str(path) + '\t['
                for box in boxes:
                    if jishu <= 5:
                        if jishu == 5:
                            s = '{"transcription": "' + box[1] + '", "points": [[' + str(box[0][0]) + ', ' + str(box[0][1]) + '], [' + str(box[0][2]) + ', ' + str(box[0][1]) + '], [' + str(box[0][2]) + ', ' + str(box[0][3]) + '], [' + str(box[0][0]) + ', ' + str(box[0][3]) + ']], "difficult": false, "key_cls": "' + box[2] + '"}'
                        else:
                           s = '{"transcription": "' + box[1] + '", "points": [[' + str(box[0][0]) + ', ' + str(
                                box[0][1]) + '], [' + str(box[0][2]) + ', ' + str(box[0][1]) + '], [' + str(
                                box[0][2]) + ', ' + str(box[0][3]) + '], [' + str(box[0][0]) + ', ' + str(
                                box[0][3]) + ']], "difficult": false, "key_cls": "' + box[2] + '"}, '
                        anno += s
                    image_cut_save(path, box[0][0], box[0][1], box[0][2], box[0][3], save_dir + '/crop_imgs/' + str(path).split('/')[-1].split('.')[0] + '_' + str(num) + '.png')
                    line = 'crop_imgs/' + str(path).split('/')[-1].split('.')[0] + '_' + str(num) + '.png' + '\t' + box[1] + '\n'
                    fp.write(line)
                    num += 1
                    jishu += 1
                anno += ']\n'
                fp1.write(anno)
                # print(anno)
            fp1.close()
            fp.close()


def image_cut_save(path, left, upper, right, lower, save_path):
    """
        所截区域图片保存
    :param path: 图片路径
    :param left: 区块左上角位置的像素点离图片左边界的距离
    :param upper：区块左上角位置的像素点离图片上边界的距离
    :param right：区块右下角位置的像素点离图片左边界的距离
    :param lower：区块右下角位置的像素点离图片上边界的距离
     故需满足：lower > upper、right > left
    :param save_path: 所截图片保存位置
    """
    img = cv2.imread(path)  # 打开图像
    cropped = img[upper:lower, left:right]
    cv2.imwrite(save_path, cropped)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--info_dir', type=str, default='val.txt')
    parser.add_argument('--save_dir', type=str, default='train')
    args = parser.parse_args()
    generate_idcard(args.info_dir, args.save_dir)
    # image_cut_save('./idcard_0_02.png', 445, 340, 540, 370, 'result.png')