def save_keypoints_detection(pred_keypoints, metainfo, class_names, skeleton_lines):
result_dir=os.path.join('result','detection')
touchdir(result_dir)
image_name = metainfo['name']
image = Image.open(image_name)
image_array = np.array(image, dtype='uint8')
gt_keypoints = metainfo['pts']
# form up gt keypoints & predict keypoints dict
gt_keypoints_dict = {}
pred_keypoints_dict = {}
for i, keypoint in enumerate(gt_keypoints):
gt_keypoints_dict[class_names[i]] = (keypoint[0], keypoint[1], 1.0)
for i, keypoint in enumerate(pred_keypoints):
pred_keypoints_dict[class_names[i]] = (keypoint[0], keypoint[1], keypoint[2])
# render gt and predict keypoints skeleton on image
image_array = render_skeleton(image_array, gt_keypoints_dict, skeleton_lines, colors=(255, 255, 255))
image_array = render_skeleton(image_array, pred_keypoints_dict, skeleton_lines)
image = Image.fromarray(image_array)
# here we handle the RGBA image
if(len(image.split()) == 4):
r, g, b, a = image.split()
image = Image.merge("RGB", (r, g, b))
image.save(os.path.join(result_dir, image_name.split(os.path.sep)[-1]))
return
def compute_AP_COCO(annotation_records, gt_classes_records, pred_classes_records, class_names, show_result=True):
'''
Compute MSCOCO AP list on AP 0.5:0.05:0.95
'''
iou_threshold_list = np.arange(0.50,0.95,0.05)
APs = {}
for iou_threshold in iou_threshold_list:
iou_threshold = round(iou_threshold, 2)
mAP = compute_mAP_PascalVOC(annotation_records, gt_classes_records, pred_classes_records, class_names, iou_threshold, show_result=False)
APs[iou_threshold] = round(mAP, 6)
#get overall AP percentage value
AP = np.mean(list(APs.values()))
if show_result:
'''
Draw MS COCO AP plot
'''
touchdir('result')
window_title = "MSCOCO AP on different IOU"
plot_title = "COCO AP = {0:.2f}%".format(AP)
x_label = "Average Precision"
output_path = os.path.join('result','COCO_AP.jpg')
draw_plot_func(APs, len(APs), window_title, plot_title, x_label, output_path, to_show=False, plot_color='royalblue', true_p_bar='')
print('\nMS COCO AP evaluation')
for (iou_threshold, AP_value) in APs.items():
print('IOU %.2f: AP %f' % (iou_threshold, AP_value))
print('total AP: %f' % (AP))
#return AP percentage value
return AP
def draw_rec_prec(rec, prec, mrec, mprec, class_name, ap):
"""
Draw plot
"""
plt.plot(rec, prec, '-o')
# add a new penultimate point to the list (mrec[-2], 0.0)
# since the last line segment (and respective area) do not affect the AP value
area_under_curve_x = mrec[:-1] + [mrec[-2]] + [mrec[-1]]
area_under_curve_y = mprec[:-1] + [0.0] + [mprec[-1]]
plt.fill_between(area_under_curve_x, 0, area_under_curve_y, alpha=0.2, edgecolor='r')
# set window title
fig = plt.gcf() # gcf - get current figure
fig.canvas.set_window_title('AP ' + class_name)
# set plot title
plt.title('class: ' + class_name + ' AP = {}%'.format(ap*100))
#plt.suptitle('This is a somewhat long figure title', fontsize=16)
# set axis titles
plt.xlabel('Recall')
plt.ylabel('Precision')
# optional - set axes
axes = plt.gca() # gca - get current axes
axes.set_xlim([0.0,1.0])
axes.set_ylim([0.0,1.05]) # .05 to give some extra space
# Alternative option -> wait for button to be pressed
#while not plt.waitforbuttonpress(): pass # wait for key display
# Alternative option -> normal display
#plt.show()
# save the plot
rec_prec_plot_path = os.path.join('result','classes')
touchdir(rec_prec_plot_path)
fig.savefig(os.path.join(rec_prec_plot_path, class_name + ".jpg"))
plt.cla() # clear axes for next plot
def compute_AP_COCO_Scale(annotation_records, scale_gt_classes_records, pred_classes_records, class_names):
'''
Compute MSCOCO AP on different scale object: small, medium, large
'''
scale_APs = {}
for scale_key in ['small','medium','large']:
gt_classes_records = scale_gt_classes_records[scale_key]
scale_AP = compute_AP_COCO(annotation_records, gt_classes_records, pred_classes_records, class_names, show_result=False)
scale_APs[scale_key] = round(scale_AP, 4)
#get overall AP percentage value
scale_mAP = np.mean(list(scale_APs.values()))
'''
Draw Scale AP plot
'''
touchdir('result')
window_title = "MSCOCO AP on different scale"
plot_title = "scale mAP = {0:.2f}%".format(scale_mAP)
x_label = "Average Precision"
output_path = os.path.join('result','COCO_scale_AP.jpg')
draw_plot_func(scale_APs, len(scale_APs), window_title, plot_title, x_label, output_path, to_show=False, plot_color='royalblue', true_p_bar='')
'''
Draw Scale Object Sum plot
'''
for scale_key in ['small','medium','large']:
gt_classes_records = scale_gt_classes_records[scale_key]
gt_classes_sum = {}
for _, class_name in enumerate(class_names):
# summarize the gt object number for every class on different scale
gt_classes_sum[class_name] = np.sum(len(gt_classes_records[class_name])) if class_name in gt_classes_records else 0
total_sum = np.sum(list(gt_classes_sum.values()))
window_title = "{} object number".format(scale_key)
plot_title = "total {} object number = {}".format(scale_key, total_sum)
x_label = "Object Number"
output_path = os.path.join('result','{}_object_number.jpg'.format(scale_key))
draw_plot_func(gt_classes_sum, len(gt_classes_sum), window_title, plot_title, x_label, output_path, to_show=False, plot_color='royalblue', true_p_bar='')
print('\nMS COCO AP evaluation on different scale')
for (scale, AP_value) in scale_APs.items():
print('%s scale: AP %f' % (scale, AP_value))
print('total AP: %f' % (scale_mAP))
def dump_saved_model(self, saved_model_path):
model = self.inference_model
touchdir(saved_model_path)
tf.keras.experimental.export_saved_model(model, saved_model_path)
print('export inference model to %s' % str(saved_model_path))
def get_prediction_class_records(model, model_format, annotation_records,
anchors, class_names, model_image_size,
conf_threshold, save_result):
'''
Do the predict with YOLO model on annotation images to get predict class dict
predict class dict would contain image_name, coordinary and score, and
sorted by score:
pred_classes_records = {
'car': [
['00001.jpg','94,115,203,232',0.98],
['00002.jpg','82,64,154,128',0.93],
...
],
...
}
'''
if model_format == 'MNN':
#MNN inference engine need create session
session = model.createSession()
pred_classes_records = {}
pbar = tqdm(total=len(annotation_records), desc='Eval model')
for (image_name, gt_records) in annotation_records.items():
image = Image.open(image_name)
image_array = np.array(image, dtype='uint8')
# support of tflite model
if model_format == 'TFLITE':
pred_boxes, pred_classes, pred_scores = yolo_predict_tflite(
model, image, anchors, len(class_names), conf_threshold)
# support of MNN model
elif model_format == 'MNN':
pred_boxes, pred_classes, pred_scores = yolo_predict_mnn(
model, session, image, anchors, len(class_names),
conf_threshold)
# support of TF 1.x frozen pb model
elif model_format == 'PB':
pred_boxes, pred_classes, pred_scores = yolo_predict_pb(
model, image, anchors, len(class_names), model_image_size,
conf_threshold)
# normal keras h5 model
elif model_format == 'H5':
pred_boxes, pred_classes, pred_scores = yolo_predict_keras(
model, image, anchors, len(class_names), model_image_size,
conf_threshold)
else:
raise ValueError('invalid model format')
#print('Found {} boxes for {}'.format(len(pred_boxes), image_name))
pbar.update(1)
if save_result:
gt_boxes, gt_classes, gt_scores = transform_gt_record(
gt_records, class_names)
result_dir = os.path.join('result', 'detection')
touchdir(result_dir)
colors = get_colors(class_names)
image_array = draw_boxes(image_array,
gt_boxes,
gt_classes,
gt_scores,
class_names,
colors=None,
show_score=False)
image_array = draw_boxes(image_array, pred_boxes, pred_classes,
pred_scores, class_names, colors)
image = Image.fromarray(image_array)
# here we handle the RGBA image
if (len(image.split()) == 4):
r, g, b, a = image.split()
image = Image.merge("RGB", (r, g, b))
image.save(
os.path.join(result_dir,
image_name.split(os.path.sep)[-1]))
# Nothing detected
if pred_boxes is None or len(pred_boxes) == 0:
continue
for box, cls, score in zip(pred_boxes, pred_classes, pred_scores):
pred_class_name = class_names[cls]
xmin, ymin, xmax, ymax = box
coordinate = "{},{},{},{}".format(xmin, ymin, xmax, ymax)
#append or add predict class item
if pred_class_name in pred_classes_records:
pred_classes_records[pred_class_name].append(
[image_name, coordinate, score])
else:
pred_classes_records[pred_class_name] = list(
[[image_name, coordinate, score]])
# sort pred_classes_records for each class according to score
for pred_class_list in pred_classes_records.values():
pred_class_list.sort(key=lambda ele: ele[2], reverse=True)
pbar.close()
return pred_classes_records
def eval_PCK(model, model_format, eval_dataset, class_names, score_threshold, normalize, conf_threshold, save_result=False, skeleton_lines=None):
if model_format == 'MNN':
#MNN inference engine need create session
session = model.createSession()
succeed_dict = {class_name: 0 for class_name in class_names}
fail_dict = {class_name: 0 for class_name in class_names}
accuracy_dict = {class_name: 0. for class_name in class_names}
# init output list for coco result json generation
# coco keypoints result is a list of following format dict:
# {
# "image_id": int,
# "category_id": int,
# "keypoints": [x1,y1,v1,...,xk,yk,vk],
# "score": float
# }
#
output_list = []
count = 0
batch_size = 1
pbar = tqdm(total=eval_dataset.get_dataset_size(), desc='Eval model')
for image_data, gt_heatmap, metainfo in eval_dataset.generator(batch_size, 8, sigma=1, is_shuffle=False, with_meta=True):
# fetch validation data from generator, which will crop out single person area, resize to input_size and normalize image
count += batch_size
if count > eval_dataset.get_dataset_size():
break
# support of tflite model
if model_format == 'TFLITE':
heatmap = hourglass_predict_tflite(model, image_data)
# support of MNN model
elif model_format == 'MNN':
heatmap = hourglass_predict_mnn(model, session, image_data)
# support of TF 1.x frozen pb model
elif model_format == 'PB':
heatmap = hourglass_predict_pb(model, image_data)
# support of ONNX model
elif model_format == 'ONNX':
heatmap = hourglass_predict_onnx(model, image_data)
# normal keras h5 model
elif model_format == 'H5':
heatmap = hourglass_predict_keras(model, image_data)
else:
raise ValueError('invalid model format')
heatmap_size = heatmap.shape[0:2]
# get predict keypoints from heatmap
pred_keypoints = post_process_heatmap(heatmap, conf_threshold)
pred_keypoints = np.array(pred_keypoints)
# get ground truth keypoints (transformed)
metainfo = metainfo[0]
gt_keypoints = metainfo['tpts']
# calculate succeed & failed keypoints for prediction
result_list = keypoint_accuracy(pred_keypoints, gt_keypoints, score_threshold, normalize)
for i, class_name in enumerate(class_names):
if result_list[i] == 0:
fail_dict[class_name] = fail_dict[class_name] + 1
elif result_list[i] == 1:
succeed_dict[class_name] = succeed_dict[class_name] + 1
# revert predict keypoints back to origin image size
reverted_pred_keypoints = revert_keypoints(pred_keypoints, metainfo, heatmap_size)
# get coco result dict with predict keypoints and image info
result_dict = get_result_dict(reverted_pred_keypoints, metainfo)
# add result dict to output list
output_list.append(result_dict)
if save_result:
# render keypoints skeleton on image and save result
save_keypoints_detection(reverted_pred_keypoints, metainfo, class_names, skeleton_lines)
pbar.update(batch_size)
pbar.close()
# save to coco result json
touchdir('result')
json_fp = open(os.path.join('result','keypoints_result.json'), 'w')
json_str = json.dumps(output_list)
json_fp.write(json_str)
json_fp.close()
# calculate accuracy for each class
for i, class_name in enumerate(class_names):
accuracy_dict[class_name] = succeed_dict[class_name] * 1.0 / (succeed_dict[class_name] + fail_dict[class_name])
#get PCK accuracy from succeed & failed keypoints
total_succeed = np.sum(list(succeed_dict.values()))
total_fail = np.sum(list(fail_dict.values()))
total_accuracy = total_succeed * 1.0 / (total_fail + total_succeed)
if save_result:
'''
Draw PCK plot
'''
window_title = "PCK evaluation"
plot_title = "PCK@{0} score = {1:.2f}%".format(score_threshold, total_accuracy)
x_label = "Accuracy"
output_path = os.path.join('result','PCK.jpg')
draw_plot_func(accuracy_dict, len(accuracy_dict), window_title, plot_title, x_label, output_path, to_show=False, plot_color='royalblue', true_p_bar='')
return total_accuracy, accuracy_dict
