yolov1搭建介绍及代码撰写详解（总结6）
可以从本人以前的文章中可以看出作者以前从事的是嵌入式控制方面相关的工作，是一个机器视觉小白，之所以开始入门机器视觉的学习主要是一个idea，想把机器视觉与控制相融合未来做一点小东西。废话不多说开始正题。（如有侵权立即删稿）
摘要
本文是介绍yolov1目标检测网络搭建,个人对其的知识总结，以及结合论文进行讲解，网络设计的知识点，代码撰写部分本人是借鉴的大佬的源码（下文会给出网址），基于pytorch编写代码。作为一个刚入门的小白怎么去学习别人的代码，一步一步的去理解每一行代码，怎么将网络设计变成代码，模仿大佬的代码去撰写。作为小白如有不足之处请批评指正哈。
YOLOV1
在网络设计之前需要明白什么是YOLOV1。
以下是我借鉴的文章的参考链接：
1.yolov1论文You Only Look Once
2.YOLOv1-v10各版本的作者（附论文及项目地址）
3.YOLOV1（pytorch） 代码讲解
4.YOLO系列算法精讲：从yolov1至yolov8的进阶之路（2万字超全整理）
5.pytorch实战13：基于pytorch实现YOLOv1（长长文）
接下来，我将结合论文进行讲解，明白论文的构思以及步骤，本人将论文的关键部分摘要出来，并且全部用中文翻译表示。
原文摘要：我们提出了YOLO，一种新的目标检测方法。先前关于对象检测的工作重新使用分类器来执行检测。相反，我们将目标检测框定为空间上分离的边界框和相关类概率的回归问题。单个神经网络在一次评估中直接从全图像预测边界框和类概率。由于整个检测管道是一个单一的网络，它可以直接在检测性能上进行端到端的优化。我们的统一架构速度极快。我们的基础YOLO模型以每秒45帧的速度实时处理图像。该网络的一个较小版本Fast YOLO每秒处理155帧的速度令人震惊，同时还实现了其他实时探测器的两倍的mAP。与最先进的检测系统相比，YOLO会产生更多的定位错误，但在背景中预测假阳性的可能性较小。最后，YOLO学会了非常一般的物体表示法。当从自然图像推广到艺术品等其他领域时，它的性能优于其他检测方法，包括DPM和R-CNN。
小结：在摘要中概述了yolo，介绍了其每秒45帧运行速度极快，与当时哪个年代的目标检测算法进行了对比。
原文内容1：使用YOLO处理图像简单明了。我们的系统
（1）将输入图像调整为448 × 448，
（2）在图像上运行单个卷积网络，
（3）通过模型的置信度对结果检测进行阈值处理
小结1：网络输入图像大小有所限制，置信度的结果输出有阈值限制，肯定是大于阈值才在图像中框选物体输出目标识别结果。
原文内容2：我们将目标检测重构为一个单一的回归问题，直接从图像像素到边界框坐标和类概率。使用我们的系统，您只需看一次（YOLO）YOU ONLY LOOK ONCE图像，就可以预测出现了哪些对象以及它们的位置。第一，YOLO的速度非常快。我们的基础网络以每秒45帧的速度运行。第二，YOLO在全图范围内对图片进行推理。YOLO在训练和测试期间看到整个图像，因此它隐式地编码了有关类及其外观的上下文信息。快速R-CNN是一种顶部检测方法[14]，它会将图像中的背景块误认为是对象，因为它无法看到更大的上下文。缺点：但它很难精确定位一些物体，尤其是小物体。
小结2：工作原理：YOLO 将整个图像划分为网格，并在每个网格中同时预测多个边界框和类标签。由于它一次性处理整个图像，因此可以获取更全面的上下文信息。
上下文信息的利用：例如，当 YOLO 看到一张包含多个对象（如一只狗在公园里玩耍，周围有许多树和其他动物）的图像时，它不仅能识别出“狗”，还可以理解“公园”的背景，意识到狗和树的关系，从而更准确地检测和分类。
原文内容3：统一检测，我们的系统将输入图像划分为S × S网格。如果对象的中心落入网格单元，则该网格单元负责检测该对象。
小结3：也就是将图像变成了S × S网格，可以从这张图中看出，S × S = 7 × 7。可以看到狗占了好几个格子，所以这些单元格都会预测狗，就会有许多的框重叠如下。
这张图中狗的大黑框的中心点如图所示，哪个交叉红点落在7 × 7的哪个网格中（感觉是第5排第2个），哪个网格就负责猜测他的种类是猫是狗还是人呢。很多框都框出来了狗，但是中心的单元格只有一个，所以框看起来重合但是每个画狗的框肯定有不重叠的部分，不然他们的中心单元格不就一样了吗？后续会讲到怎么消除这些重叠框。
原文内容4：每个网格单元预测B个边界框和这些框的置信度得分。这些置信度分数反映了模型对该框包含对象的置信度，以及模型认为该框的预测准确度。形式上，我们将置信定义为Pr（Object）* IOU真实预测。如果该单元格中不存在对象，则置信度分数应为零。否则，我们希望置信度得分等于预测框和地面真实值之间的并集交集（IOU）。
小结4：一个单元格内可能画B（yolo是用的2个）个边界框，这两个边界框的中心一定是落在这个单元格中的。每个框都会有每个物种的置信度得分，20个物种那就会有20个得分值。这里引入了IOU的概念叫做intersection over union（交并比）。
原文内容5：边界框概述：每个边界框由5个预测组成：x、y、w、h和置信度。（x，y）坐标表示相对于网格单元格边界的框的中心。宽度和高度是相对于整个图像预测的。最后，置信度预测表示预测框和任何地面实况框之间的IOU。
我们的系统将检测建模为一个回归问题。它将图像划分为S×S网格，并为每个网格单元预测B边界框，这些框的置信度和C类概率。这些预测被编码为S × S ×（B×5 +C）tensor。为了在PASCAL VOC上评估YOLO，我们使用S = 7，B = 2。PASCAL VOC有20个标记类别，因此C = 20。我们的最终预测是一个7 × 7 × 30的张量。
小结5：tensor张量的构成在此处，30 =【2*（x、y、w、h和置信度）+20】，这也是真正要训练的东西，要训练的和GoogleNet、VGG16等不同，不再是根据标签物种训练，要训练的目标是tensor【2*（x、y、w、h和置信度）+20】，所以为此代码撰写部分得要增加文件处理部分，txt部分内容为【2*（x、y、w、h和置信度）+20】。
原文内容6：2.1.网络设计，网络的初始卷积层从图像中提取特征，而全连通层预测输出概率和坐标。我们的网络架构受到了GoogleNet图像分类模型的启发[34]。我们的网络有24个卷积层，后面是2个完全连接的层。我们不使用GoogleNet使用的初始模块，而是简单地使用1×1约简层，然后是3×3卷积层，类似于Lin等人[22]的方法。完整的网络如图3所示。我们还训练了YOLO的快速版本，旨在推动快速目标检测的边界。Fast YOLO使用的神经网络具有较少的卷积层（9层而不是24层）和较少的滤波器。除了网络规模之外，YOLO和Fast YOLO的所有培训和测试参数都相同。
小结6：网络的设计是参考了GoogleNet，可以见我的前文对GoogleNet的讲解。
原文内容7：我们的检测网络有24个卷积层，后面是2个完全连接层。交替的1 × 1卷积层减少了前一层的特征空间。我们在ImageNet分类任务中以一半的分辨率（224 × 224输入图像）预训练卷积层，然后将分辨率加倍以进行检测。然后，我们转换模型以执行检测。Ren等人指出，在预训练网络中添加卷积层和连通层可以提高性能[29]。按照他们的例子，我们添加了四个卷积层和两个完全连接层与随机初始化的权重。检测通常需要细粒度的视觉信息，因此我们将网络的输入分辨率从224 × 224提高到448 × 448。
小结7:上文图片很清晰的描述了网络设计，正如原文所述。输入图片为448 × 448的图片。
原文内容8：我们的最后一层预测了类概率和边界框坐标。我们通过图像的宽度和高度来规范化边界框的宽度和高度，使它们落在0和1之间。我们将边界框的x和y坐标参数化为特定网格单元位置的偏移，因此它们也被限制在0和1之间。我们对最后一层使用线性激活函数。
小结8:边界框在最后一层里面。
原文内容9：YOLO预测每个网格单元有多个边界框。在训练时，我们只需要一个边界框预测器来负责每个对象。我们基于哪个预测具有最高的当前IOU以及地面真实值来分配一个预测器来“负责”预测对象。这导致了边界框预测器之间的专门化。每个预测器都能更好地预测特定的大小、长宽比或对象类别，从而提高整体回忆率。
小结9：每个单元格会有多个边界框输出，但是只取IOU最高的概率计算值最高的来负责预测。
原文内容10： 训练部分损失函数。请注意，损失函数仅在对象存在于该网格单元中时才惩罚分类错误（因此前面讨论了条件类概率）。它也仅在预测器“负责”地面实况框（即，在该网格单元中具有任何预测器中的最高IOU）的情况下惩罚边界框坐标误差。
- 边界框位置损失：
• 这个损失项计算了预测边界框的中心点坐标（x, y）以及宽度和高度（w, h）与真实边界框之间的差异。
• 通常使用均方误差（MSE）来计算这个损失，目的是让模型学会准确预测目标的边界框位置和大小。 - 分类损失：
• 这是计算目标类别的预测概率与真实类别标签之间的差异。通常使用交叉熵损失。
• 该部分损失帮助模型学习如何区分不同的类别。
3.置信度损失（Confidence Loss）：
• 置信度损失计算的是模型对于边界框内存在目标的置信度与真实值之间的差异。
• 当边界框确实包含目标时，confidence 的真实值为1；当边界框不包含目标时，confidence 的真实值为0。
• 这个损失项的作用是鼓励模型在目标框内尽可能高的置信度，同时在没有目标的框内保持低的置信度。
小结:可以看出损失函数也是围绕x、y、w、h和置信度+20目标类别构造的。根据论文已知，λcoord = 5 and λnoobj =0.5 ，xi 和 yi真实边界框的中心点坐标 (x_i ) ̂,(y_i ) ̂预测边界框的中心点坐标。
首先红框的第一行真实坐标中心点减去预测坐标中心点计算中心点偏差。第二行中是框大小误差的计算，越大的项开平方越小。
如图所示的红线，随着平方根自变量X的增加，斜率越来越小。从而实现对小数值更加敏感。
蓝色框中是预测BBOX的误差，第一行是预测概率高的BBOX，而小概率的BBOX则乘以了0.5。绿色部分也就是最常见的标签损失预测。
原文内容11：就像在训练中一样，预测测试图像的检测只需要一次网络评估。在PASCAL VOC上，该网络预测每个图像98个边界框和每个框的类概率。YOLO在测试时非常快，因为它只需要一个网络评估，不像基于分类器的方法。网格设计在边界框预测中强制执行空间多样性。通常很清楚一个物体福尔斯落在哪个网格单元中，网络只为每个物体预测一个盒子。
小结10：49个单元格，每个格子画两个边界框，不会随着物体的增加速度大幅度减缓，物体增加，原本不需要大量计算的方格出现了物体，一定程度上速度还是会减缓的。
原文内容12:缺点:YOLO对边界框预测施加了很强的空间约束，因为每个网格单元只能预测两个框，并且只能有一个类。这种空间约束限制了我们的模型可以预测的附近对象的数量。我们的模型很难处理成组出现的小对象，如鸟群。由于我们的模型学习从数据中预测边界框，它很难推广到新的或不寻常的纵横比或配置的对象。我们的模型还使用相对粗糙的特征来预测边界框，因为我们的架构具有来自输入图像的多个下采样层。
小结11：作者概述了yolo的不足。
文件处理部分
本次采用的数据集是VOC数据集，VOC数据集主要使用的是图像和xml文件。图像不用多说，xml文件格式如下。
从这个xml文件中可以看出，文件的图片是2007_000061.jpg,图片大小为500，333
物体1：x1位置274，y1位置11，x2位置437，y2位置279，物种类别是船
物体2：x1位置184，y1位置214，x2位置281，y2位置252，物种类别是船
所以txt文件中应该书写为：2007_000061.jpg 274 11 437 279 3 184 214 281 252 3
软件代码构思
有了以上理论基础后，开始构建代码思路，整体构建思路如下图所示，写代码之前一定要构思好大致思路，代码永远是为你思路框架服务的。整体代码构思，代码为5个py文件撰写。训练部分中1.voc数据集文件处理为txt，这一块主要是将voc数据集的图片对应好每一个txt文件，上文介绍了如何构造txt文件。2.网络采用resnet，网络设计方面并不采用原文中的网络而是采用resnet50。在YOLOv1论文中，作者自己构建了一个架构，但是该架构还需要在ImageNet分类数据集上预训练，比较麻烦。因此大部分人都选用的ResNet作为自己的backbone，因为可以方便调用官方的预训练权重，这里我也是采用的ResNet。3.损失函数构建根据原文搭建。4.训练部分，后续详细概述。5.dataset构建主要是将提取出来的txt文件进行处理，将每个图片对应其文本内容，如2007_000027.jpg，对应标签tensor30，内容153 15 500 369 15 。
1.voc数据集文件处理为txt
这部分代码主要分为3个部分，首先定义基本参数，识别物体的类别名等等。其次是解析xml文件的内容，读取出需要的内容。最后将在xml中读取的内容排序好后放入voctrain.txt、voctest.txt两个文件中分别用作训练和测试。
from xml.etree import ElementTree as ET
import os
import random
# -------------------1. 定义一些基本的参数------------------
# 定义所有的类名
VOC_CLASSES = (
'aeroplane', 'bicycle', 'bird', 'boat',
'bottle', 'bus', 'car', 'cat', 'chair',
'cow', 'diningtable', 'dog', 'horse',
'motorbike', 'person', 'pottedplant',
'sheep', 'sofa', 'train', 'tvmonitor')
# 训练集和测试集文件名字
train_set = open('voctrain.txt', 'w')
test_set = open('voctest.txt', 'w')
# 要读取的xml文件路径，记得自己修改路径
Annotations = 'D:/Pycharm/YOLOv1_pytorch/dataset/VOCtrainval_11-May-2012/VOCdevkit/VOC2012/Annotations/'
# 列出所有的xml文件
xml_files = os.listdir(Annotations)
# 打乱数据集
random.shuffle(xml_files)
# 训练集数量
train_num = int(len(xml_files) * 0.7)
# 训练列表
train_lists = xml_files[:train_num]
# 测测试列表
test_lists = xml_files[train_num:]
# -------------------2. 定义解析xml文件的函数-------------------------
def parse_rec(filename):
# 参数：输入xml文件名
# 创建xml对象
tree = ET.parse(filename)
objects = []
# 迭代读取xml文件中的object节点，即物体信息
for obj in tree.findall('object'):
obj_struct = {}
# difficult属性，即这里不需要那些难判断的对象
difficult = int(obj.find('difficult').text)
if difficult == 1: # 若为1则跳过本次循环
continue
# 开始收集信息
obj_struct['name'] = obj.find('name').text
bbox = obj.find('bndbox')
obj_struct['bbox'] = [int(float(bbox.find('xmin').text)),
int(float(bbox.find('ymin').text)),
int(float(bbox.find('xmax').text)),
int(float(bbox.find('ymax').text))]
objects.append(obj_struct)
return objects
# -------------------------3. 把信息保存入文件中-------------------------
def write_txt():
# # 生成训练集txt
count = 0
for train_list in train_lists:
count += 1
# 获取图片名字
image_name = train_list.split('.')[0] + '.jpg' # 图片文件名
# 对他进行解析
results = parse_rec(Annotations + train_list)
# 如果返回的对象为空，表示张图片难以检测，因此直接跳过
if len(results) == 0:
print(train_list)
continue
# 否则，则写入文件中
# 先写入文件名字
train_set.write(image_name)
# 接着指定下面写入的格式
for result in results:
class_name = result['name']
bbox = result['bbox']
class_name = VOC_CLASSES.index(class_name)
train_set.write(' ' + str(bbox[0]) +
' ' + str(bbox[1]) +
' ' + str(bbox[2]) +
' ' + str(bbox[3]) +
' ' + str(class_name))
train_set.write('\n')
train_set.close()
# 生成测试集txt
# 原理同上面
for test_list in test_lists:
count += 1
image_name = test_list.split('.')[0] + '.jpg' # 图片文件名
results = parse_rec(Annotations + test_list)
if len(results) == 0:
print(test_list)
continue
test_set.write(image_name)
for result in results:
class_name = result['name']
bbox = result['bbox']
class_name = VOC_CLASSES.index(class_name)
test_set.write(' ' + str(bbox[0]) +
' ' + str(bbox[1]) +
' ' + str(bbox[2]) +
' ' + str(bbox[3]) +
' ' + str(class_name))
test_set.write('\n')
test_set.close()
# 5. 运行
if __name__ == '__main__':
write_txt()
运行完成之后分别生成
voctrain.txt、voctest.txt这两个文件，文件内容如图所示。
2.网络采用resnet
import numpy as np
# 1. 导入所需要的包
import torch
import math
from torch import nn
import torch.utils.model_zoo as model_zoo
import torch.nn.functional as F
# 2. 构建block: 不含有1*1
class Base_Block(nn.Module):
# 用于扩充的变量，表示扩大几倍
expansion = 1
def __init__(self,in_planes,out_planes,stride=1,downsample=None):
'''
:param in_planes: 输入的通道数
:param planes: 输出的通道数
:param stride: 默认步长
:param downsample: 是否进行下采样
'''
super(Base_Block, self).__init__()
# 定义网络结构 + 初始化参数
self.conv1 = nn.Conv2d(in_planes,out_planes,kernel_size=3,stride=stride,padding=1,bias=False)
self.bn1 = nn.BatchNorm2d(out_planes)
self.relu = nn.ReLU(inplace=True) # inplace为True表示直接改变原始参数值
self.conv2 = nn.Conv2d(out_planes,out_planes,kernel_size=3,stride=1,padding=1,bias=False)
self.bn2 = nn.BatchNorm2d(out_planes)
self.downsample = downsample
self.stride = stride
def forward(self,x):
# 前向传播
res = x # 残差
# 正常传播
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
# 判断是否下采样
if self.downsample is not None:
res = self.downsample(res)
# 相加
out += res
# 返回结果
out = self.relu(out)
return out
# 3. 构建Block：含有1*1
class Senior_Block(nn.Module):
expansion = 4
def __init__(self,in_planes,planes,stride=1,downsample=None):
'''
:param in_planes: 输入通道数
:param planes: 中间通道数，最终的输出通道数还需要乘以扩大系数，即expansion
:param stride: 步长
:param downsample: 下采样方法
'''
super(Senior_Block, self).__init__()
self.conv1 = nn.Conv2d(in_planes,planes,kernel_size=1,bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(planes,planes,kernel_size=3,stride=stride,padding=1,bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
self.bn3 = nn.BatchNorm2d(planes*4)
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
def forward(self,x):
# 残差
res = x
# 前向传播
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.relu(out)
out = self.conv3(out)
out = self.bn3(out)
# 是否下采样
if self.downsample is not None:
res = self.downsample(x)
# 相加
out += res
out = self.relu(out)
return out
# 4. 构建输出层
class Output_Block(nn.Module):
expansion = 1
def __init__(self, in_planes, planes, stride=1, block_type='A'):
'''
:param in_planes: 输入通道数
:param planes: 中间通道数
:param stride: 步长
:param block_type: block类型，为A表示不需要下采样，为B则需要
'''
super(Output_Block, self).__init__()
# 定义卷积
self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=2, bias=False,dilation=2)
self.bn2 = nn.BatchNorm2d(planes)
self.conv3 = nn.Conv2d(planes, self.expansion*planes, kernel_size=1, bias=False)
self.bn3 = nn.BatchNorm2d(self.expansion*planes)
# 判断是否需要下采样，相比于普通的判断方式，多了一个block类型
self.downsample = nn.Sequential()
if stride != 1 or in_planes != self.expansion*planes or block_type=='B':
self.downsample = nn.Sequential(
nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(self.expansion*planes)
)
def forward(self, x):
# 前向传播
out = F.relu(self.bn1(self.conv1(x)))
out = F.relu(self.bn2(self.conv2(out)))
out = self.bn3(self.conv3(out))
# 相加与下采样
out += self.downsample(x)
out = F.relu(out)
return out
# 5. 构建ResNet
class ResNet(nn.Module):
def __init__(self, block, layers):
'''
:param block: 即基本的Block块对象
:param layers: 指的是创建的何种ResNet，以及其对应的各个层的个数，比如ResNet50，传入的就是[3, 4, 6, 3]
'''
super(ResNet, self).__init__()
# 最开始的通道数，为64
self.inplanes = 64
# 最开始大家都用到的卷积层和池化层
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
# 开始定义不同的block块
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
# 不要忘记我们最后定义的output_block
self.layer5 = self._make_out_layer(in_channels=2048)
# 接上最后的卷积层即可，将输出变为30个通道数，shape为7*7*30
self.avgpool = nn.AvgPool2d(2) # kernel_size = 2 , stride = 2
self.conv_end = nn.Conv2d(256, 30, kernel_size=3, stride=1, padding=1, bias=False)
self.bn_end = nn.BatchNorm2d(30)
# 进行参数初始化
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
# 根据传入的layer个数和block创建
def _make_layer(self, block, planes, blocks, stride=1):
'''
:param block: Block对象
:param planes: 输入的通道数
:param blocks: 即需要搭建多少个一样的块
:param stride: 步长
'''
# 初始化下采样变量
downsample = None
# 判断是否需要进行下采样，即根据步长或者输入与输出通道数是否匹配
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
# 如果需要下采样，目的肯定是残差和输出可以加在一起
nn.Conv2d(self.inplanes, planes * block.expansion,kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(planes * block.expansion),
)
# 开始创建
layers = []
# 第一个block需要特别处理：
# 比如第一个，传入的channel为512，但是最终的输出为256，那么是需要下采样的
# 但是对于第二个block块，传入的肯定是第一个的输出即256，而最终输出也为256，因此不需要下采样
layers.append(block(self.inplanes, planes, stride, downsample))
self.inplanes = planes * block.expansion
for i in range(1, blocks):
# 重复指定次数
layers.append(block(self.inplanes, planes))
return nn.Sequential(*layers) # *表示解码，即将列表解开为一个个对象
# 输出层的构建
def _make_out_layer(self, in_channels):
layers = []
# 根据需求，构建出类似与block的即可
layers.append(Output_Block(in_planes=in_channels, planes=256, block_type='B'))
layers.append(Output_Block(in_planes=256, planes=256, block_type='A'))
layers.append(Output_Block(in_planes=256, planes=256, block_type='A'))
return nn.Sequential(*layers)
def forward(self, x):
# 经历共有的卷积和池化层
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
# 经历各个block块
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.layer5(x)
# 经历最终的输出
x = self.avgpool(x)
x = self.conv_end(x)
x = self.bn_end(x)
x = F.sigmoid(x) # 归一化到0-1
# 将输出构建为正确的shape
x = x.permute(0, 2, 3, 1) # (-1,7,7,30)
return x
# 6. 构建不同的ResNet函数
# 预训练下载链接
model_urls = {
'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
}
# 构建ResNet50
def resnet50(pretrained=False, **kwargs):
model = ResNet(Senior_Block, [3, 4, 6, 3], **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
return model
# 7.主函数测试功能
def main():
# 生成一个随机张量，值在[0, 1)之间
tensor = np.random.rand(448, 448, 3)
tensor = torch.tensor(tensor, dtype=torch.float32) # 转换为 PyTorch tensor
tensor = tensor.permute(2, 0, 1).unsqueeze(0) # 改变形状为 (1, 3, 448, 448)
model = resnet50() # 创建自己的resnet50，
output = model(tensor)
print(output.shape)
if __name__ == '__main__':
main()
输出效果如图所示。
3.损失函数构建
损失函数就是围绕1.边界框位置损失；2.分类损失；3.置信度损失而构造。代码中详细概述了它的步骤。
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
# 2. 损失函数类
class Yolo_Loss(nn.Module):
def __init__(self,S=7, B=2, l_coord=5, l_noobj=0.5):
'''
:param S: Yolov1论文中的S，即划分的网格，默认为 7
:param B: Yolov1论文中的B，即多少个框预测，默认为 2
:param l_coord: 损失函数中的超参数，默认为 5
:param l_noobj: 同上，默认为 0.5
'''
super(Yolo_Loss, self).__init__()
# 初始化各个参数
self.S = S
self.B = B
self.l_coord = l_coord
self.l_noobj = l_noobj
# 前向传播
def forward(self,pred_tensor,target_tensor):
# 获取batchsize大小
N = pred_tensor.size()[0]
# 具有目标标签的索引值，此时shape为[batch,7,7]
coo_mask = target_tensor[:, :, :, 4] > 0
# 不具有目标的标签索引值，此时shape为[batch,7,7]
noo_mask = target_tensor[:, :, :, 4] == 0
# 将shape变为[batch,7,7,30]
coo_mask = coo_mask.unsqueeze(-1).expand_as(target_tensor)
noo_mask = noo_mask.unsqueeze(-1).expand_as(target_tensor)
# 获取预测值中包含对象的所有点（共7*7个点），并转为[x,30]的形式，其中x表示有多少点中的框包含有对象
coo_pred = pred_tensor[coo_mask].view(-1, 30)
# 对上面获取的值进行处理
# 1. 转为box形式：box[x1,y1,w1,h1,c1]，shape为[2x,5]，因为每个单元格/点有两个预测框
box_pred = coo_pred[:, :10].contiguous().view(-1, 5)
# 2. 转为class信息，即30中后面的20个值
class_pred = coo_pred[:, 10:]
# 同理，对真实值进行操作，方便对比计算损失值
coo_target = target_tensor[coo_mask].view(-1, 30)
box_target = coo_target[:, :10].contiguous().view(-1, 5)
class_target = coo_target[:, 10:]
# 同上的操作，获取不包含对象的预测值、真实值
noo_pred = pred_tensor[noo_mask].view(-1, 30)
noo_target = target_tensor[noo_mask].view(-1, 30)
# 不包含物体grid ceil的置信度损失：即图中的D部分
# 1. 自己创建一个索引
noo_pred_mask = torch.cuda.ByteTensor(noo_pred.size())
noo_pred_mask.zero_() # 将全部元素变为Flase的意思
# 2. 将其它位置的索引置为0，唯独两个框的置信度位置变为1
noo_pred_mask[:, 4] = 1
noo_pred_mask[:, 9] = 1
# 3. 获取对应的值
noo_pred_c = noo_pred[noo_pred_mask] # noo pred只需要计算 c 的损失 size[-1,2]
noo_target_c = noo_target[noo_pred_mask]
# 4. 计算损失值：均方误差
nooobj_loss = F.mse_loss(noo_pred_c, noo_target_c, size_average=False)
# 计算包含物体的损失值
# 创建几个全为False/0的变量，用于后期存储值
coo_response_mask = torch.cuda.ByteTensor(box_target.size()) # 负责预测框
coo_response_mask.zero_()
coo_not_response_mask = torch.cuda.ByteTensor(box_target.size()) # 不负责预测的框的索引（因为一个cell两个预测框，而只有IOU最大的负责索引）
coo_not_response_mask.zero_()
box_target_iou = torch.zeros(box_target.size()).cuda() # 具体的IOU值存放处
# 由于一个单元格两个预测框，因此step=2
for i in range(0, box_target.size()[0], 2): # choose the best iou box
# 获取预测值中的两个box
box1 = box_pred[i:i + 2] # [x,y,w,h,c]
# 创建一个临时变量，用于存储中左上角+右下角坐标值，因为计算IOU需要
box1_xyxy = Variable(torch.FloatTensor(box1.size()))
# 下面将中心坐标+高宽 转为 左上角+右下角坐标的形式，并归一化
box1_xyxy[:, :2] = box1[:, :2] / float(self.S) - 0.5 * box1[:, 2:4] # 原本(xc,yc)为7*7 所以要除以7
box1_xyxy[:, 2:4] = box1[:, :2] / float(self.S) + 0.5 * box1[:, 2:4]
# 用同样的思路对真实值进行处理，不过不同的是真实值一个对象只有一个框
box2 = box_target[i].view(-1, 5)
box2_xyxy = Variable(torch.FloatTensor(box2.size()))
box2_xyxy[:, :2] = box2[:, :2] / float(self.S) - 0.5 * box2[:, 2:4]
box2_xyxy[:, 2:4] = box2[:, :2] / float(self.S) + 0.5 * box2[:, 2:4]
# 计算两者的IOU
iou = self.compute_iou(box1_xyxy[:, :4], box2_xyxy[:, :4]) # 前者shape为[2,4]，后者为[1,4]
# 获取两者IOU最大的值和索引，因为一个cell有两个预测框，一般而言取IOU最大的作为预测框
max_iou, max_index = iou.max(0)
max_index = max_index.data.cuda()
# 将IOU最大的索引设置为1，即表示这个框负责预测
coo_response_mask[i + max_index] = 1
# 将不是IOU最大的索引设置为1，即表示这个预测框不负责预测
coo_not_response_mask[i + 1 - max_index] = 1
# 获取具体的IOU值
box_target_iou[i + max_index, torch.LongTensor([4]).cuda()] = (max_iou).data.cuda()
box_target_iou = Variable(box_target_iou).cuda()
# 获取负责预测框的值、IOU值和真实框的值
box_pred_response = box_pred[coo_response_mask].view(-1, 5)
box_target_response_iou = box_target_iou[coo_response_mask].view(-1, 5)
box_target_response = box_target[coo_response_mask].view(-1, 5)
# 这个对应的是图中的部分C，负责预测框的损失
contain_loss = F.mse_loss(box_pred_response[:, 4], box_target_response_iou[:, 4], size_average=False)
# 1. 计算坐标损失，即图中的A和B部分
loc_loss = F.mse_loss(box_pred_response[:, :2], box_target_response[:, :2], size_average=False) + F.mse_loss(
torch.sqrt(box_pred_response[:, 2:4]), torch.sqrt(box_target_response[:, 2:4]), size_average=False)
# 获取不负责预测框的值、真实值
box_pred_not_response = box_pred[coo_not_response_mask].view(-1, 5)
box_target_not_response = box_target[coo_not_response_mask].view(-1, 5)
box_target_not_response[:, 4] = 0 # 将真实值置为0
# 2. 计算不负责预测框的损失值,即图中的部分C
not_contain_loss = F.mse_loss(box_pred_not_response[:, 4], box_target_not_response[:, 4], size_average=False)
# 3. 类别损失，即图中的E部分
class_loss = F.mse_loss(class_pred, class_target, size_average=False)
return (self.l_coord * loc_loss + contain_loss + not_contain_loss + self.l_noobj * nooobj_loss + class_loss) / N
# 计算IOU的函数
def compute_iou(self, box1, box2):
'''
:param box1: 预测的box，一般为[2,4]
:param box2: 真实的box，一般为[1,4]
:return:
'''
# 获取各box个数
N = box1.size(0)
M = box2.size(0)
# 计算两者中左上角左边较大的
lt = torch.max(
box1[:, :2].unsqueeze(1).expand(N, M, 2), # [N,2] -> [N,1,2] -> [N,M,2]
box2[:, :2].unsqueeze(0).expand(N, M, 2), # [M,2] -> [1,M,2] -> [N,M,2]
)
# 计算两者右下角左边较小的
rb = torch.min(
box1[:, 2:].unsqueeze(1).expand(N, M, 2), # [N,2] -> [N,1,2] -> [N,M,2]
box2[:, 2:].unsqueeze(0).expand(N, M, 2), # [M,2] -> [1,M,2] -> [N,M,2]
)
# 计算两者相交部分的长、宽
wh = rb - lt # [N,M,2]
# 如果长、宽中有小于0的，表示可能没有相交趋于，置为0即可
wh[wh < 0] = 0 # clip at 0
inter = wh[:, :, 0] * wh[:, :, 1] # [N,M]
# 计算各个的面积
# box1的面积
area1 = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1]) # [N,]
# box2的面积
area2 = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1]) # [M,]
area1 = area1.unsqueeze(1).expand_as(inter) # [N,] -> [N,1] -> [N,M]
area2 = area2.unsqueeze(0).expand_as(inter) # [M,] -> [1,M] -> [N,M]
# IOu值，交集除以并集，其中并集为两者的面积和减去交集部分
iou = inter / (area1 + area2 - inter)
return iou
4.训练部分
1.导入头文件
# -------------------------1.导入头文件---------------------------------
import warnings
warnings.filterwarnings("ignore")
from tqdm import tqdm
import torch
from torch.utils.data import DataLoader
import torchvision.transforms as T
from torchvision import models
from resnet50 import resnet50
from yolov1_loss import Yolo_Loss
from dataset import Yolo_Dataset
- 定义基本参数
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
batch_size = 4 # 根据自己的电脑设定
epochs = 1
lr = 0.01
file_root = 'D:\\Pycharm\\YOLOv1_pytorch\\dataset\\VOCtrainval_11-May-2012\\VOCdevkit\\VOC2012\\JPEGImages\\' # 需要根据的实际路径修改
- 创建模型并继承预训练参数
# -------------3. 创建模型并继承预训练参数------------------
# 加载预训练模型
# 接下来就是让自己的模型去继承权重参数
model = resnet50() # 创建模型实例
model.load_state_dict(torch.load('D:\Pycharm\yolov1\save_weights\yolo.pth')) # 加载权重
print('预训练模型加载完成')
model.to(device) # 转移到设备
- 将模型等放入GPU中
# -------------------4. 将模型等放入GPU中--------------
loss = Yolo_Loss()
optimizer = torch.optim.SGD(model.parameters(),lr=lr,momentum=0.9,weight_decay=5e-4)
model.to(device)
loss.to(device)
# 5. 加载数据
train_dataset = Yolo_Dataset(root=file_root,list_file='./voctrain.txt',train=True,transforms = [T.ToTensor()])
train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True,drop_last=True)
test_dataset = Yolo_Dataset(root=file_root,list_file='./voctest.txt',train=False,transforms = [T.ToTensor()])
test_loader = DataLoader(test_dataset,batch_size=batch_size,shuffle=True,drop_last=True)
- 训练
# -------------------------5. 训练 ---------------------------
# 打印一些基本的信息
print('starting train the model')
print('the train_dataset has %d images' % len(train_dataset))
print('the batch_size is ',batch_size)
# 定义一个最佳损失值
best_test_loss = 0
# 开始训练
for e in range(epochs):
model.train()
# 调整学习率
if e == 20:
print('change the lr')
optimizer.param_groups[0]['lr'] /= 10
if e == 35:
print('change the lr')
optimizer.param_groups[0]['lr'] /= 10
# 进度条显示
tqdm_tarin = tqdm(train_loader)
# 定义损失变量
total_loss = 0.
for i,(images,target) in enumerate(tqdm_tarin):
# 将变量放入设备中
images,target = images.to(device),target.to(device)
# 训练--损失等
pred = model(images)
loss_value = loss(pred,target)
total_loss += loss_value.item()
optimizer.zero_grad()
loss_value.backward()
optimizer.step()
# 打印一下损失值
if (i+1) % 5 == 0:
tqdm_tarin.desc = 'train epoch[{}/{}] loss:{:.6f}'.format(e+1,epochs,total_loss/(i+1))
# 启用验证模式
model.eval()
validation_loss = 0.0
tqdm_test = tqdm(test_loader)
for i, (images, target) in enumerate(tqdm_test):
images, target = images.cuda(), target.cuda()
pred = model(images)
loss_value = loss(pred, target)
validation_loss += loss_value.item()
validation_loss /= len(test_loader)
# 显示验证集的损失值
print('In the test step,the average loss is %.6f' % validation_loss)
# --------------------- 6.保存模型---------------------------------
if best_test_loss > validation_loss:
best_test_loss = validation_loss
print('get best test loss %.5f' % best_test_loss)
torch.save(model.state_dict(), './save_weights/best.pth')
# 记得最后保存参数
torch.save(model.state_dict(), './save_weights/yolo1.pth')
5.dataset构建
首先这一部分代码的作用你得明白，它要实现的功能是输入任意大小的图片，将其统一变成[1, 3, 448, 448]供网络输入。图像处理时包括一系列变化，随机调整图像的亮度，随机调整图像的饱和度，随机对图像进行模糊处理，随机缩放图像的宽度，保持高度不变，随机裁剪图像等等。以上为图像处理部分，而标签处理要实现的是读取txt文件，将标签信息编码成[1, 7, 7, 30]。从而实现对图像和标签的预处理。dataset整体代码如下。
import os
import random
import numpy as np
import torch
import torchvision.transforms as T
from torch.utils.data import Dataset
import cv2
# 2. 构建数据加载器,负责加载数据、进行预处理和增强操作
class Yolo_Dataset(Dataset):
# 默认图片大小
image_size = 448
def __init__(self,root,list_file,train=True,transforms=None):
'''
:param root: 根目录，比如`./data/`，这个目录必须包含真正的图片
:param list_file: 即我们之前使用generate_txt_file生成的txt文件，这里刚好和本文件在同一目录
:param train: 是否为训练集，默认为True
:param transforms: 预处理方法，默认为None
'''
# 初始化各个参数
self.root = root
self.train = train
self.transform = transforms
self.fnames = []
self.boxes = []
self.labels = []
self.mean = (123, 117, 104) # RGB
# 打开文件，读取内容
with open(list_file) as f:
lines = f.readlines()
# 针对每一行进行处理
for line in lines:
splited = line.strip().split()
self.fnames.append(splited[0]) # 文件名字
# 读取文件的类别和box信息
# 每五个数据表示一个对象，因此下面都是关于5的处理
num_boxes = (len(splited) - 1) // 5 #判断里面有几个物体 （长度-标签）// 5
box=[]
label=[]
for i in range(num_boxes):
x = float(splited[1+5*i])
y = float(splited[2+5*i])
x2 = float(splited[3+5*i])
y2 = float(splited[4+5*i])
c = splited[5+5*i] # c是种类标签
box.append([x,y,x2,y2])
label.append(int(c)+1)
self.boxes.append(torch.Tensor(box))
self.labels.append(torch.LongTensor(label))
# 记录一下对象个数
self.num_samples = len(self.boxes)
def __len__(self):
# 返回对象个数
return self.num_samples
def __getitem__(self, idx):
# 随便获取一张图片
fname = self.fnames[idx]
# 打开图片，记得把路径拼接完整
# print(os.path.join(self.root + fname))
img = cv2.imread(os.path.join(self.root + fname))
# 获取图片的相关信息
boxes = self.boxes[idx].clone()
labels = self.labels[idx].clone()
# 如果是训练模式，需要进行图像的增强
# 需要注意的是，同时处理图像和box
if self.train:
img, boxes = self.random_flip(img, boxes) # 随机翻转
img, boxes = self.randomScale(img, boxes) # 随机缩放
img = self.randomBlur(img) # 随机模糊
img = self.RandomBrightness(img) # 随机调整亮度
img = self.RandomHue(img) # 随机调整色调
img = self.RandomSaturation(img) # 随机调整饱和度
img, boxes, labels = self.randomShift(img, boxes, labels) # 随机移动
img, boxes, labels = self.randomCrop(img, boxes, labels) # 随机裁剪
# 进行基本的处理
h, w, _ = img.shape
# 归一化
boxes /= torch.Tensor([w, h, w, h]).expand_as(boxes)
# 主要是CV2打开模式默认为BGR，因此需要转为RGB模式
img = self.BGR2RGB(img)
# 减去均值
img = self.subMean(img, self.mean)
# 将图片缩放到指定大小，即448*448
img = cv2.resize(img, (self.image_size, self.image_size))
# 需要特别处理，将各种信息变为yolov1需要的7*7*30
target = self.encoder(boxes, labels) # 7x7x30
# 最后，进行预处理操作
for t in self.transform:
img = t(img)
return img, target # 最终的返回变量
def encoder(self,boxes,labels):
'''
将边界框（boxes）和标签（labels）转换为YOLO模型所需的目标格式
boxes (tensor) [[x1,y1,x2,y2],[]]
labels (tensor) [...]
return 7x7x30
'''
grid_num = 7
# 先创建一个全为0的张量，后面进行填充即可
target = torch.zeros((grid_num,grid_num,30)) #7，7，30
# 缩放因子
cell_size = 1./grid_num
# 计算出w、h和中心点坐标 boxes:[[x1, y1, x2, y2], ...],(x2, y2) 是右下角坐标
# wh: 通过相减，计算得到每个边界框的宽度和高度，结果格式为 [[width, height], ...]，其中 width = x2 - x1，height = y2 - y1
# 2:从第二号位开始，:2到第二号位结束
wh = boxes[:,2:]-boxes[:,:2] # boxes[:, 2:]: 提取所有边界框的右下角坐标,boxes[:, :2]:提取所有边界框的左上角坐标
cxcy = (boxes[:,2:]+boxes[:,:2])/2 # cx = (x1 + x2) / 2，cy = (y1 + y2) / 2 tensor（1，2）
for i in range(cxcy.size()[0]):
cxcy_sample = cxcy[i] # 中心坐标
ij = (cxcy_sample/cell_size).ceil()-1 # 左上角坐标，需要乘以缩放因子得到归一化后的坐标
target[int(ij[1]),int(ij[0]),4] = 1
target[int(ij[1]),int(ij[0]),9] = 1
target[int(ij[1]),int(ij[0]),int(labels[i])+9] = 1
# 匹配到的网格的左上角相对坐标
xy = ij*cell_size
# 相对偏移量
delta_xy = (cxcy_sample -xy)/cell_size
target[int(ij[1]),int(ij[0]),2:4] = wh[i] # 设置当前网格的宽度和高度 box1
target[int(ij[1]),int(ij[0]),:2] = delta_xy # 设置相对偏移量
target[int(ij[1]),int(ij[0]),7:9] = wh[i] # 再次设置当前网格的宽度和高度 box2
target[int(ij[1]),int(ij[0]),5:7] = delta_xy
return target
# 下面是各种预处理算法，来自别人的代码直接拷贝过来的
def BGR2RGB(self, img):
return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
def BGR2HSV(self, img):
return cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
def HSV2BGR(self, img):
return cv2.cvtColor(img, cv2.COLOR_HSV2BGR)
# 随机调整图像的亮度
def RandomBrightness(self, bgr):
if random.random() < 0.5:
hsv = self.BGR2HSV(bgr)
h, s, v = cv2.split(hsv)
adjust = random.choice([0.5, 1.5])
v = v * adjust
v = np.clip(v, 0, 255).astype(hsv.dtype)
hsv = cv2.merge((h, s, v))
bgr = self.HSV2BGR(hsv)
return bgr
# 随机调整图像的饱和度
def RandomSaturation(self, bgr):
if random.random() < 0.5:
hsv = self.BGR2HSV(bgr)
h, s, v = cv2.split(hsv)
adjust = random.choice([0.5, 1.5])
s = s * adjust
s = np.clip(s, 0, 255).astype(hsv.dtype)
hsv = cv2.merge((h, s, v))
bgr = self.HSV2BGR(hsv)
return bgr
# 随机调整图像的色调
def RandomHue(self, bgr):
if random.random() < 0.5:
hsv = self.BGR2HSV(bgr)
h, s, v = cv2.split(hsv)
adjust = random.choice([0.5, 1.5])
h = h * adjust
h = np.clip(h, 0, 255).astype(hsv.dtype)
hsv = cv2.merge((h, s, v))
bgr = self.HSV2BGR(hsv)
return bgr
# 随机对图像进行模糊处理
def randomBlur(self, bgr):
if random.random() < 0.5:
bgr = cv2.blur(bgr, (5, 5))
return bgr
# 平移变换
def randomShift(self, bgr, boxes, labels):
center = (boxes[:, 2:] + boxes[:, :2]) / 2
if random.random() < 0.5:
height, width, c = bgr.shape
after_shfit_image = np.zeros((height, width, c), dtype=bgr.dtype)
after_shfit_image[:, :, :] = (104, 117, 123) # bgr
shift_x = random.uniform(-width * 0.2, width * 0.2)
shift_y = random.uniform(-height * 0.2, height * 0.2)
# print(bgr.shape,shift_x,shift_y)
# 原图像的平移
if shift_x >= 0 and shift_y >= 0:
after_shfit_image[int(shift_y):, int(shift_x):, :] = bgr[:height - int(shift_y), :width - int(shift_x),
:]
elif shift_x >= 0 and shift_y < 0:
after_shfit_image[:height + int(shift_y), int(shift_x):, :] = bgr[-int(shift_y):, :width - int(shift_x),
:]
elif shift_x < 0 and shift_y >= 0:
after_shfit_image[int(shift_y):, :width + int(shift_x), :] = bgr[:height - int(shift_y), -int(shift_x):,
:]
elif shift_x < 0 and shift_y < 0:
after_shfit_image[:height + int(shift_y), :width + int(shift_x), :] = bgr[-int(shift_y):,
-int(shift_x):, :]
shift_xy = torch.FloatTensor([[int(shift_x), int(shift_y)]]).expand_as(center)
center = center + shift_xy
mask1 = (center[:, 0] > 0) & (center[:, 0] < width)
mask2 = (center[:, 1] > 0) & (center[:, 1] < height)
mask = (mask1 & mask2).view(-1, 1)
boxes_in = boxes[mask.expand_as(boxes)].view(-1, 4)
if len(boxes_in) == 0:
return bgr, boxes, labels
box_shift = torch.FloatTensor([[int(shift_x), int(shift_y), int(shift_x), int(shift_y)]]).expand_as(
boxes_in)
boxes_in = boxes_in + box_shift
labels_in = labels[mask.view(-1)]
return after_shfit_image, boxes_in, labels_in
return bgr, boxes, labels
# 随机缩放图像的宽度，保持高度不变
def randomScale(self, bgr, boxes):
# 固定住高度，以0.8-1.2伸缩宽度，做图像形变
if random.random() < 0.5:
scale = random.uniform(0.8, 1.2)
height, width, c = bgr.shape
bgr = cv2.resize(bgr, (int(width * scale), height))
scale_tensor = torch.FloatTensor([[scale, 1, scale, 1]]).expand_as(boxes)
boxes = boxes * scale_tensor
return bgr, boxes
return bgr, boxes
# 随机裁剪图像
def randomCrop(self, bgr, boxes, labels):
if random.random() < 0.5:
center = (boxes[:, 2:] + boxes[:, :2]) / 2
height, width, c = bgr.shape
h = random.uniform(0.6 * height, height)
w = random.uniform(0.6 * width, width)
x = random.uniform(0, width - w)
y = random.uniform(0, height - h)
x, y, h, w = int(x), int(y), int(h), int(w)
center = center - torch.FloatTensor([[x, y]]).expand_as(center)
mask1 = (center[:, 0] > 0) & (center[:, 0] < w)
mask2 = (center[:, 1] > 0) & (center[:, 1] < h)
mask = (mask1 & mask2).view(-1, 1)
boxes_in = boxes[mask.expand_as(boxes)].view(-1, 4)
if (len(boxes_in) == 0):
return bgr, boxes, labels
box_shift = torch.FloatTensor([[x, y, x, y]]).expand_as(boxes_in)
boxes_in = boxes_in - box_shift
boxes_in[:, 0] = boxes_in[:, 0].clamp_(min=0, max=w)
boxes_in[:, 2] = boxes_in[:, 2].clamp_(min=0, max=w)
boxes_in[:, 1] = boxes_in[:, 1].clamp_(min=0, max=h)
boxes_in[:, 3] = boxes_in[:, 3].clamp_(min=0, max=h)
labels_in = labels[mask.view(-1)]
img_croped = bgr[y:y + h, x:x + w, :]
return img_croped, boxes_in, labels_in
return bgr, boxes, labels
# 从图像中减去均值，用于图像标准化
def subMean(self, bgr, mean):
mean = np.array(mean, dtype=np.float32)
bgr = bgr - mean
return bgr
# 随机水平翻转图像
def random_flip(self, im, boxes):
if random.random() < 0.5:
im_lr = np.fliplr(im).copy()
h, w, _ = im.shape
xmin = w - boxes[:, 2]
xmax = w - boxes[:, 0]
boxes[:, 0] = xmin
boxes[:, 2] = xmax
return im_lr, boxes
return im, boxes
# 随机调整图像的亮度
def random_bright(self, im, delta=16):
alpha = random.random()
if alpha > 0.3:
im = im * alpha + random.randrange(-delta, delta)
im = im.clip(min=0, max=255).astype(np.uint8)
return im
# 3. 调试代码
def main():
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
file_root = 'D:/Pycharm/YOLOv1_pytorch/dataset/VOCtrainval_11-May-2012/VOCdevkit/VOC2012/JPEGImages/' # 记得改为自己的路径
train_dataset = Yolo_Dataset(root=file_root,list_file='voctrain.txt',train=True,transforms = [T.ToTensor()] )
train_loader = DataLoader(train_dataset,batch_size=8,shuffle=False,num_workers=0)
train_iter = iter(train_loader)
# torch.Size([1, 3, 448, 448])
# torch.Size([1, 7, 7, 30]) 这个就好比你真正要训练的东西。你要训练的是[1, 7, 7, 30]的这么一个list
img,target = next(train_iter)
print(img.shape)
print(target.shape)
if __name__ == '__main__':
main()
训练效果
测试部分
# -----------------1. 导入头文件--------------------
import os
import random
import torch
from torch.autograd import Variable
import torchvision.transforms as transforms
import cv2
from matplotlib import pyplot as plt
import numpy as np
import warnings
warnings.filterwarnings('ignore')
from resnet50 import resnet50
# --------------2. 定义一些基本参数-------------------
# 类别索引
VOC_CLASSES = (
'aeroplane', 'bicycle', 'bird', 'boat',
'bottle', 'bus', 'car', 'cat', 'chair',
'cow', 'diningtable', 'dog', 'horse',
'motorbike', 'person', 'pottedplant',
'sheep', 'sofa', 'train', 'tvmonitor')
# -------------------画矩形框的时候用到的颜色变量-------------------------
Color = [[0, 0, 0],
[128, 0, 0],
[0, 128, 0],
[128, 128, 0],
[0, 0, 128],
[128, 0, 128],
[0, 128, 128],
[128, 128, 128],
[64, 0, 0],
[192, 0, 0],
[64, 128, 0],
[192, 128, 0],
[64, 0, 128],
[192, 0, 128],
[64, 128, 128],
[192, 128, 128],
[0, 64, 0],
[128, 64, 0],
[0, 192, 0],
[128, 192, 0],
[0, 64, 128]]
# -----------------------------3. 解码函数----------------------------------
def decoder(pred):
'''
:param pred: batchx7x7x30，但是预测的时候一般一张图片一张的放，因此batch=1
:return: box[[x1,y1,x2,y2]] label[...]
'''
# 定义一些基本的参数
grid_num = 7 # 网格划分标准大小
boxes=[]
cls_indexs=[]
probs = []
cell_size = 1./grid_num # 缩放因子
# 获取一些值
pred = pred.data # 预测值的数据：1*7*7*30
pred = pred.squeeze(0) # 预测值的数据：7x7x30
contain1 = pred[:,:,4].unsqueeze(2) # 先获取第一个框的置信度，然后升维变为7*7*1
contain2 = pred[:,:,9].unsqueeze(2) # 同上，只是为第二个框
contain = torch.cat((contain1,contain2),2) # 拼接在一起，变为7*7*2
mask1 = contain > 0.05 #大于阈值0.1，设置为True
mask2 = (contain==contain.max()) # 找出置信度最大的，设置为True
mask = (mask1+mask2).gt(0) # 将mask1+mask2，让其中大于0的设置为True
# 开始迭代每个单元格，即7*7个
for i in range(grid_num):
for j in range(grid_num):
# 迭代两个预测框
for b in range(2):
# 如果mask为1，表示这个框是最大的置信度框
if mask[i,j,b] == 1:
# 获取坐标值
box = pred[i,j,b*5:b*5+4]
# 获取置信度值
contain_prob = torch.FloatTensor([pred[i,j,b*5+4]])
# 将7*7的坐标，归一化
xy = torch.FloatTensor([j,i])*cell_size #cell左上角 up left of cell
#
box[:2] = box[:2]*cell_size + xy
# 将[cx,cy,w,h]转为[x1,xy1,x2,y2]
box_xy = torch.FloatTensor(box.size()) # 重新创建一个变量存储值
box_xy[:2] = box[:2] - 0.5*box[2:] # 这个就是中心坐标加减宽度/高度得到左上角/右下角坐标
box_xy[2:] = box[:2] + 0.5*box[2:]
# 获取最大的概率和类别索引值
max_prob,cls_index = torch.max(pred[i,j,10:],0)
# 如果置信度 * 类别概率 > 0.1，即说明有一定的可信度
# 那么把值加入各个变量列表中
if float((contain_prob*max_prob)[0]) > 0.1:
boxes.append(box_xy.view(1,4))
cls_indexs.append(torch.tensor([cls_index.item()]))
probs.append(contain_prob*max_prob)
# 如果boxes为0，表示没有框，返回0
if len(boxes) ==0:
boxes = torch.zeros((1,4))
probs = torch.zeros(1)
cls_indexs = torch.zeros(1)
# 否则，进行处理，就是简单把原来的列表值[tensor,tensor]转为tensor的形式
# 里面的值不变
else:
boxes = torch.cat(boxes,0) #(n,4)
probs = torch.cat(probs,0) #(n,)
cls_indexs = torch.cat(cls_indexs,0) #(n,)
# 后处理——NMS
keep = nms(boxes,probs)
# 返回值
return boxes[keep],cls_indexs[keep],probs[keep]
# --------------------------------4. NMS处理---------------------------
def nms(bboxes,scores,threshold=0.5):
'''
:param bboxes: bboxes(tensor) [N,4]
:param scores: scores(tensor) [N,]
:param threshold: 阈值
:return: 返回过滤后的框
'''
# 获取各个框的坐标值
x1 = bboxes[:,0]
y1 = bboxes[:,1]
x2 = bboxes[:,2]
y2 = bboxes[:,3]
# 计算面积
areas = (x2-x1) * (y2-y1)
# 将置信度按照降序排序，并获取排序后的各个置信度在这个顺序中的索引
_,order = scores.sort(0,descending=True)
keep = []
# 判断order中的元素个数是否大于0
while order.numel() > 0:
# 如果元素个数只剩下一个了，结束循环
if order.numel() == 1:
i = order.item()
keep.append(i)
break
# 获取最大置信度的索引
i = order[0]
keep.append(i)
# 对后面的元素坐标进行截断处理
xx1 = x1[order[1:]].clamp(min=x1[i]) # min指的是小于它的设置为它的值，大于它的不管
yy1 = y1[order[1:]].clamp(min=y1[i])
xx2 = x2[order[1:]].clamp(max=x2[i])
yy2 = y2[order[1:]].clamp(max=y2[i])
# 此时的xx1，yy1等是排除了目前选中的框的，即假设x1有三个元素，那么xx1只有2个元素
# 获取排序后的长和宽以及面积，如果小于0则设置为0
w = (xx2-xx1).clamp(min=0)
h = (yy2-yy1).clamp(min=0)
inter = w*h
# 准备更新order、
# 计算选中的框和剩下框的IOU值
ovr = inter / (areas[i] + areas[order[1:]] - inter)
# 如果 IOU小于设定的阈值，说明需要保存下来继续筛选（NMS原理）
ids = (ovr<=threshold).nonzero().squeeze()
if ids.numel() == 0:
break
order = order[ids+1]
return torch.LongTensor(keep)
# -----------------------5. 预测函数------------------------------
def predict_single(model, image_name, root_path=''):
result = [] # 保存结果的变量
# 打开图片
image = cv2.imread(root_path + image_name)
h, w, _ = image.shape
# resize为模型的输入大小，即448*448
img = cv2.resize(image, (448, 448))
# 由于我们模型那里定义的颜色模式为RGB，因此这里需要转换
mean = (123, 117, 104) # RGB均值
img = img - np.array(mean, dtype=np.float32)
# 预处理
transform = transforms.Compose([transforms.ToTensor(), ])
img = transform(img)
img = Variable(img[None, :, :, :], volatile=True)
img = img.cuda()
# 开始预测
pred = model(img) # 1x7x7x30
pred = pred.cpu()
# 解码
boxes, cls_indexs, probs = decoder(pred)
# 开始迭代每个框
for i, box in enumerate(boxes):
# 获取相关坐标，只是需要把原来归一化后的坐标转回去
x1 = int(box[0] * w)
x2 = int(box[2] * w)
y1 = int(box[1] * h)
y2 = int(box[3] * h)
# 获取类别索引、概率等值
cls_index = cls_indexs[i]
cls_index = int(cls_index) # convert LongTensor to int
prob = probs[i]
prob = float(prob)
# 把这些值集中放入一个变量中返回
result.append([(x1, y1), (x2, y2), VOC_CLASSES[cls_index], image_name, prob])
return result
if __name__ == '__main__':
# 慢慢的显示
import time
# 创建模型，加载参数
model = resnet50()
model.load_state_dict(torch.load('./save_weights/yolo.pth'))
model.eval()
model.cuda()
# 设置图片路径
# base_path = './test_images/1.jpg'
# image_name_list = [base_path]
base_path = './test_images/'
image_name_list = [base_path+i for i in os.listdir(base_path)]
# 打乱顺序
random.shuffle(image_name_list)
print('stating predicting....')
for image_name in image_name_list:
image = cv2.imread(image_name)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
result = predict_single(model, image_name)
# 画矩形框和对应的类别信息
for left_up, right_bottom, class_name, _, prob in result:
# 获取颜色
color = Color[VOC_CLASSES.index(class_name)]
# 画矩形
cv2.rectangle(image, left_up, right_bottom, color, 2)
# 获取类型信息和对应概率，此时为str类型
label = class_name + str(round(prob, 2))
# 把类别和概率信息写上，还要为这个信息加上一个矩形框
text_size, baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.4, 1)
p1 = (left_up[0], left_up[1] - text_size[1])
cv2.rectangle(image, (p1[0] - 2 // 2, p1[1] - 2 - baseline), (p1[0] + text_size[0], p1[1] + text_size[1]),
color, -1)
cv2.putText(image, label, (p1[0], p1[1] + baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1, 8)
# 显示图片
plt.figure()
plt.imshow(image)
plt.show()
time.sleep(2)
# 是否保存结果图片
# cv2.imwrite('./test_images/result.jpg', image)
测试效果
对大物体也有识别不出来的时候。
以上就是本人的心得与总结，如有不足之处请多多包涵。
百度网盘链接代码权值及测试图片，训练图片链接:
链接: https://pan.baidu.com/s/1EAVYEXPfkJIGlG0a56gKhw?pwd=fbqg
提取码: fbqg
原文地址：https://blog.csdn.net/weixin_52531699/article/details/143233014
免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！