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')