diff --git a/tools/fonts/NanumSquare_acR.ttf b/tools/fonts/NanumSquare_acR.ttf new file mode 100644 index 0000000..3b43e02 Binary files /dev/null and b/tools/fonts/NanumSquare_acR.ttf differ diff --git a/tools/make_wifi_image.py b/tools/make_wifi_image.py new file mode 100644 index 0000000..7364fc8 --- /dev/null +++ b/tools/make_wifi_image.py @@ -0,0 +1,266 @@ +import argparse +from PIL import Image, ImageFont, ImageDraw +import json +import string +import random +import os +from tqdm import tqdm +import requests +import numpy as np + + +def random_text(): + string_pool = string.ascii_letters + string.digits + string.punctuation + string_pool2 = string.ascii_letters + string.digits + cnt = 0 + length = random.randint(8, 20) + result = '' + for _ in range(length): + if cnt < 3: + rand_char = random.choice(string_pool) + else: + rand_char = random.choice(string_pool2) + if rand_char in string.punctuation: + cnt += 1 + result += rand_char + + return result + +def get_ann(img_path,api_url="http://118.222.179.32:30000/ocr/") -> dict: + headers = {"secret": "Boostcamp0000"} + file_dict = {"file": open(img_path , "rb")} + response = requests.post(api_url, headers=headers, files=file_dict) + return response.json() + +def get_keys(ann): + id_list = ['ID', 'ID:', '아이디', '아이디:', 'NETWORK'] + pw_list = ['PW', 'PW:', '비밀번호', '비밀번호:', 'PASSCODE:', 'password', 'password:', 'PASSWORD', 'PASSWORD:'] + words = ann['ocr']['word'] + position_id = [0,0,0,0] + position_pw = [0,0,0,0] + for word in words: + if word['text'] in id_list: + position_id = word['points'] + if word['text'] in pw_list: + position_pw = word['points'] + + return position_id, position_pw + +def get_text_color(template_img, pos_id_key, gen_id_bbox): + pix = template_img.load() + x1, y1 = pos_id_key[0] + x2, y2 = pos_id_key[2] + x1_, y1_ = gen_id_bbox[0], gen_id_bbox[1] + x2_, y2_ = gen_id_bbox[2], gen_id_bbox[3] + + color_list = [] + for i in range(x1, x2+1): # ID bbox의 pixel값들 저장 + for j in range(y1, y2+1): + if not isinstance(pix[i,j], int) and len(pix[i,j]) == 4: + pix_c = (pix[i,j][0:3]) + else: + pix_c = pix[i,j] + color_list.append(pix_c) + + for i in range(x1_, x2_+1): # 생성할 위치의 bbox의 pixel값들 저장(배경 관련 rgb값을 더해주기 위한 목적) + for j in range(y1_, y2_+1): + if not isinstance(pix[i,j], int) and len(pix[i,j]) == 4: + pix_c = (pix[i,j][0:3]) + else: + pix_c = pix[i,j] + color_list.append(pix_c) + color_list_ = np.array(color_list) + mean_color = np.mean(color_list_, 0) # 배경색과 가까운 rgb 평균값을 얻을 수 있음. + d_list = [] + for color in color_list_: + D = np.sqrt(np.sum((color - mean_color) ** 2)) # 평균과 각 색과의 유클리드 거리 계산 + d_list.append(D) + ind = d_list.index(max(d_list)) # 평균 색에서 가장 거리가 먼 색을 text color라 가정 + return color_list[ind] + +def wifi_image_generation(): + + img_path = arg.img_path + font_path = arg.font_path + num_img = arg.num_img + text_id = arg.text_id + text_pw = arg.text_pw + text_size = 10 + text_color = (0, 0, 0) + + # Load image & annotation + image_id = 0 + box_id = 1 + coco_img_list = [] + coco_ann_list = [] + + img_list = os.listdir(img_path) + for img_name in tqdm(img_list): + + filename = os.path.join(img_path,img_name) + template_ann = get_ann(filename) + pos_id_key, pos_pw_key = get_keys(template_ann) + pos_x_gen_id = pos_id_key[1][0] + pos_x_gen_pw = pos_pw_key[1][0] + pos_y_gen_id = min(pos_id_key[0][1], pos_id_key[1][1]) + pos_y_gen_pw = min(pos_pw_key[0][1], pos_pw_key[1][1]) + + for n in range(num_img): + image_id += 1 + template_img = Image.open(filename) + img_width, img_height = template_img.size + + # 생성할 문자열의 좌상단 좌표, key와의 margin을 random값으로 지정(10~20 pixels) + pos_id_value = [pos_x_gen_id + random.randint(10, 20), pos_y_gen_id] + pos_pw_value = [pos_x_gen_id + random.randint(10, 20), pos_y_gen_pw] + + # Font Selection + font_list = os.listdir(font_path) + fontfile = font_path + '/' + random.choice(font_list) + text_size = max(pos_id_key[3][1] - pos_id_key[0][1], pos_id_key[2][1] - pos_id_key[1][1]) + font = ImageFont.truetype(fontfile, text_size) + + # Render the Text + image_editable = ImageDraw.Draw(template_img) # Convert the image to an editable format + if arg.text_id is None: + text_id = random_text() + if arg.text_pw is None: + text_pw = random_text() + + # Starting Coordinates:(0,0) in the upper left corner, Text, Text color: RGB, Font style + gen_id_bbox = image_editable.textbbox(pos_id_value, text_id, font=font) # (x0,y0,x2,y2) + gen_pw_bbox = image_editable.textbbox(pos_pw_value, text_pw, font=font) # (x0,y0,x2,y2) + text_color = get_text_color(template_img, pos_id_key, gen_id_bbox) + image_editable.text(gen_id_bbox, text_id, text_color, font=font) # ID + image_editable.text(gen_pw_bbox, text_pw, text_color, font=font) # PW + + # Export the result + export_dir = './gen_imgs' + if os.path.isdir(export_dir) == False: + os.mkdir(export_dir) + export_filename = img_name[0:img_name.rfind('.')] + '_' + str(n).zfill(2) +'.jpg' + template_img.save('/'.join([export_dir, export_filename])) + + # Make COCO format annotation + coco_img_list.append(coco_img(image_id, img_width, img_height, export_filename)) + box_id, ann_info = coco_ann(gen_id_bbox, gen_pw_bbox, box_id, image_id) + coco_ann_list.extend(ann_info) + + print(f'{len(img_list) * num_img} wifi images generated') + + coco_dataset = coco_template() + coco_dataset['images'] = coco_img_list + coco_dataset['annotations'] = coco_ann_list + anno_dir = './gen_imgs/annotations' + if os.path.isdir(anno_dir) == False: + os.mkdir(anno_dir) + with open('./gen_imgs/annotations/anno.json','w') as f: + json.dump(coco_dataset, f, indent=4) + print('COCO format annotations file generated') + +def coco_template(): + coco_dataset = { + "licenses": [ + { + "name": "", + "id": 0, + "url": "" + } + ], + "info": { + "contributor": "", + "date_created": "", + "description": "", + "url": "", + "version": "", + "year": "" + }, + "categories": [ + { + "id": 1, + "name": "ID", + "supercategory": "" + }, + { + "id": 2, + "name": "PW", + "supercategory": "" + } + ] + } + return coco_dataset + +def convert_coco_bbox(bbox): + width = bbox[2] - bbox[0] + height = bbox[3] - bbox[1] + + return [bbox[0], bbox[1], width, height], width * height + +def coco_img(image_id, width, height, filename): + img_info = { + "id": image_id, + "width": width, + "height": height, + "file_name": filename, + "license": 0, + "flickr_url": "", + "coco_url": "", + "date_captured": 0 + } + return img_info + +def coco_ann(id_bbox, pw_bbox, box_id, image_id): + + id_bbox, id_area = convert_coco_bbox(id_bbox) + pw_bbox, pw_area = convert_coco_bbox(pw_bbox) + + ann_info = [{ + "id": box_id, + "image_id": image_id, + "category_id": 1, # ID + "segmentation": [ + [ + id_bbox[0], id_bbox[1], + id_bbox[0] + id_bbox[2], id_bbox[1], + id_bbox[0] + id_bbox[2], id_bbox[1] + id_bbox[3], + id_bbox[0], id_bbox[1] + id_bbox[3] + ] + ], + "area": id_area, + "bbox": id_bbox, + "iscrowd": 0, + "attributes": { + "occluded": False + } + }, + { + "id": box_id+1, + "image_id": image_id, + "category_id": 2, # PW + "segmentation": [ + [ + pw_bbox[0], pw_bbox[1], + pw_bbox[0] + pw_bbox[2], pw_bbox[1], + pw_bbox[0] + pw_bbox[2], pw_bbox[1] + pw_bbox[3], + pw_bbox[0], pw_bbox[1] + pw_bbox[3] + ] + ], + "area": pw_area, + "bbox": pw_bbox, + "iscrowd": 0, + "attributes": { + "occluded": False + } + }] + return box_id + 2, ann_info + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument('-img','--img_path',default='./templates') + parser.add_argument('-font','--font_path',default='./fonts') + parser.add_argument('-num','--num_img',default=3) + parser.add_argument('-id','--text_id',default=None) + parser.add_argument('-pw','--text_pw',default=None) + arg = parser.parse_args() + + wifi_image_generation() \ No newline at end of file diff --git a/tools/templates/template1.jpeg b/tools/templates/template1.jpeg new file mode 100644 index 0000000..2e25d95 Binary files /dev/null and b/tools/templates/template1.jpeg differ