目标检测-训练与测试

目标检测网络的训练大致是如下的流程:

  • 设置各种超参数
  • 定义数据加载模块 dataloader
  • 定义网络 model
  • 定义损失函数 loss
  • 定义优化器 optimizer
  • 遍历训练数据,预测-计算loss-反向传播

首先,我们导入必要的库,然后设定各种超参数

后处理

目标框信息解码

之前我们的提到过,模型不是直接预测的目标框信息,而是预测的基于anchor的偏移,且经过了编码。因此后处理的第一步,就是对模型的回归头的输出进行解码,拿到真正意义上的目标框的预测结果。

后处理还需要做什么呢?由于我们预设了大量的先验框,因此预测时在目标周围会形成大量高度重合的检测框,而我们目标检测的结果只希望保留一个足够准确的预测框,所以就需要使用某些算法对检测框去重。这个去重算法叫做NMS

NMS的大致算法步骤如下:

  1. 按照类别分组,依次遍历每个类别。

  2. 当前类别按分类置信度排序,并且设置一个最低置信度阈值如0.05,低于这个阈值的目标框直接舍弃。

  3. 当前概率最高的框作为候选框,其它所有与候选框的IOU高于一个阈值(自己设定,如0.5)的框认为需要被抑制,从剩余框数组中删除。

  4. 然后在剩余的框里寻找概率第二大的框,其它所有与第二大的框的IOU高于设定阈值的框被抑制。

  5. 依次类推重复这个过程,直至遍历完所有剩余框,所有没被抑制的框即为最终检测框。

    目标检测-训练与测试_第1张图片

    单图预测推理

当模型已经训练完成后,下面我们来看下如何对单张图片进行推理,得到目标检测结果。

首先我们需要导入必要的python包,然后加载训练好的模型权重。

随后我们需要定义预处理函数。为了达到最好的预测效果,测试环节的预处理方案需要和训练时保持一致,仅去除掉数据增强相关的变换即可。

因此,这里我们需要进行的预处理为:

  • 将图片缩放为 224 * 224 的大小
  • 转换为 Tensor 并除 255
  • 进行减均值除方差的归一化
  • 接着我们就来进行推理,过程很简单,核心流程可以概括为:

  • 读取一张图片
  • 预处理
  • 模型预测
  • 对模型预测进行后处理

最后,我们将最终得到的检测框结果进行绘制,得到类似如下图的检测结果:

目标检测-训练与测试_第2张图片

你可能感兴趣的:(目标检测)