271 lines
8.3 KiB
Python
271 lines
8.3 KiB
Python
#!/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)
|
|
|