DocumentAIKit/idcardgenerator/idcard_det_rec_generate_back.py

import argparse
import os.path
import random

import cv2
from PIL import Image, ImageDraw, ImageFont
from tqdm import tqdm


def add_space(str):
    temp_str = ''
    if len(str) <= 6:
        num_of_space = (3 * (6 - len(str))) / (len(str) - 1)
        for i in range(0, len(str) - 1):
            temp_str += str[i]
            for j in range(0, int(num_of_space)):
                temp_str += ' '
        temp_str += str[len(str) - 1]
        return temp_str


def generate_idcard(name_dir, bingyi_dir, addr_dir, id_dir, save_dir, generate_num, bg_dir, font_dir):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    fp_name = open(name_dir, 'r', encoding='utf-8')
    fp_bingyi = open(bingyi_dir, 'r', encoding='utf-8')
    fp_addr= open(addr_dir, 'r', encoding='utf-8')
    fp_id = open(id_dir, 'r', encoding='utf-8')

    # 姓名列表
    name_list = []
    for name in fp_name.readlines():
        name_list.append(name.rstrip('\n').replace(" ", ''))
    # print(name_list)
    # 兵役選項
    bingyi_list = []
    for bingyi in fp_bingyi.readlines():
        bingyi_list.append(bingyi[2:].rstrip('\n'))

    # 出生地及住址
    addr_list = []
    for addr in fp_addr.readlines():
        addr = addr.rstrip('\n')
        ll = addr.split(',')
        addr_temp = ll[1] + ll[0]
        addr_list.append(addr_temp)

    # 編號
    id_list = []
    for id in fp_id.readlines():
        id_list.append(id.rstrip('\n'))
    cnt = 2401
    for i in tqdm(range(0, generate_num)):
        name_father_id = random.randint(0, len(name_list) - 1)
        name_mother_id = random.randint(0, len(name_list) - 1)
        name_couple_id = random.randint(0, len(name_list) - 1)
        bingyi_id = random.randint(0, len(bingyi_list) - 1)
        addr_id = random.randint(0, len(addr_list) - 1)
        id_id = random.randint(0, len(id_list) - 1)

        # 父母親及配偶姓名
        name_father = name_list[name_father_id]
        name_mother = name_list[name_mother_id]
        name_couple = name_list[name_couple_id]

        # 兵役
        bingyi = bingyi_list[bingyi_id]

        # 住址
        addr = addr_list[addr_id]

        n = random.randint(0, 1)
        if n == 0:
            # 出生地
            birth_place = '臺灣省' + addr[0:3]
        # 县或市为出生地
        else:
            birth_place = addr[0:3]

        # 編號
        id = id_list[id_id]

        # 获取背景图列表
        bg_img_name_list = os.listdir(bg_dir)
        bg_path_list = []
        for bg_img_name in bg_img_name_list:
            bg_path_list.append(bg_dir + '/' + bg_img_name)
        # 每次随机取背景图
        bg_path = bg_path_list[random.randint(0, len(bg_path_list) - 1)]
        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_font_size = 35
        # 姓名中加入间隔
        name_father = add_space(name_father)
        name_mother = add_space(name_mother)
        name_couple = add_space(name_couple)

        temp_str = ''
        for i in range(0, len(addr) - 1):
            temp_str += addr[i]
            temp_str += ' '
        temp_str += addr[len(addr) - 1]

        addr = temp_str

        font_name_list = os.listdir(font_dir)
        font_path_list = []
        for font_name in font_name_list:
            font_path_list.append(font_dir + '/' + font_name)
        # 每次随机字体
        font_path = font_path_list[random.randint(0, len(font_path_list) - 1)]

        name_font = ImageFont.truetype(font_path, name_font_size)  # 设置名字字体和大小
        addr_font = ImageFont.truetype(font_path, 23)
        id_font = ImageFont.truetype('OpticalBBold.otf', 28)
        # # 计算出要写入的文字占用的像素
        # # 创建一个可以在给定图像上绘图的对象
        draw = ImageDraw.Draw(bg)

        # 设置配偶及役别数据项为空
        rand_bingyi = random.randint(0, 1)
        rang_spouse = random.randint(0, 1)
        if rand_bingyi == 0:
            bingyi=''
        if rang_spouse == 0:
            name_couple = ''

        draw.text((110, 17), name_father, fill="#444444", font=name_font)
        draw.text((415, 17), name_mother, fill="#444444", font=name_font)
        draw.text((115, 82), name_couple, fill="#444444", font=name_font)
        draw.text((415, 82), bingyi, fill="#444444", font=name_font)
        draw.text((115, 147), birth_place, fill="#444444", font=name_font)
        draw.text((115, 210), addr[0:24], fill="#444444", font=addr_font)
        draw.text((115, 240), addr[24:], fill="#444444", font=addr_font)
        draw.text((455, 377), id, fill="#42A329", font=id_font)

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

        path_list.append(img_path)

        boxes = []

        box_name_father = [105, 17, 305, 62]
        box0 = []
        box0.append(box_name_father)
        box0.append(name_father.replace(' ', ''))
        box0.append('father_name')
        boxes.append(box0)

        box_name_mother = [410, 17, 610, 62]
        box1 = []
        box1.append(box_name_mother)
        box1.append(name_mother.replace(' ', ''))
        box1.append('mother_name')
        boxes.append(box1)

        box_name_spouse = [110, 82, 310, 127]
        box2 = []
        box2.append(box_name_spouse)
        box2.append(name_couple.replace(' ', ''))
        box2.append('spouse_name')
        boxes.append(box2)

        box_birth_place = [110, 147, 335, 190]
        box3 = []
        box3.append(box_birth_place)
        box3.append(birth_place.replace(' ', ''))
        box3.append('birth_place')
        boxes.append(box3)

        box_bingyi = [410, 82, 610, 125]
        box4 = []
        box4.append(box_bingyi)
        box4.append(bingyi.replace(' ', ''))
        box4.append('level')
        boxes.append(box4)

        box_id = [450, 377, 655, 410]
        box5 = []
        box5.append(box_id)
        box5.append(id.replace(' ', ''))
        box5.append('number')
        boxes.append(box5)

        num = 1
        # 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).split('/')[-1] + '\t['
            for box in boxes:
                if jishu <= 6:
                    if jishu == 6:
                        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()
    fp_name.close()
    fp_bingyi.close()
    fp_addr.close()
    fp_id.close()
    return


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('--name_dir', type=str, default='./materials/name_cht.txt')
    parser.add_argument('--bingyi_dir', type=str, default='./materials/bingyi.txt')
    parser.add_argument('--addr_dir', type=str, default='./materials/addr.txt')
    parser.add_argument('--id_dir', type=str, default='./materials/id_back.txt')
    parser.add_argument('--save_dir', type=str, default=r'C:\Users\KDAN\Desktop\workspace\文字检测数据集（台湾身份证）\tw_idcard_back_det_0215\imgs')
    parser.add_argument('--bg_dir', type=str, default='./test_bgs')
    parser.add_argument('--generate_num', type=int, default=700)
    parser.add_argument('--font_dir', type=str, default='./fonts')

    args = parser.parse_args()
    generate_idcard(args.name_dir, args.bingyi_dir, args.addr_dir, args.id_dir, args.save_dir, args.generate_num, args.bg_dir, args.font_dir)
    # image_cut_save('./idcard_0_02.png', 445, 340, 540, 370, 'result.png')