yolo.py(Yolov1开源项目代码详细理解)

最新推荐文章于 2024-04-17 15:58:08 发布

A940551792

最新推荐文章于 2024-04-17 15:58:08 发布

阅读量427

点赞数

文章标签： python pytorch 目标检测

本文链接：https://blog.csdn.net/A940551792/article/details/130030310

版权

项目代码: 来自github的YOLOv1开源项目
本文是关于yolo.py的详细理解。

该文件写了一个myYOLO类，继承了nn.Module，包含7个函数

init
create_grid
set_grid
decode_boxes
nms
postprocess
forward

init

    def __init__(self, device, input_size=None# 输入图像的尺寸，例如(416, 416)
    , num_classes=20, trainable=False, conf_thresh=0.01, nms_thresh=0.5, hr=False):
        super(myYOLO, self).__init__()
        self.device = device#选择设备，例如torch.device('cuda')
        self.num_classes = num_classes#目标检测任务中要检测的类别数量
        self.trainable = trainable#是否处于训练模式
        self.conf_thresh = conf_thresh#置信度阈值，小于此阈值的预测框将被过滤掉
        self.nms_thresh = nms_thresh#非极大值抑制阈值，用于去除重复的预测框
        self.stride = 32#每个 32x32 像素块在检测模型中会被映射为一个单元
        self.grid_cell = self.create_grid(input_size)#记录划分的网格信息，详细看后文
        self.input_size = input_size#输入图像的尺寸，例如(416, 416)
        self.scale = np.array([[[input_size[1], input_size[0], input_size[1], input_size[0]]]])
        self.scale_torch = torch.tensor(self.scale.copy(), device=device).float()

        # we use resnet18 as backbone
        self.backbone = resnet18(pretrained=True)

        # neck
        self.SPP = nn.Sequential(
            Conv(512, 256, k=1),
            SPP(),
            BottleneckCSP(256*4, 512, n=1, shortcut=False)
        )
        self.SAM = SAM(512)
        self.conv_set = BottleneckCSP(512, 512, n=3, shortcut=False)

        self.pred = nn.Conv2d(512, 1 + self.num_classes + 4, 1)

这段代码是构建模型的主体部分，其中包括三个模块：
backbone：使用的是 ResNet18 预训练模型。
neck：包括 SPP、SAM 和 BottleneckCSP 三个部分，用于进一步提取特征和增强模型的表达能力。
pred：用于预测物体类别和边界框。
backbone 用于提取图像的高级特征。neck是在backbone网络之后添加的一组网络层，这些提取backbone网络的特征之后，对其进行进一步处理和改进。
neck的设计是基于特定的目标检测算法和问题，优化特征提取和物体检测的准确性和效率。
SPP (Spatial Pyramid Pooling)：是一种空间金字塔池化方法，用于捕捉不同尺度的特征。这里使用的是一个 Conv(512, 256, k=1) 层，一个 SPP 层，和一个BottleneckCSP(256*4, 512, n=1, shortcut=False) 层的结构。

create_grid

    def create_grid(self, input_size):
        w, h = input_size[1], input_size[0]
        # generate grid cells
        ws, hs = w // self.stride, h // self.stride
        grid_y, grid_x = torch.meshgrid([torch.arange(hs), torch.arange(ws)])
        grid_xy = torch.stack([grid_x, grid_y], dim=-1).float()
        grid_xy = grid_xy.view(1, hs*ws, 2).to(self.device)
        
        return grid_xy

set_grid

这个函数是用来生成YOLOv1算法中的网格的，input_size是一个元组，保存了图片的宽度和高度信息，ws，hs是该图片横向和纵向的网格数量。torch.arange(hs)会生成一个一维的tensor，值由0到hs-1，torch.meshgrid用于生成网格的坐标点，最后会输出两个大小为hs*ws的tensor，其中grid_y中的数字是网格的列坐标，grid_x中的数字是网格中的行坐标， torch.stack([grid_x, grid_y], dim=-1).float()这个函数是将[grid_x, grid_y]这两个tensor，按照第二维来堆叠，将会生成一个形状为（hs,ws,2）的tensor，grid_xy.view(1, hs*ws, 2).to(self.device)最后通过view函数调整了tensor的维度，并且返回。

    def set_grid(self, input_size):
        self.input_size = input_size
        self.grid_cell = self.create_grid(input_size)
        self.scale = np.array([[[input_size[1], input_size[0], input_size[1], input_size[0]]]])
        self.scale_torch = torch.tensor(self.scale.copy(), device=self.device).float()

这个函数更新了一下几个属性的值。input_size 是输入图片的尺寸，是一个元组，input_size[0]是图片的宽，input_size[1]是图片的高。self.grid_cell里面保存的是上面创建的形状为（hs,ws,2）的tensor，具体的值是每个网格的坐标。

decode_boxes

    def decode_boxes(self, pred):
        """
        input box :  [tx, ty, tw, th]
        output box : [xmin, ymin, xmax, ymax]
        """
        output = torch.zeros_like(pred)
        pred[:, :, :2] = torch.sigmoid(pred[:, :, :2]) + self.grid_cell
        pred[:, :, 2:] = torch.exp(pred[:, :, 2:])

        # [c_x, c_y, w, h] -> [xmin, ymin, xmax, ymax]
        output[:, :, 0] = pred[:, :, 0] * self.stride - pred[:, :, 2] / 2
        output[:, :, 1] = pred[:, :, 1] * self.stride - pred[:, :, 3] / 2
        output[:, :, 2] = pred[:, :, 0] * self.stride + pred[:, :, 2] / 2
        output[:, :, 3] = pred[:, :, 1] * self.stride + pred[:, :, 3] / 2
        
        return output

这是一个用于将模型输出的边界框预测pred，转换为实际边界框坐标output的函数，这里tensor的维度是(batch_size, grid_size, 4)。 pred[:, :, :2] = torch.sigmoid(pred[:, :, :2]) + self.grid_cell这里的 pred[:, :, :2] 表示取 pred 张量中第三维的前两个元素，即网络的 $t_x, t_y$ 预测值。然后通过 torch.sigmoid 函数进行激活，将预测值映射到 [0, 1] 的范围内，再加上网格的坐标，得到中心点坐标在整张图片中的相对位置。pred[:, :, 2:] = torch.exp(pred[:, :, 2:])将预测框的后两个元素取指数，使它们为正数，代表目标框的宽和高（参考Yolo V1论文）。将预测框的前两个元素乘以 self.stride，使它们转化为中心点在整张图片中的绝对坐标，再加上或者减去宽高除以2，得到边界框点的具体位置。

nms

    def nms(self, dets, scores):
        """"Pure Python NMS baseline."""
        x1 = dets[:, 0]  #xmin
        y1 = dets[:, 1]  #ymin
        x2 = dets[:, 2]  #xmax
        y2 = dets[:, 3]  #ymax

        areas = (x2 - x1) * (y2 - y1)                 # the size of bbox
        order = scores.argsort()[::-1]                        # sort bounding boxes by decreasing order

        keep = []                                             # store the final bounding boxes
        while order.size > 0:
            i = order[0]                                      #the index of the bbox with highest confidence
            keep.append(i)                                    #save it to keep
            xx1 = np.maximum(x1[i], x1[order[1:]])
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])

            w = np.maximum(1e-28, xx2 - xx1)#第一个参数 1e-28 是一个非常小的数，目的是避免除数为0或接近0时的错误
            h = np.maximum(1e-28, yy2 - yy1)
            inter = w * h

            # Cross Area / (bbox + particular area - Cross Area)
            ovr = inter / (areas[i] + areas[order[1:]] - inter)
            #reserve all the boundingbox whose ovr less than thresh
            inds = np.where(ovr <= self.nms_thresh)[0]
            order = order[inds + 1]

        return keep

NMS（非极大值抑制）算法，用于在检测框中去除冗余的框，保留置信度最高的框。
dets 是一个二维数组，形状为 (num_detections, 4)，其中 num_detections 是检测到的目标数。每一行代表一个检测到的目标，包含四个值 [xmin, ymin, xmax, ymax] 分别表示目标框的左上角和右下角坐标。
scores是一个一维数组，和dets是一一对应的，保存的是该框置信度最高类别的得分。
order = scores.argsort()[::-1]是对scores按照分数大小排序，并且返回索引的排序，[::-1]的意思的进行降序排列。keep.append(i)将目前置信度最高的框的索引（整型）添加到keep这个list里。
x1[i]是当前置信度最大的框的左上角坐标，x1[order[1:]]是剩下的框的左上角坐标，这里依次对比，得到的xx1是一个一维数组。xx1和xx2分别是计算当前框与其他框的交集时，交集框的左上角和右下角的x坐标，所以xx2表示右边界更小的那个框的右边界，xx1表示左边界更大的那个框的左边界，需要取二者的较大值作为交集框的左边界。
值得注意的是这里的xx1等是一维数组。
ovr = inter / (areas[i] + areas[order[1:]] - inter)，这个代码的意思是计算交并比，假设橙色框是最大置信度框i，那么areas[i]就是橙色框的面积， areas[order[1:]]是黄色框的面积。
np.where(ovr <= self.nms_thresh)这里的where函数会输出满足条件的ovr中元素的索引，如果ovr是二维数组就会输出一个tuple，因此这里作者写上了[0]来避免出错，但是实际上这里ovr是一维数组，可以不写[0]。+ 1是因为在inds中存储的是除去第一个元素的剩余元素的下标，所以在从原来的order数组中取出这些下标的元素时，需要将这些下标都加1，才能保证取到正确的元素。

postprocess

    def postprocess(self, all_local, all_conf, exchange=True, im_shape=None):
        """
        bbox_pred: (HxW, 4), bsize = 1
        prob_pred: (HxW, num_classes), bsize = 1
        """
        bbox_pred = all_local
        prob_pred = all_conf

        cls_inds = np.argmax(prob_pred, axis=1)
        prob_pred = prob_pred[(np.arange(prob_pred.shape[0]), cls_inds)]
        scores = prob_pred.copy()

这个函数是一个完整的目标检测模型的后处理函数，是针对一张图片的预测结果进行处理的。all_local 是所有框的位置信息，可以看做是预测出的边界框的坐标和宽高，all_conf 是所有框的每个类别的预测得分。np.argmax(prob_pred, axis=1)，prob_pred的维度是(HxW, num_classes)，每一行都是框的每个类别的预测得分。这个函数中 axis=1表示在每一行的方向上进行操作，也就是计算每一个框最可能是什么类别的物体，并且将物体类别索引保存在cls_inds中。
prob_pred = prob_pred[(np.arange(prob_pred.shape[0]), cls_inds)]使用了NumPy的高级索引（fancy indexing）方式，具体体来说就是传入坐标来对prob_pred进行切片。prob_pred.shape[0]生成了一个一维的数组，和prob_pred具有相同的行数，也就是行数保持不变，cls_inds作为列索引，得到了一个元组，将该元组作为坐标值，从原先的prob_pred取出数值，最后得到了一个一维的prob_pred。
scores = prob_pred.copy()最后得到了置信度得分scores。

# threshold
        keep = np.where(scores >= self.conf_thresh)
        bbox_pred = bbox_pred[keep]
        scores = scores[keep]
        cls_inds = cls_inds[keep]

keep 中保存的超过分数阈值的检测框的索引号，然后从bbox_pred scores cls_inds中取出对应的元素。

        # NMS
        keep = np.zeros(len(bbox_pred), dtype=np.int)#记录了哪些框被保留下来的flag
        for i in range(self.num_classes):#对每一个i类都进行nms
            inds = np.where(cls_inds == i)[0]#找到i类别的检测框索引
            if len(inds) == 0:#如果没有这一类就进行下一个类的处理
                continue
            c_bboxes = bbox_pred[inds]
            c_scores = scores[inds]#取出对应元素
            c_keep = self.nms(c_bboxes, c_scores)#nms
            keep[inds[c_keep]] = 1#修改flag

这段代码是对同一个类别的目标检测框进行非极大值抑制。如果对整个图像进行非极大值抑制，会导致有些检测框被错误的抑制掉，比如不同类别的目标重合度比较高，其检测框IOU也就很高，这时可能会抑制掉其中一个。这里的keep记录了哪些框被保留下来，对每个目标框都先记0，然后经过nms以后将保留的目标框记录为1。

        keep = np.where(keep > 0)#挑选出flag=1的框，他们都是被nms保留下来的框
        bbox_pred = bbox_pred[keep]
        scores = scores[keep]
        cls_inds = cls_inds[keep]#选出对应的元素

        if im_shape != None:
            # clip
            bbox_pred = self.clip_boxes(bbox_pred, im_shape)#方法可能是用来剪裁（clip）或调整bounding boxes的大小和位置，以使它们适合于图像的尺寸和范围，代码里没写这个函数。

        return bbox_pred, scores, cls_inds

目标检测模型的后处理函数的输出是三个数组，分别为最终的预测框bbox_pred、得分scores和类别cls_inds，通过list下标一一对应。

forward

 def forward(self, x, target=None):
        # backbone
        _, _, C_5 = self.backbone(x)#C_5是backbone处理后的结果中的一个特定的特征图

        # head
        C_5 = self.SPP(C_5)
        C_5 = self.SAM(C_5)
        C_5 = self.conv_set(C_5)

首先，输入数据通过 self.backbone（Resnet18）进行 backbone 处理，得到 C_5，这是 backbone 处理后的结果中的一个特定的特征图。然后，C_5 这个特征图通过一系列操作在 head 中被进一步处理，包括 Spatial Pyramid Pooling (self.SPP)、Spatial Attention Mechanism (self.SAM)、卷积操作 self.conv_set 等。最终，得到的输出结果即为预测的目标框的位置和类别信息。

        # pred
        prediction = self.pred(C_5)
        prediction = prediction.view(C_5.size(0), 1 + self.num_classes + 4, -1).permute(0, 2, 1)
        B, HW, C = prediction.size()#B 表示 batch size，HW 表示特征图的空间大小，C 表示每个预测框的维度数。

prediction是经过全连接层处理后的特征图，包含了每个空间位置的类别概率、边界框位置和置信度预测。这里prediction的维度为(batch_size, 1 + num_classes + 4, H, W)，为了便于后续处理，将prediction的维度调整为(batch_size,1 + num_classes + 4，H*W )，PyTorch 中，如果某个维度上的大小不确定，可以使用 -1 来表示该维度的大小应该根据其他维度的大小和 tensor 的总大小自动计算。然后通过permute交换第二列和第三列，得到维度为(batch_size, HW, 1 + num_classes + 4)的tensor。

        # Divide prediction to obj_pred, txtytwth_pred and cls_pred   
        # [B, H*W, 1]
        conf_pred = prediction[:, :, :1]
        # [B, H*W, num_cls]
        cls_pred = prediction[:, :, 1 : 1 + self.num_classes]
        # [B, H*W, 4]
        txtytwth_pred = prediction[:, :, 1 + self.num_classes:]

conf_pred的维度为(batch_size, HxW, 1)，表示每个bbox是否为目标物体的置信度。
cls_pred的维度为(batch_size, HxW, num_cls)，表示每个bbox属于不同类别的概率值。
txtytwth_pred的维度为(batch_size, HxW, 4)，表示每个bbox的位置偏移量，其中tx和ty表示相对于bbox中心点的偏移量，tw和th表示相对于bbox的宽度和高度的缩放系数。

 # test
        if not self.trainable:#测试模式的情况下
            with torch.no_grad():
                # batch size = 1
                all_conf = torch.sigmoid(conf_pred)[0]           # 0 is because that these is only 1 batch.
                                     label=target)

with torch.no_grad(): 是一个上下文管理器（Context Manager），用于执行一些不需要计算梯度的代码块。在这个上下文中，PyTorch将不会为任何操作构建计算图，也不会计算任何梯度，这可以提高代码的运行效率和降低内存消耗。
conf_pred 的维度为 [B, H*W, 1]，torch.sigmoid(conf_pred) 将对每个元素进行 Sigmoid 操作，输出维度为 [B, H*W, 1]，使用 [0] 取出第一个样本的置信度。最终，all_conf 的维度为 [H*W, 1]，表示特征图上所有先验框的置信度预测值。

                all_bbox = torch.clamp((self.decode_boxes(txtytwth_pred) / self.scale_torch)[0], 0., 1.)
                all_class = (torch.softmax(cls_pred[0, :, :], 1) * all_conf)

首先对txtytwth_pred进行decode转换为实际坐标值，然后解码后的坐标值除以 self.scale_torch 进行归一化，而这里的clamp函数将所有超出这个范围的坐标值限制在[0, 1]内，确保了所有的边界框坐标都在图像内部。最终，all_bbox 的维度为 [H*W, 4]，表示特征图上所有先验框的预测坐标值。这行代码是将预测出的类别分数cls_pred经过softmax归一化处理后乘以置信度all_conf，得到每个类别在所有检测框中的得分，其中all_conf是在之前的后处理中保留下来的置信度信息，表示每个检测框被认为是目标的置信度大小。

                all_conf = all_conf.to('cpu').numpy()
                all_class = all_class.to('cpu').numpy()
                all_bbox = all_bbox.to('cpu').numpy()
                
                bboxes, scores, cls_inds = self.postprocess(all_bbox, all_class)

最后调用postprocess()函数得到检测到的目标框、置信度和类别预测结果。

        else:#在训练模式下
            conf_loss, cls_loss, txtytwth_loss, total_loss = tools.loss(pred_conf=conf_pred, pred_cls=cls_pred,
                                                                        pred_txtytwth=txtytwth_pred,
                                                                        label=target)

            return conf_loss, cls_loss, txtytwth_loss, total_loss

如果trainable是True，表示当前模式为training（训练），代码会执行目标检测的前向计算，计算置信度预测损失、类别预测损失和目标框预测损失，最后返回。