#!/usr/bin/env python # coding: utf-8 # In[26]: import os import json from utils import YOLO import pickle import numpy as np import shutil from PIL import Image, ImageDraw, ImageFont import argparse import random import colorsys import shutil # In[27]: def generate_colors(classes): # added color_file.txt so we can customize bad colors hsv_tuples = [(float(x) / classes, 1., 1.) for x in range(classes)] colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples)) colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors)) random.seed(10101) # Fixed seed for consistent colors across runs. random.shuffle(colors) # Shuffle colors to decorrelate adjacent classes. random.seed(None) # Reset seed to default. return colors # In[27]: def main_(args): emulator_res_folder_path = args.input_folder # In[28]: image_size = int(args.input_size) output_num = int(args.output_num) score_thresh = float(args.score) keep_aspect_ratio = args.keep_aspect_ratio if keep_aspect_ratio.lower() == "true": keep_aspect_ratio = True else: keep_aspect_ratio = False # In[29]: yolo = YOLO(image_size, output_num, score_thresh, keep_aspect_ratio) fnt = ImageFont.truetype(args.font_path, 15) # In[30]: temp_file_path = os.path.join(emulator_res_folder_path,"detection_results.txt") tot_f = open(temp_file_path,"w+") for folder in os.listdir(emulator_res_folder_path): img_res_folder = os.path.join(emulator_res_folder_path, folder) if not os.path.isdir(img_res_folder): continue print(img_res_folder) data_lst = [] w, h = 0, 0 img_path = "" for file in os.listdir(img_res_folder): if "temp" in file: output_path = os.path.join(img_res_folder, file) data = np.loadtxt(output_path) data_lst.append(data) if (file.endswith(".png") or file.endswith(".jpg") or file.endswith(".jpeg")) and not "_thresh_" in file: img_path_ = os.path.join(img_res_folder, file) if "thresh" in img_path_: continue else: img_path = img_path_ print(img_path) w, h = Image.open(img_path).size if len(data_lst) > 0 and w > 0 and h > 0: res_lst = yolo.detect_image(data_lst, False, w, h) # yolo.close_session() summary_path = os.path.join(img_res_folder, "result.txt") with open(summary_path, 'w+') as f: for res in res_lst: f.write(" ".join(res) +'\n') #draw on img #print(res_lst) dict_class_info = {} for res in res_lst: name, score, x1, y1, x2, y2 = res score = float(score) x1 = int(x1) x2 = int(x2) y1 = int(y1) y2 = int(y2) w = int(x2-x1) h = int(y2-y1) if name not in dict_class_info.keys(): dict_class_info[name] = [] dict_class_info[name].append([score, x1, y1, x2, y2, w, h]) colors = generate_colors(len(dict_class_info.keys())) thresh_list = [] if score_thresh == 0.01: thresh_list = [0.01, 0.1, 0.15, 0.2, 0.25, 0.3] else: thresh_list = [score_thresh] for thresh in thresh_list: print("score thresh: {}".format(thresh)) img = Image.open(img_path) draw = ImageDraw.Draw(img) for i in range(len(dict_class_info.keys())): name = list(dict_class_info.keys())[i] for location in dict_class_info[name]: score, x1, y1, x2, y2, w, h = location print("{:>10} {:10.4f} {} {} {} {}".format(name, score, x1, y1, x2, y2)) if score < thresh: continue label = '{} {:.2f}'.format(name,score) labelsize = draw.textsize(label,font=fnt) thickness = int((img.size[0] + img.size[1]) / 300) for k in range(thickness): print(colors[i]) draw.rectangle([x1+k, y1+k, x2+k, y2+k],outline=colors[i][0]) if thresh == 0.2: data = img_path.split('/')[4]+" "+name+' {} {} {} {} {}\n'.format(score, x1, y1, x2, y2) tot_f.write(data) draw.rectangle([x1, y2, x1+2*thickness+labelsize[0], y2+thickness+labelsize[1]],fill=colors[i][0],outline=colors[i][0]) draw.text([x1+2*thickness, y2+thickness], label, fill=0, font=fnt) img_name = img_path.split("/")[-1].split(".")[0] ext = img_path.split(".")[-1] new_img_name = img_name + "_thresh_{}".format(thresh) + "." + ext new_img_path = os.path.join(img_res_folder, new_img_name) img.save(new_img_path) for i in range(len(dict_class_info.keys())): name = list(dict_class_info.keys())[i] class_folder = os.path.join(img_res_folder, name) if os.path.exists(class_folder): shutil.rmtree(class_folder) os.mkdir(class_folder) for thresh in thresh_list: print("score thresh: {}".format(thresh)) img = Image.open(img_path) draw = ImageDraw.Draw(img) for location in dict_class_info[name]: score, x1, y1, x2, y2, w, h = location print("{:>10} {:10.4f} {} {} {} {}".format(name, score, x1, y1, x2, y2)) if score < thresh: continue label = '{} {:.2f}'.format(name,score) labelsize = draw.textsize(label,font=fnt) thickness = int((img.size[0] + img.size[1]) / 300) for k in range(thickness): draw.rectangle([x1+k, y1+k, x2+k, y2+k],outline=colors[i][0]) draw.rectangle([x1, y2, x1+2*thickness+labelsize[0], y2+thickness+labelsize[1]],fill=colors[i][0],outline=colors[i][0]) draw.text([x1+2*thickness, y2+thickness], label, fill=0,font=fnt) img_name = img_path.split("/")[-1].split(".")[0] ext = img_path.split(".")[-1] new_img_name = img_name + "_thresh_{}".format(thresh) + "." + ext new_img_path = os.path.join(class_folder, new_img_name) print(new_img_path) img.save(new_img_path) print("=" * 10) tot_f.close() # In[ ]: if __name__ == "__main__": argparser = argparse.ArgumentParser( description="convert emulator result to box and score" ) argparser.add_argument( '-i', '--input_folder', default="/data1/emulator", help="path of input folder containing emulator output" ) argparser.add_argument( '-is', '--input_size', default=416, help="yolo input size" ) argparser.add_argument( '-s', '--score', default=0.01, help="score threshold" ) argparser.add_argument( '-f', '--font_path', default="/workspace/scripts/utils/Microsoft Sans Serif.ttf", help="path of font" ) argparser.add_argument( '-ot', '--output_num', help="the number of output nodes" ) argparser.add_argument( '-k', '--keep_aspect_ratio', default="True", help="keep_aspect_ratio" ) args = argparser.parse_args() main_(args)