DocumentAIKit/idcardgenerator/generate_roi_imgs.py

import argparse
import os.path
import cv2
from PIL import Image, ImageDraw, ImageFont
from tqdm import tqdm


def generate_idcard(info_dir, save_dir):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    f1 = open('materials/rec_gt.txt', 'w+', encoding='utf-8')
    with open(info_dir, 'r', encoding='utf-8') as fp:
        text_lines = fp.readlines()
        cnt = 0
        for text_line in tqdm(text_lines):
            info = text_line.rstrip('\n').split(',')
            if cnt < len(text_lines) * 0.2:
                bg_path = "bgs/tw_idcard_bg_1.jpg"
            elif len(text_lines) * 0.2 <= cnt < len(text_lines) * 0.6:
                bg_path = "bgs/tw_idcard_bg_2.png"
            else:
                bg_path = "bgs/tw_idcard_bg_3.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[0]  # 写入的文字
            content_name = name
            # print(name)

            name_font_size = 45
            temp_str = ''
            if len(name) <= 6:
                num_of_space = (2 * (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/TW-Sung-98-1-2.ttf"  # 一个字体文件
            name_font = ImageFont.truetype(font, name_font_size)  # 设置名字字体和大小
            # 计算出要写入的文字占用的像素
            # 创建一个可以在给定图像上绘图的对象
            draw = ImageDraw.Draw(bg)

            draw.text((5, 0), name, fill="#07090E", font=name_font)

            # 保存画布
            img_path = save_dir + '/name_' + str(cnt) + '.png'
            f1.write(img_path + '\t' + content_name + '\n')
            cnt += 1
            bg.save(img_path, "PNG")
    f1.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--info_dir', type=str, default='./materials/name.txt')
    parser.add_argument('--save_dir', type=str, default='./imgs')
    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')