代码收藏家技术教程 2023-06-05

YOLOv5 v5.0：解读边界框预测中的对象概率和分类置信度

首先使用torch.hub.load加载已经训练好的模型

模型的训练过程可以参考此链接

	import torch
	
	# 从本地加载自定义的YOLOv5模型
	model_path = ".."  # yolov5根目录
	weight_path = model_path+"/weights/best.pt"
	model = torch.hub.load(model_path, "custom", weight_path, source="local")

给模型一张图片进行预测

	imgs = [model_path+'/dataset/images/test/29998.jpg']
	results = model(imgs)
	
	from PIL import Image
	import matplotlib.pyplot as plt
	
	results.print()
	results.show()  # 这两句用于看一下模型检测结果

检测结果
查看results中所包含的信息

	results.pandas().xyxy  # Pandas DataFrame

输出：

输出的内容中，每一行是一个边界框，前4列是边界框4条边的位置，confidence是总体置信度（此处的confidence是标题中所说的probability和confidence的乘积），class是分类类别的index（其在yolov5根目录/data/你的配置文件.yaml，或者任何自定义位置中配置），name是分类类别的名字

注意，此处模型所输出的results的类型是yolov5根目录/models/common.py中定义的Detections类，其中有pandas()方法和xyxy等成员的定义，其值来自Detections实例化时的第2个参数pred。

这个类的唯一调用是在这个文件中，上面一点的autoShape类，其传入的对应参数是y，y来自y = non_max_suppression(y, conf_thres=self.conf, iou_thres=self.iou, classes=self.classes) # NMS，这里面的第一个参数y来自上面的y = self.model(x, augment, profile)[0]，即forward得到的网络输出。

注意在nom_max_suppression方法中，obj_conf和cls_conf被相乘成为conf，并且y在从nom_max_suppression中出来以后，其结构没有动过，因此应该修改nom_max_suppression方法。

对yolov5源码的修改

首先看到yolov5根目录/utils/generals.py中的nom_max_suppression方法，其返回值output即为上面results.pandas().xyxy中得到的返回结果，因此此处修改这个返回值，以便使用上面的方法将其获取到。

找到这个方法中的下面两行：

	# Compute conf
    x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf

此处将对象概率和分类置信度相乘了。注意，x[:, 5:]是所有类别的分类置信度，比如我这里一共42类，那他就有42列，而x[:, 4:5]是边界框的对象概率。

为了能够将其输出，在上面那行将两者相乘的代码之前将其保存在新变量中：

    # 在计算conf之前将两个值拿出来
    cls_conf_all = x[:, 5:].clone() # 这是所有类别的置信度
    # 注意由于有42个类别，取出来的obj_conf有42列。
    
    obj_conf = x[:, 4:5].clone()  # 对象概率

    # Compute conf
    x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf

注意此处的 .clone()，使用clone方法是为了保证在代码对变量 x 进行修改时不会改动到这里取出来的cls_conf_all 和obj_conf 。

然后找到以下代码块进行修改：

    # Detections matrix nx6 (xyxy, conf, cls)
    if multi_label:  # 如果允许一个盒子（？）多label
        i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
        x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
    else:  # best class only 每个盒子仅取置信度最高的label
        conf, j = x[:, 5:].max(1, keepdim=True)  # 注意conf是最大值，j是其对应的index
        x = torch.cat((box, conf, j.float(), obj_conf, cls_conf_all), 1)[conf.view(-1) > conf_thres]
        # 注意此处的修改版 x ，其0-5列保持原样，6列是对象概率obj_conf，7-48列是这42个类别分别的置信度，一共49列

        # x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]  # 这里之后，x的0-3列是xyxy，然后是conf（已经乘以对象概率）和j

这样就得到了一个n行，7+nc列的output，其中n是输出边界框的数量，nc是类别总数。
之所以要在这里修改，而不是将对象概率和分类置信度单独拿出来最后拼接，是因为 x 在这一块代码之后被根据置信度排序了，这导致提前拿出来的cls_conf_all 和obj_conf 与 x 的行对不上。

然后再回到刚才common.py中的Detections类，此处由于yolov5原码中很多地方都写死了，直接将这个shape被改变过的output作为pred参数传给Detections类会出一些问题，因此我们可以在这个pred被赋值给成员之前将其存储到一个新的成员中，如下：

class Detections:
    # detections class for YOLOv5 inference results
    def __init__(self, imgs, pred, files, times=None, names=None, shape=None):
        super(Detections, self).__init__()
        self.newnew = pred  # 新成员
        if(pred[0].shape[1] == 49):
            pred = pred[0][:, :-43]  # 这里将最后那些列给删掉了，同时也去掉了第1维
            pred = pred[None, ...]  # 用这一句来将pred的第1维给补上
        d = pred[0].device  # device
        # ...

注意此处的49和-43，49是因为我这里分类类别一共42列，拼接之后是49列。