MindSpore:YOLOv3人体目标检测模型实现(四)
5.模型评估
首先是导包以及设置context:
from mindspore import context
from dataset.voc2012_dataset import create_voc2012_dataset
from model.yolo import YOLOV3DarkNet53
from utils.utils_yolov3 import DetectionEngine, load_yolov3
from yolov3_eval_conf import EvalYOLOv3Conf
# set context
context.set_context(mode=context.GRAPH_MODE,
device_target="GPU", save_graphs=False)
# set_graph_kernel_context
context.set_context(enable_graph_kernel=True)
context.set_context(graph_kernel_flags="--enable_parallel_fusion "
"--enable_trans_op_optimize "
"--disable_cluster_ops=ReduceMax,Reshape "
"--enable_expand_ops=Conv2D")
# Set mempool block size for improving memory utilization, which will not take effect in GRAPH_MODE
if context.get_context("mode") == context.PYNATIVE_MODE:
context.set_context(mempool_block_size="31GB")
# config
config = EvalYOLOv3Conf()
评估模型时的配置如下(也可见 附件\yolov3_eval_conf.py):
class EvalYOLOv3Conf():
def __init__(self):
# ---- dataset ----
self.data_path = "VOCdevkit/VOC2012/"
self.data_usage = "my_person_val" # 评估数据集
self.data_training = False
self.num_classes = 1
self.class_to_idx = {}
self.anchor_scales = [[15, 38],
[34, 86],
[84, 127],
[51, 192],
[91, 257],
[173, 195],
[142, 319],
[221, 339],
[351, 365]]
self.batch_size = 48 # 评估时的batch
self.max_box = 32
# test
self.test_img_shape = [416, 416] # 图片缩放大小
# ---- Model ----
self.out_channel = 3 * ( 5 + self.num_classes)
self.keep_detect = True
self.ckpt_path = "./train_ckpt/yolov3-???.ckpt" # 评估的模型
# ---- Detect ----
self.nms_thresh = 0.5 # nms 算法去重的IoU阈值
self.eval_ignore_threshold = 0.001 # 检测置信度阈值
self.detcte_result_dir = "det_res/" # 分类检测框的临时结果存放
self.image_id_idx = [] # xml编号的列表,与图片序号对应,eval_yolov3.py中设置
self.anno_path = "VOCdevkit/VOC2012/Annotations/{}.xml" # Annotations目录的位置,注意最后的格式
self.val_path = "VOCdevkit/VOC2012/ImageSets/Main/my_person_val.txt" # 模型评估xml文件
读数据集:
# dataset
voc2012_dat, data_size = create_voc2012_dataset(config, 2)
config.steps_per_epoch = int(data_size / config.batch_size)
image_id_idx = {}
with open(config.val_path) as f:
lines = f.readlines()
for i, line in enumerate(lines):
image_id_idx[i] = line.strip()
config.image_id_idx = image_id_idx
print("dataset size: ",data_size)
print("bath num in 1 epoch: ", config.steps_per_epoch)
定义网络和检测器:
# network
network = YOLOV3DarkNet53(is_training=config.data_training, config=config)
load_yolov3(network,config.ckpt_path)
network.set_train(False)
# init detection engine
detection = DetectionEngine(config)
这里我们不需要损失函数,所以只用了网络结构YOLOV3DarkNet53,然后用load_yolov3将前面训练好的模型加载进来。
DetectionEngine类主要处理模型前向传播后的输出,包括置信度阈值筛选、NMS算法去除重叠等等工作,最后计算AP的工作也是它完成,对这些计算感兴趣的朋友可以看 附件\utils\utils_yolov3.py 以及 附件\utils\eval_utils.py。
最后对模型进行测试评估:
print('Start inference....')
for i, data in enumerate(voc2012_dat.create_dict_iterator(num_epochs=1)):
image = data["image"]
image_shape = data["image_shape"]
image_id = data["img_id"]
prediction = network(image)
output_big, output_me, output_small = prediction
output_big = output_big.asnumpy()
output_me = output_me.asnumpy()
output_small = output_small.asnumpy()
image_id = image_id.asnumpy()
image_shape = image_shape.asnumpy()
detection.detect([output_small, output_me, output_big], config.batch_size, image_shape, image_id)
if i % 2 == 0:
print('Processing... {:.2f}% '.format(i * config.batch_size / data_size * 100))
print("Finish")
print('Calculating mAP...')
detection.do_nms_for_results()
result_file_path = detection.write_result()
print('result file path: ', result_file_path)
detection.get_eval_result()
for k, t in detection.eval_res.items():
print(k, " AP : ", t['ap'])
import matplotlib.pyplot as plt
plt.title("P-R curve",fontsize=14)
plt.xlabel("recall", fontsize=14)
plt.ylabel("precision", fontsize=14)
res = detection.eval_res['person']
plt.plot(res['prec'], res['rec'])
plt.savefig("p_r.png")
最后我设置检测置信度的阈值为 0.01,nms 算法中的 IoU 阈值为 0.5时,得到的结果是AP为0.609,P-R曲线如下: