在YOLO-Fastest和YOLO-FastestV2中添加评价指标acc@iou

1. 引言

         YOLO-Fastest是第一个把YOLO系列目标检测算法的参数量降到1M以内的检测算法，模型精度相对于其他YOLO系列有所降低，但是参数量大幅下降，不足1M，因此很适合移动平台的模型部署。本文将目标检测算法评价指标acc@iou添加到YOLO-Fastest和YOLO-FastestV2中的评价指标中。

        YOLO-Fastest官方GitHub：https://github.com/dog-qiuqiu/Yolo-Fastest

        YOLO-FastestV2官方GitHub：https://github.com/dog-qiuqiu/Yolo-FastestV2

2. 步骤

        由于V1和V2源码非常相似，本文只对V2版本的代码进行改进

1. 在utils/utils.py文件中，新建函数evaluation_iou，此函数是evaluation函数的改进版本，新增了计算acc@iou的模块，并将验证集的这个指标输出。代码如下：

def evaluation_iou(val_dataloader, cfg, model, device, conf_thres = 0.01, nms_thresh = 0.4, iou_thres = 0.5):

    labels = []
    sample_metrics = []  # List of tuples (TP, confs, pred)
    pbar = tqdm(val_dataloader)


    global thres_lst
    global hit_lst
    global num
    thres_lst = [0.5,0.55,0.6,0.65,0.7,0.75,0.8,0.85,0.9,0.95]
    hit_lst = {}
    num = 0
    for thres in thres_lst:
        hit_lst[thres] = 0

    for imgs, targets in pbar:  #images (64,3,224,224) targets(79,6)
        imgs = imgs.to(device).float() / 255.0
        targets = targets.to(device)       

        # Extract labels
        labels += targets[:, 1].tolist()  #读取所有的分类
        # Rescale target
        targets[:, 2:] = xywh2xyxy(targets[:, 2:])
        targets[:, 2:] *= torch.tensor([cfg["width"], cfg["height"], cfg["width"], cfg["height"]]).to(device)

        #对预测的anchorbox进行nms处理
        with torch.no_grad():
            preds = model(imgs)

            #特征图后处理:生成anchorbox
            output = handel_preds(preds, cfg, device)
            output_boxes = non_max_suppression(output, conf_thres = conf_thres, iou_thres = nms_thresh)

        sample_metrics += get_batch_statistics_iou(output_boxes, targets, iou_thres, device)
        pbar.set_description("Evaluation model:") 

    acc_mean = 0
    for thres in hit_lst.keys():
        hit_lst[thres] /= num
        acc_mean += hit_lst[thres]
    acc_mean /= len(hit_lst)

    acc_log = {}
    for p in ['0.5','0.75','0.95','acc_mean']:
        acc_log[p] = 0
    acc_log['0.5'] = hit_lst[0.5]
    acc_log['0.75'] = hit_lst[0.75]
    acc_log['0.95'] = hit_lst[0.95]
    acc_log['acc_mean'] = acc_mean

    acc_info="<<<<<<<<<Validation>>>>>>>>>>\nACC@iou0.5: {}\nACC@iou0.75: {}\nACC@iou0.95: {}\nACC@iou0.5:0.95: {}\n".format(
               hit_lst[0.5],hit_lst[0.75],hit_lst[0.95], acc_mean)
    #LOGGER.info(acc_info)
    print("acc@iou result>>>>>>>>>>>>>>>>>>>>>>\n", acc_log)
    if len(sample_metrics) == 0:  # No detections over whole validation set.
        print("---- No detections over whole validation set ----")
        return None

    # Concatenate sample statistics
    true_positives, pred_scores, pred_labels = [np.concatenate(x, 0) for x in list(zip(*sample_metrics))]
    metrics_output = ap_per_class(true_positives, pred_scores, pred_labels, labels)
    
    return metrics_output

2. 还是在上一步的代码中新增函数get_batch_statistics_iou，用于计算每一个batch_size的iou值，并返回给全局变量hit_lst中。代码如下：

def get_batch_statistics_iou(outputs, targets, iou_threshold, device):
    """ Compute true positives, predicted scores and predicted labels per sample """
    batch_metrics = []

    global num

    for sample_i in range(len(outputs)):  #对个batch中的每个output进行分析
        num += 1
        if outputs[sample_i] is None:
            continue

        output = outputs[sample_i]
        pred_boxes = output[:, :4]  # N, 4
        pred_scores = output[:, 4]  # 4
        pred_labels = output[:, -1]  # 4

        true_positives = np.zeros(pred_boxes.shape[0])

        annotations = targets[targets[:, 0] == sample_i][:, 1:]
        target_labels = annotations[:, 0] if len(annotations) else []
        
        
        if len(annotations):
            detected_boxes = []
            target_boxes = annotations[:, 1:]
            iou_list = []
            for pred_i, (pred_box, pred_label) in enumerate(zip(pred_boxes, pred_labels)):
                
                pred_box = pred_box.to(device)
                pred_label = pred_label.to(device)

                # If no targets are found break
                if len(detected_boxes) == len(annotations):  #已经对比的预测框达到标记的真实框个数
                    break

                # Ignore if label is not one of the target labels
                if pred_label.to(device) not in target_labels:
                    continue

                iou, box_index = bbox_iou(pred_box.unsqueeze(0), target_boxes).max(0)

                iou_list.append(iou)
                

                if iou >= iou_threshold and box_index not in detected_boxes:
                    true_positives[pred_i] = 1
                    detected_boxes += [box_index]

            max_iou = max(iou_list) if len(iou_list) != 0 else 0
            for thres in thres_lst:
                if max_iou > thres:
                    hit_lst[thres] += 1
        iou_list = []  # 二次置零，避免在没有标记的情况下计算出iou值
        batch_metrics.append([true_positives, pred_scores, pred_labels])
    return batch_metrics

3. 在evaluation.py文件末尾处，将两处utils.utils.evaluation替换为utils.utils.evaluation_iou

    #模型评估
    print("computer mAP...")
    _, _, AP, _ = utils.utils.evaluation_iou(val_dataloader, cfg, model, device)
    print("computer PR...")
    precision, recall, _, f1 = utils.utils.evaluation_iou(val_dataloader, cfg, model, device, 0.3)
    print("Precision:%f Recall:%f AP:%f F1:%f"%(precision, recall, AP, f1))

4. 运行evaluation.py文件，查看输出的acc@iou指标结果。