学习YOLOV3源码
YOLOV3的网络结构
一、左边是Darknet-53结构
二、右边可分为FPN特征金字塔进行加强特征提取和Yolo Head对三个有效特征层进行预测
a、构建FPN特征金字塔进行加强特征提取
在特征利用部分,YoloV3提取多特征层进行目标检测,一共提取三个特征层。
三个特征层位于主干部分Darknet53的不同位置,分别位于中间层,中下层,底层,三个特征层的shape分别为(52,52,256)、(26,26,512)、(13,13,1024)。
在获得三个有效特征层后,我们利用这三个有效特征层进行FPN层的构建,构建方式为:
1.13x13x1024的特征层进行5次卷积处理,处理完后利用YoloHead获得预测结果,一部分用于进行上采样UmSampling2d后与26x26x512特征层进行结合,结合特征层的shape为(26,26,768)。
2.结合特征层再次进行5次卷积处理,处理完后利用YoloHead获得预测结果,一部分用于进行上采样UmSampling2d后与52x52x256特征层进行结合,结合特征层的shape为(52,52,384)。
3.结合特征层再次进行5次卷积处理,处理完后利用YoloHead获得预测结果。
特征金字塔可以将不同shape的特征层进行特征融合,有利于提取出更好的特征。
b、利用Yolo Head获得预测结果
利用FPN特征金字塔,我们可以获得三个加强特征,这三个加强特征的shape分别为(13,13,512)、(26,26,256)、(52,52,128),然后我们利用这三个shape的特征层传入Yolo Head获得预测结果。
Yolo Head本质上是一次3x3卷积加上一次1x1卷积,3x3卷积的作用是特征整合,1x1卷积的作用是调整通道数。
对三个特征层分别进行处理,假设我们预测是的VOC数据集,我们的输出层的shape分别为(13,13,75),(26,26,75),(52,52,75),最后一个维度为75是因为该图是基于voc数据集的,它的类为20种,YoloV3针对每一个特征层的每一个特征点存在3个先验框,所以预测结果的通道数为3x25;
如果使用的是coco训练集,类则为80种,最后的维度应该为255 = 3x85,三个特征层的shape为(13,13,255),(26,26,255),(52,52,255)
其实际情况就是,输入N张416x416的图片,在经过多层的运算后,会输出三个shape分别为(N,13,13,255),(N,26,26,255),(N,52,52,255)的数据,对应每个图分为13x13、26x26、52x52的网格上3个先验框的位置。
Darknet-53结构代码: #(darknet.py文件)
import math
from collections import OrderedDict
import torch.nn as nn
#---------------------------------------------------------------------#
# 残差结构
# 利用一个1x1卷积下降通道数,然后利用一个3x3卷积提取特征并且上升通道数
# 最后接上一个残差边
#---------------------------------------------------------------------#
class BasicBlock(nn.Module):
def __init__(self, inplanes, planes):
super(BasicBlock, self).__init__()
self.conv1 = nn.Conv2d(inplanes, planes[0], kernel_size=1, stride=1, padding=0, bias=False)
self.bn1 = nn.BatchNorm2d(planes[0])
self.relu1 = nn.LeakyReLU(0.1)
self.conv2 = nn.Conv2d(planes[0], planes[1], kernel_size=3, stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes[1])
self.relu2 = nn.LeakyReLU(0.1)
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu1(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.relu2(out)
out += residual
return out
class DarkNet(nn.Module):
def __init__(self, layers): #layers = [1, 2, 8, 8, 4]
super(DarkNet, self).__init__()
self.inplanes = 32
# 416,416,3 -> 416,416,32
self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False) #通道数变化:3-->32
self.bn1 = nn.BatchNorm2d(self.inplanes) #参数inplanes=32表示特征通道数量
self.relu1 = nn.LeakyReLU(0.1)
# 416,416,32 -> 208,208,64
self.layer1 = self._make_layer(planes=[32, 64], blocks=layers[0]) #使用残差结构 1
# 208,208,64 -> 104,104,128
self.layer2 = self._make_layer([64, 128], layers[1]) # 2
# 104,104,128 -> 52,52,256
self.layer3 = self._make_layer([128, 256], layers[2]) # 8
# 52,52,256 -> 26,26,512
self.layer4 = self._make_layer([256, 512], layers[3]) # 8
# 26,26,512 -> 13,13,1024
self.layer5 = self._make_layer([512, 1024], layers[4]) # 4
self.layers_out_filters = [64, 128, 256, 512, 1024]
# 进行权值初始化
for m in self.modules():
if isinstance(m, nn.Conv2d): #isinstance(实例对象,直接或者间接类名、基本类型或者由它们组成的元组) 如isinstance(12,int) = Ture
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里面,首先利用一个步长为2的3x3卷积进行下采样
# 然后进行残差结构的堆叠
#---------------------------------------------------------------------#
def _make_layer(self, planes, blocks):
layers = []
# 下采样,步长为2,卷积核大小为3
layers.append(("ds_conv", nn.Conv2d(self.inplanes, planes[1], kernel_size=3, stride=2, padding=1, bias=False)))
layers.append(("ds_bn", nn.BatchNorm2d(planes[1])))
layers.append(("ds_relu", nn.LeakyReLU(0.1)))
# 加入残差结构
self.inplanes = planes[1]
for i in range(0, blocks): #layers = [1, 2, 8, 8, 4] , 第一次做self.layer1传入的blocks=layers[0]=1
layers.append(("residual_{}".format(i), BasicBlock(self.inplanes, planes))) #计算blocks个残差块,BasicBlock输出out
return nn.Sequential(OrderedDict(layers))
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu1(x)
x = self.layer1(x)
x = self.layer2(x)
out3 = self.layer3(x)
out4 = self.layer4(out3)
out5 = self.layer5(out4)
return out3, out4, out5
def darknet53():
model = DarkNet(layers = [1, 2, 8, 8, 4]) #1,2,8,8,4对应的是残差块的个数
return model
darknet.py文件解读:
该文件里有两个类(BasicBlock残差块和DarkNet网络结构)
1.BasicBlock:残差结构中分为主通道和捷径(shortcut),输入0分两条路分别进入主通道和捷径,主通道经过两次卷积、BN、Relu(得到主通道的输出1),捷径则是不对输入执行任何操作(捷径的输出2=输入0),最终的输出out=输出1+输出2
介绍里面的几种方法:
1. _make_layer方法:该方法传入的参数是planes(两个数的数组)和blocks(传入的是要残差块的数量,结构的残差数据是layers = [1, 2, 8, 8, 4],借助索引取出残差数)。首先是对输入进行一层卷积、BN、Relu,然后执行BasicBlock的残差操作,执行的残差次数是blocks
首先看到DarkNet类中的forward方法
1.外部传入图像x,按照步骤,先对输入图像经过一层卷积、BN、Relu
2.接着执行self.layer1(x)的操作,在layer1中使用到了_make_layer方法,接着执行layer2,再执行layer3(得到第一个输出out3),接着执行layer4(得到第二个输出out4),最后执行layer5(得到第三个输出out5)
注意到Darknet中传入的参数是残差块数量的数组layers = [1, 2, 8, 8, 4],在layer1中传入的参数blocks=layers[0]=1,即执行一次残差块的操作(一次BasicBlock),同理,layer2执行两次,layer3执行八次,layer4执行八次,layer5执行四次。
YOLO特征提取及预测代码: #(yolo.py文件)
from collections import OrderedDict
import torch
import torch.nn as nn
from nets.darknet import darknet53
def conv2d(filter_in, filter_out, kernel_size):
pad = (kernel_size - 1) // 2 if kernel_size else 0 #当卷积核kernel_size=3时,padding=2; kernel_size=1时,padding=0
return nn.Sequential(OrderedDict([ #正常进行卷积、正则化、relu操作
("conv", nn.Conv2d(filter_in, filter_out, kernel_size=kernel_size, stride=1, padding=pad, bias=False)),
("bn", nn.BatchNorm2d(filter_out)),
("relu", nn.LeakyReLU(0.1)),
]))
#------------------------------------------------------------------------#
# make_last_layers里面一共有七个卷积,前五个用于提取特征。
# 后两个用于获得yolo网络的预测结果
#------------------------------------------------------------------------#
def make_last_layers(filters_list, in_filters, out_filter):
m = nn.Sequential(
conv2d(filter_in=in_filters, filter_out=filters_list[0], kernel_size=1),
conv2d(filters_list[0], filters_list[1], 3),
conv2d(filters_list[1], filters_list[0], 1),
conv2d(filters_list[0], filters_list[1], 3),
conv2d(filters_list[1], filters_list[0], 1),
conv2d(filters_list[0], filters_list[1], 3),
nn.Conv2d(filters_list[1], out_filter, kernel_size=1, stride=1, padding=0, bias=True)
)
return m #那三个out最后经过的7次卷积(5+1+1),最终的维度为(batch_size,13或26或52,13或26或52,75)
class YoloBody(nn.Module):
def __init__(self, anchors_mask, num_classes, pretrained = False):
super(YoloBody, self).__init__()
#---------------------------------------------------#
# 生成darknet53的主干模型
# 获得三个有效特征层,他们的shape分别是:
# 52,52,256
# 26,26,512
# 13,13,1024
#---------------------------------------------------#
self.backbone = darknet53() #返回三个out
if pretrained:
self.backbone.load_state_dict(torch.load("model_data/darknet53_backbone_weights.pth"))
#---------------------------------------------------#
# out_filters : [64, 128, 256, 512, 1024]
#---------------------------------------------------#
out_filters = self.backbone.layers_out_filters
#------------------------------------------------------------------------#
# 计算yolo_head的输出通道数,对于voc数据集而言
# final_out_filter0 = final_out_filter1 = final_out_filter2 = 75
#------------------------------------------------------------------------#
self.last_layer0 = make_last_layers([512, 1024], out_filters[-1], len(anchors_mask[0]) * (num_classes + 5)) #filters_list=[512, 1024]
self.last_layer1_conv = conv2d(512, 256, 1) #通道数减半操作
self.last_layer1_upsample = nn.Upsample(scale_factor=2, mode='nearest') #scale_factor=指定输出的尺寸是输入尺寸的倍数(注意是尺寸不是通道数),这里是尺寸翻倍:13*13*256-->26*26*256
self.last_layer1 = make_last_layers([256, 512], out_filters[-2] + 256, len(anchors_mask[1]) * (num_classes + 5))
self.last_layer2_conv = conv2d(256, 128, 1)
self.last_layer2_upsample = nn.Upsample(scale_factor=2, mode='nearest')
self.last_layer2 = make_last_layers([128, 256], out_filters[-3] + 128, len(anchors_mask[2]) * (num_classes + 5))
def forward(self, x):
#---------------------------------------------------#
# 获得三个有效特征层,他们的shape分别是:
# 52,52,256;26,26,512;13,13,1024
#---------------------------------------------------#
x2, x1, x0 = self.backbone(x)
#---------------------------------------------------#
# 第一个特征层
# out0 = (batch_size,255,13,13)
#---------------------------------------------------#
# 13,13,1024 -> 13,13,512 -> 13,13,1024 -> 13,13,512 -> 13,13,1024 -> 13,13,512
out0_branch = self.last_layer0[:5](x0)
out0 = self.last_layer0[5:](out0_branch)
# 13,13,512 -> 13,13,256 -> 26,26,256
x1_in = self.last_layer1_conv(out0_branch)
x1_in = self.last_layer1_upsample(x1_in)
# 26,26,256 + 26,26,512 -> 26,26,768
x1_in = torch.cat([x1_in, x1], 1)
#---------------------------------------------------#
# 第二个特征层
# out1 = (batch_size,255,26,26)
#---------------------------------------------------#
# 26,26,768 -> 26,26,256 -> 26,26,512 -> 26,26,256 -> 26,26,512 -> 26,26,256
out1_branch = self.last_layer1[:5](x1_in)
out1 = self.last_layer1[5:](out1_branch)
# 26,26,256 -> 26,26,128 -> 52,52,128
x2_in = self.last_layer2_conv(out1_branch)
x2_in = self.last_layer2_upsample(x2_in)
# 52,52,128 + 52,52,256 -> 52,52,384
x2_in = torch.cat([x2_in, x2], 1)
#---------------------------------------------------#
# 第一个特征层
# out3 = (batch_size,255,52,52)
#---------------------------------------------------#
# 52,52,384 -> 52,52,128 -> 52,52,256 -> 52,52,128 -> 52,52,256 -> 52,52,128
out2 = self.last_layer2(x2_in)
return out0, out1, out2
yolo.py文件解读:
里面有一个YoloBody类,执行的是YOLO操作。
介绍里面的几种方法:
1.make_last_layers,执行7次卷积、BN、Relu操作(第七次只进行卷积操作)
2.conv2d,这个是里面定义的一个方法,不是pytorch里面的卷积操作,因为在YOLOV3里面的卷积核都是3*3或1*1,YOLOV3里没有maxpooling用的是卷积操作(3*3的卷积核,padding=1),
在conv2d这个方法中对卷积核kernel_size=3时,设置padding=2; kernel_size=1时,设置padding=0,并执行一次卷积、BN、Relu操作。
先看到forward中
1.首先将输入x放入darknet53()中返回三个返回值x2, x1, x0。
2.执行last_layer0[:5](x0),对x0执行make_last_layers的前五层卷积、BN、Relu操作,返回out0_branch。然后执行last_layer0[5:](out0_branch),将out0_branch作为输入执行make_last_layers的后两层卷积、BN、Relu操作,返回out0(维度是13,13,75)。
3.执行last_layer1_conv(out0_branch),返回x1_in,再对这个返回值执行last_layer1_upsample上采样操作(对输入尺寸翻倍),最后执行拼接操作torch.cat([x1_in, x1], 1),输出x1_in。然后执行last_layer1[:5](x1_in)操作,对x1_in执行make_last_layers的前五层卷积、BN、Relu操作,返回out1_branch。然后执行last_layer1[5:](out0_branch),将out1_branch作为输入执行make_last_layers的后两层卷积、BN、Relu操作,返回out1(维度是26,26,75)。
4.执行last_layer2_conv(out1_branch),返回x2_in,再对这个返回值执行last_layer2_upsample上采样操作(对输入尺寸翻倍),最后执行拼接操作torch.cat([x2_in, x2], 1),输出x2_in。然后执行last_layer2(x2_in)操作,对x2_in执行make_last_layers的七层卷积、BN、Relu操作,返回out2_branch,返回out2(维度是26,26,75)。
最终返回三个输出out0、out1、out2
yolo_training.py文件
import torch
import torch.nn as nn
import math
import numpy as np
class YOLOLoss(nn.Module): #我的注释以感受野最大的输入out0为例 out0.shape=(bs,75,13,13)
def __init__(self, anchors, num_classes, input_shape, cuda, anchors_mask = [[6,7,8], [3,4,5], [0,1,2]]):
super(YOLOLoss, self).__init__()
#-----------------------------------------------------------#
# 13x13的特征层对应的anchor是[116,90],[156,198],[373,326]
# 26x26的特征层对应的anchor是[30,61],[62,45],[59,119]
# 52x52的特征层对应的anchor是[10,13],[16,30],[33,23]
#-----------------------------------------------------------#
self.anchors = anchors #[116,90],[156,198],[373,326]
self.num_classes = num_classes #
self.bbox_attrs = 5 + num_classes
self.input_shape = input_shape #batch_size,3*(5+num_classes),13,13
self.anchors_mask = anchors_mask #[[6,7,8], [3,4,5], [0,1,2]] 框的索引
self.ignore_threshold = 0.5
self.cuda = cuda
def clip_by_tensor(self, t, t_min, t_max):
t = t.float()
result = (t >= t_min).float() * t + (t < t_min).float() * t_min
result = (result <= t_max).float() * result + (result > t_max).float() * t_max
return result
def MSELoss(self, pred, target):
return torch.pow(pred - target, 2)
def BCELoss(self, pred, target):
epsilon = 1e-7
pred = self.clip_by_tensor(pred, epsilon, 1.0 - epsilon)
output = - target * torch.log(pred) - (1.0 - target) * torch.log(1.0 - pred)
return output
def forward(self, l, input, targets=None):
#----------------------------------------------------#
# l代表的是,当前输入进来的有效特征层,是第几个有效特征层
# input的shape为 bs, 3*(5+num_classes), 13, 13 yolo.py里的三个输出 out0, out1, out2 out0.shape=(bs,75,13,13)
# bs, 3*(5+num_classes), 26, 26
# bs, 3*(5+num_classes), 52, 52
# targets代表的是真实框。
#----------------------------------------------------#
#--------------------------------#
# 获得图片数量,特征层的高和宽
# 13和13
#--------------------------------#
bs = input.size(0) #batch_size
in_h = input.size(2) #13
in_w = input.size(3) #13
#-----------------------------------------------------------------------#
# 计算步长
# 每一个特征点对应原来的图片上多少个像素点
# 如果特征层为13x13的话,一个特征点就对应原来的图片上的32个像素点
# 如果特征层为26x26的话,一个特征点就对应原来的图片上的16个像素点
# 如果特征层为52x52的话,一个特征点就对应原来的图片上的8个像素点
# stride_h = stride_w = 32、16、8
# stride_h和stride_w都是32。
#-----------------------------------------------------------------------#
stride_h = self.input_shape[0] / in_h #32 原始图片的高(416)/特征图的网格行数(特征图的高13) = 感受野(特征图上一个网格对应原始图像的大小32)
stride_w = self.input_shape[1] / in_w #32
#-------------------------------------------------#
# 此时获得的scaled_anchors大小是相对于特征层的
#-------------------------------------------------#
scaled_anchors = [(a_w / stride_w, a_h / stride_h) for a_w, a_h in self.anchors] #anchors对应的是三种框[116,90],[156,198],[373,326],计算后的框位置(相当于缩小32倍)是相对于13*13的图片里的框
#-----------------------------------------------#
#scaled_anchors是三种框的位置[3.626,2.8125],[4.875,6.1875],[11.656,10.187]
# 输入的input一共有三个,他们的shape分别是
# bs, 3*(5+num_classes), 13, 13 => batch_size, 3, 13, 13, 5 + num_classes
# batch_size, 3, 26, 26, 5 + num_classes
# batch_size, 3, 52, 52, 5 + num_classes
#-----------------------------------------------#
prediction = input.view(bs, len(self.anchors_mask[l]), self.bbox_attrs, in_h, in_w).permute(0, 1, 3, 4, 2).contiguous() #view()的作用相当于numpy中的reshape,改变input的维度 .permute(0, 1, 3, 4, 2)按照维度进行装换 anchors_mask = [[6,7,8], [3,4,5], [0,1,2]]
#len(self.anchors_mask[l])=3 bbox_attrs = 5 + num_classes=25 in_h=13 in_w=13 装换后的维度是(bs,3,13,13,25)
#-----------------------------------------------#
# 先验框的中心位置的调整参数
#在bbox_attrs中,前五个元素代表x,y,w,h,conf,后20个元素代表20个类的概率
#-----------------------------------------------#
x = torch.sigmoid(prediction[..., 0]) #[..., 0]表示全部数据的x ,将x调整至0到1
y = torch.sigmoid(prediction[..., 1])
#-----------------------------------------------#
# 先验框的宽高调整参数
#-----------------------------------------------#
w = prediction[..., 2]
h = prediction[..., 3]
#-----------------------------------------------#
# 获得置信度,是否有物体
#-----------------------------------------------#
conf = torch.sigmoid(prediction[..., 4])
#-----------------------------------------------#
# 种类置信度
#-----------------------------------------------#
pred_cls = torch.sigmoid(prediction[..., 5:])
#-----------------------------------------------#
#获得网络应该有的预测结果
#-----------------------------------------------#
y_true, noobj_mask, box_loss_scale = self.get_target(l, targets, scaled_anchors, in_h, in_w) #scaled_anchors是三种框的位置[3.626,2.8125],[4.875,6.1875],[11.656,10.187] in_h=13 in_w=13
#---------------------------------------------------------------#
# 将预测结果进行解码,判断预测结果和真实值的重合程度
# 如果重合程度过大则忽略,因为这些特征点属于预测比较准确的特征点
# 作为负样本不合适
#----------------------------------------------------------------#
noobj_mask = self.get_ignore(l, x, y, h, w, targets, scaled_anchors, in_h, in_w, noobj_mask) #get_ignore()是用来得到预测的先验框
if self.cuda:
y_true = y_true.cuda()
noobj_mask = noobj_mask.cuda()
box_loss_scale = box_loss_scale.cuda()
#-----------------------------------------------------------#
# reshape_y_true[...,2:3]和reshape_y_true[...,3:4]
# 表示真实框的宽高,二者均在0-1之间
# 真实框越大,比重越小,小框的比重更大。
#-----------------------------------------------------------#
box_loss_scale = 2 - box_loss_scale
#-----------------------------------------------------------#
# 计算中心偏移情况的loss,使用BCELoss效果好一些
#-----------------------------------------------------------#
loss_x = torch.sum(self.BCELoss(x, y_true[..., 0]) * box_loss_scale * y_true[..., 4])
loss_y = torch.sum(self.BCELoss(y, y_true[..., 1]) * box_loss_scale * y_true[..., 4])
#-----------------------------------------------------------#
# 计算宽高调整值的loss
#-----------------------------------------------------------#
loss_w = torch.sum(self.MSELoss(w, y_true[..., 2]) * 0.5 * box_loss_scale * y_true[..., 4])
loss_h = torch.sum(self.MSELoss(h, y_true[..., 3]) * 0.5 * box_loss_scale * y_true[..., 4])
#-----------------------------------------------------------#
# 计算置信度的loss
#-----------------------------------------------------------#
loss_conf = torch.sum(self.BCELoss(conf, y_true[..., 4]) * y_true[..., 4]) + \
torch.sum(self.BCELoss(conf, y_true[..., 4]) * noobj_mask)
loss_cls = torch.sum(self.BCELoss(pred_cls[y_true[..., 4] == 1], y_true[..., 5:][y_true[..., 4] == 1]))
loss = loss_x + loss_y + loss_w + loss_h + loss_conf + loss_cls
num_pos = torch.sum(y_true[..., 4])
num_pos = torch.max(num_pos, torch.ones_like(num_pos))
return loss, num_pos
def calculate_iou(self, _box_a, _box_b): #计算IoU值
#-----------------------------------------------------------#
# 计算真实框的左上角和右下角
#-----------------------------------------------------------#
b1_x1, b1_x2 = _box_a[:, 0] - _box_a[:, 2] / 2, _box_a[:, 0] + _box_a[:, 2] / 2
b1_y1, b1_y2 = _box_a[:, 1] - _box_a[:, 3] / 2, _box_a[:, 1] + _box_a[:, 3] / 2
#-----------------------------------------------------------#
# 计算先验框获得的预测框的左上角和右下角
#-----------------------------------------------------------#
b2_x1, b2_x2 = _box_b[:, 0] - _box_b[:, 2] / 2, _box_b[:, 0] + _box_b[:, 2] / 2
b2_y1, b2_y2 = _box_b[:, 1] - _box_b[:, 3] / 2, _box_b[:, 1] + _box_b[:, 3] / 2
#-----------------------------------------------------------#
# 将真实框和预测框都转化成左上角右下角的形式
#-----------------------------------------------------------#
box_a = torch.zeros_like(_box_a)
box_b = torch.zeros_like(_box_b)
box_a[:, 0], box_a[:, 1], box_a[:, 2], box_a[:, 3] = b1_x1, b1_y1, b1_x2, b1_y2
box_b[:, 0], box_b[:, 1], box_b[:, 2], box_b[:, 3] = b2_x1, b2_y1, b2_x2, b2_y2
#-----------------------------------------------------------#
# A为真实框的数量,B为先验框的数量
#-----------------------------------------------------------#
A = box_a.size(0)
B = box_b.size(0)
#-----------------------------------------------------------#
# 计算交的面积
#-----------------------------------------------------------#
max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2), box_b[:, 2:].unsqueeze(0).expand(A, B, 2)) #右下角最小的x,y
min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2), box_b[:, :2].unsqueeze(0).expand(A, B, 2)) #左上角最大的x,y
inter = torch.clamp((max_xy - min_xy), min=0) #torch.clamp作用是将输入压缩到指定区间 ,这里指定最小为0,即小于0的输入输出为0,大于0的输入不变
inter = inter[:, :, 0] * inter[:, :, 1]
#-----------------------------------------------------------#
# 计算预测框和真实框各自的面积
#-----------------------------------------------------------#
area_a = ((box_a[:, 2]-box_a[:, 0]) * (box_a[:, 3]-box_a[:, 1])).unsqueeze(1).expand_as(inter) # [A,B]
area_b = ((box_b[:, 2]-box_b[:, 0]) * (box_b[:, 3]-box_b[:, 1])).unsqueeze(0).expand_as(inter) # [A,B]
#-----------------------------------------------------------#
# 求IOU
#-----------------------------------------------------------#
union = area_a + area_b - inter
return inter / union # [A,B]
def get_target(self, l, targets, anchors, in_h, in_w): #scaled_anchors是三种框的位置[3.626,2.8125],[4.875,6.1875],[11.656,10.187] in_h=13 in_w=13
#-----------------------------------------------------#
# 计算一共有多少张图片
#-----------------------------------------------------#
bs = len(targets)
#-----------------------------------------------------#
# 用于选取哪些先验框不包含物体
#-----------------------------------------------------#
noobj_mask = torch.ones(bs, len(self.anchors_mask[l]), in_h, in_w, requires_grad = False) #anchors_mask=[[6,7,8], [3,4,5], [0,1,2]] 框的索引 ,len(self.anchors_mask[l])=3
#-----------------------------------------------------#
# 让网络更加去关注小目标
#-----------------------------------------------------#
box_loss_scale = torch.zeros(bs, len(self.anchors_mask[l]), in_h, in_w, requires_grad = False) #(bs,3,13,13)
#-----------------------------------------------------#
# batch_size, 3, 13, 13, 5 + num_classes
#-----------------------------------------------------#
y_true = torch.zeros(bs, len(self.anchors_mask[l]), in_h, in_w, self.bbox_attrs, requires_grad = False)
for b in range(bs):
if len(targets[b])==0:
continue #if中条件满足时,不再执行循环体后面的代码,而是重新开始循环
batch_target = torch.zeros_like(targets[b])
#-------------------------------------------------------#
# 计算出正样本在特征层上的中心点
#-------------------------------------------------------#
batch_target[:, [0,2]] = targets[b][:, [0,2]] * in_w #[:, [0,2]]意思是选中第0列和第2列 targets是真实框
batch_target[:, [1,3]] = targets[b][:, [1,3]] * in_h
batch_target[:, 4] = targets[b][:, 4]
batch_target = batch_target.cpu()
#-------------------------------------------------------#
# 将真实框转换一个形式
# num_true_box, 4
#-------------------------------------------------------#
gt_box = torch.FloatTensor(torch.cat((torch.zeros((batch_target.size(0), 2)), batch_target[:, 2:4]), dim=1)) #batch_target.size(0)=bs 输入的维度是二维,如:(1,2) dim=1表示对列进行拼接,拼接成2+2=4列
#-------------------------------------------------------#
# 将先验框转换一个形式
# 9, 4
#-------------------------------------------------------#
anchor_shapes = torch.FloatTensor(torch.cat((torch.zeros((len(anchors), 2)), torch.FloatTensor(anchors)), 1))
#-------------------------------------------------------#
# 计算交并比
# self.calculate_iou(gt_box, anchor_shapes) = [num_true_box, 9]每一个真实框和9个先验框的重合情况
# best_ns:
# [每个真实框最大的重合度max_iou, 每一个真实框最重合的先验框的序号]
#-------------------------------------------------------#
best_ns = torch.argmax(self.calculate_iou(gt_box, anchor_shapes), dim=-1) #返回指定维度最大值的序号
for t, best_n in enumerate(best_ns):
if best_n not in self.anchors_mask[l]: #anchors_mask=[[6,7,8], [3,4,5], [0,1,2]] 框的索引
continue
#----------------------------------------#
# 判断这个先验框是当前特征点的哪一个先验框
#----------------------------------------#
k = self.anchors_mask[l].index(best_n)
#----------------------------------------#
# 获得真实框属于哪个网格点
#----------------------------------------#
i = torch.floor(batch_target[t, 0]).long() #torch.floor() 向下取整
j = torch.floor(batch_target[t, 1]).long()
#----------------------------------------#
# 取出真实框的种类
#----------------------------------------#
c = batch_target[t, 4].long()
#----------------------------------------#
# noobj_mask代表无目标的特征点
#----------------------------------------#
noobj_mask[b, k, j, i] = 0
#----------------------------------------#
# tx、ty代表中心调整参数的真实值
#----------------------------------------#
y_true[b, k, j, i, 0] = batch_target[t, 0] - i.float()
y_true[b, k, j, i, 1] = batch_target[t, 1] - j.float()
y_true[b, k, j, i, 2] = math.log(batch_target[t, 2] / anchors[best_n][0])
y_true[b, k, j, i, 3] = math.log(batch_target[t, 3] / anchors[best_n][1])
y_true[b, k, j, i, 4] = 1
y_true[b, k, j, i, c + 5] = 1
#----------------------------------------#
# 用于获得xywh的比例
# 大目标loss权重小,小目标loss权重大
#----------------------------------------#
box_loss_scale[b, k, j, i] = batch_target[t, 2] * batch_target[t, 3] / in_w / in_h
return y_true, noobj_mask, box_loss_scale #noobj_mask代表无目标的特征点
def get_ignore(self, l, x, y, h, w, targets, scaled_anchors, in_h, in_w, noobj_mask):
#-----------------------------------------------------#
# 计算一共有多少张图片
#-----------------------------------------------------#
bs = len(targets)
FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor
#-----------------------------------------------------#
# 生成网格,先验框中心,网格左上角
#-----------------------------------------------------#
grid_x = torch.linspace(0, in_w - 1, in_w).repeat(in_h, 1).repeat( #torch.linspace()该函数的参数包括起始值、终止值和张量中元素的数量 repeat的参数是对应维度的复制个数
int(bs * len(self.anchors_mask[l])), 1, 1).view(x.shape).type(FloatTensor)
grid_y = torch.linspace(0, in_h - 1, in_h).repeat(in_w, 1).t().repeat(
int(bs * len(self.anchors_mask[l])), 1, 1).view(y.shape).type(FloatTensor)
# 生成先验框的宽高
scaled_anchors_l = np.array(scaled_anchors)[self.anchors_mask[l]]
anchor_w = FloatTensor(scaled_anchors_l).index_select(1, LongTensor([0])) #index_select() 参数1表示从第几维挑选数据,参数2表示从第一个参数维度中的哪个位置挑选数据
anchor_h = FloatTensor(scaled_anchors_l).index_select(1, LongTensor([1]))
anchor_w = anchor_w.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(w.shape)
anchor_h = anchor_h.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(h.shape)
#-------------------------------------------------------#
# 计算调整后的先验框中心与宽高
#-------------------------------------------------------#
pred_boxes_x = torch.unsqueeze(x.data + grid_x, -1)
pred_boxes_y = torch.unsqueeze(y.data + grid_y, -1)
pred_boxes_w = torch.unsqueeze(torch.exp(w.data) * anchor_w, -1)
pred_boxes_h = torch.unsqueeze(torch.exp(h.data) * anchor_h, -1)
pred_boxes = torch.cat([pred_boxes_x, pred_boxes_y, pred_boxes_w, pred_boxes_h], dim = -1) #生成了预测的先验框
for b in range(bs):
#-------------------------------------------------------#
# 将预测结果转换一个形式
# pred_boxes_for_ignore num_anchors, 4
#-------------------------------------------------------#
pred_boxes_for_ignore = pred_boxes[b].view(-1, 4)
#-------------------------------------------------------#
# 计算真实框,并把真实框转换成相对于特征层的大小
# gt_box num_true_box, 4
#-------------------------------------------------------#
if len(targets[b]) > 0:
batch_target = torch.zeros_like(targets[b])
#-------------------------------------------------------#
# 计算出正样本在特征层上的中心点
#-------------------------------------------------------#
batch_target[:, [0,2]] = targets[b][:, [0,2]] * in_w
batch_target[:, [1,3]] = targets[b][:, [1,3]] * in_h
batch_target = batch_target[:, :4]
#-------------------------------------------------------#
# 计算交并比
# anch_ious num_true_box, num_anchors
#-------------------------------------------------------#
anch_ious = self.calculate_iou(batch_target, pred_boxes_for_ignore)
#-------------------------------------------------------#
# 每个先验框对应真实框的最大重合度
# anch_ious_max num_anchors
#-------------------------------------------------------#
anch_ious_max, _ = torch.max(anch_ious, dim = 0)
anch_ious_max = anch_ious_max.view(pred_boxes[b].size()[:3])
noobj_mask[b][anch_ious_max > self.ignore_threshold] = 0
return noobj_mask
def weights_init(net, init_type='normal', init_gain = 0.02):
def init_func(m):
classname = m.__class__.__name__
if hasattr(m, 'weight') and classname.find('Conv') != -1:
if init_type == 'normal':
torch.nn.init.normal_(m.weight.data, 0.0, init_gain)
elif init_type == 'xavier':
torch.nn.init.xavier_normal_(m.weight.data, gain=init_gain)
elif init_type == 'kaiming':
torch.nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
elif init_type == 'orthogonal':
torch.nn.init.orthogonal_(m.weight.data, gain=init_gain)
else:
raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
elif classname.find('BatchNorm2d') != -1:
torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
torch.nn.init.constant_(m.bias.data, 0.0)
print('initialize network with %s type' % init_type)
net.apply(init_func)
YoloBady输出的三个out的先验框是基于三种尺寸的特征图上的先验框,需对其进行解码,将先验框显现在原图上。
1.首先得到基于特征图的锚框scaled_anchors
2.计算先验框的中心位置的调整参数x、y和宽高调整参数w、h
3.通过get_target方法计算真实框与先验框的交并比
预测部分:
主函数部分调用YOLO()类,传入图片至detect_image。
1.首先对图片进行预处理,对非正方形的图片增加灰条(防止调整图片时发生失真),将数据 类型转为tensor类型。
2.将经过预处理后的图片传入YoloBody中,返回三个经过深层网络的特征图输出。
3.对三个特征图输出进行解码操作,返回一个list,里面有三个tensor,对应三种特征图上每个特征点的85个关于先验框数据(4个先验框的位置尺寸数据+1个物体置信度+80个种类置信度)的输出。
4.进行非极大值抑制,输出最终预测框的数据(每一个预测框有七个数据,框的左上角和右下角坐标、物体置信度、种类置信度、种类索引(x1, y1, x2, y2, obj_conf, class_conf, class_pred))。
5。最后在图片上绘制预测框。
解码代码:
import torch
import torch.nn as nn
from torchvision.ops import nms
import numpy as np
class DecodeBox():
def __init__(self, anchors, num_classes, input_shape, anchors_mask = [[6,7,8], [3,4,5], [0,1,2]]): #锚框原文件尺寸:10,13, 16,30, 33,23, 30,61, 62,45, 59,119, 116,90, 156,198, 373,326
super(DecodeBox, self).__init__()
self.anchors = anchors
self.num_classes = num_classes
self.bbox_attrs = 5 + num_classes
self.input_shape = input_shape
#-----------------------------------------------------------#
# 13x13的特征层对应的anchor是[116,90],[156,198],[373,326]
# 26x26的特征层对应的anchor是[30,61],[62,45],[59,119]
# 52x52的特征层对应的anchor是[10,13],[16,30],[33,23]
#-----------------------------------------------------------#
self.anchors_mask = anchors_mask
def decode_box(self, inputs):
outputs = []
for i, input in enumerate(inputs):
#-----------------------------------------------#
# 输入的input一共有三个,他们的shape分别是
# batch_size, 255, 13, 13
# batch_size, 255, 26, 26
# batch_size, 255, 52, 52
#-----------------------------------------------#
batch_size = input.size(0)
input_height = input.size(2)
input_width = input.size(3)
#-----------------------------------------------#
# 输入为416x416时
# stride_h = stride_w = 32、16、8
#-----------------------------------------------#
stride_h = self.input_shape[0] / input_height
stride_w = self.input_shape[1] / input_width
#-------------------------------------------------#
# 此时获得的scaled_anchors大小是相对于特征层的
#-------------------------------------------------#
scaled_anchors = [(anchor_width / stride_w, anchor_height / stride_h) for anchor_width, anchor_height in self.anchors[self.anchors_mask[i]]]
#-----------------------------------------------#
# 输入的input一共有三个,他们的shape分别是
# batch_size, 3, 13, 13, 85
# batch_size, 3, 26, 26, 85
# batch_size, 3, 52, 52, 85
#-----------------------------------------------#
prediction = input.view(batch_size, len(self.anchors_mask[i]), #anchors_mask[0] = [6,7,8] 对应锚框尺寸为:116,90, 156,198, 373,326
self.bbox_attrs, input_height, input_width).permute(0, 1, 3, 4, 2).contiguous() #改变维度顺序,改为(bs,3,13,13,85)
#-----------------------------------------------#
# 先验框的中心位置的调整参数
#-----------------------------------------------#
x = torch.sigmoid(prediction[..., 0]) #85列中的第0列 维度为:(bs,3,13,13,1)或(bs,3,13,13)
y = torch.sigmoid(prediction[..., 1])
#-----------------------------------------------#
# 先验框的宽高调整参数
#-----------------------------------------------#
w = prediction[..., 2]
h = prediction[..., 3]
#-----------------------------------------------#
# 获得置信度,是否有物体
#-----------------------------------------------#
conf = torch.sigmoid(prediction[..., 4])
#-----------------------------------------------#
# 种类置信度
#-----------------------------------------------#
pred_cls = torch.sigmoid(prediction[..., 5:]) #维度为:(bs,3,13,13,80)
FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor
#----------------------------------------------------------#
# 生成网格,先验框中心,网格左上角
# batch_size,3,13,13
#----------------------------------------------------------#
grid_x = torch.linspace(0, input_width - 1, input_width).repeat(input_height, 1).repeat(
batch_size * len(self.anchors_mask[i]), 1, 1).view(x.shape).type(FloatTensor) #维度为:(bs,3,13,13)
grid_y = torch.linspace(0, input_height - 1, input_height).repeat(input_width, 1).t().repeat(
batch_size * len(self.anchors_mask[i]), 1, 1).view(y.shape).type(FloatTensor)
#----------------------------------------------------------#
# 按照网格格式生成先验框的宽高
# batch_size,3,13,13
#----------------------------------------------------------#
anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0])) #维度为:(3,1)
anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))
anchor_w = anchor_w.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(w.shape) #维度为:(bs,3,13,13)
anchor_h = anchor_h.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(h.shape)
#----------------------------------------------------------#
# 利用预测结果对先验框进行调整
# 首先调整先验框的中心,从先验框中心向右下角偏移
# 再调整先验框的宽高。
#----------------------------------------------------------#
pred_boxes = FloatTensor(prediction[..., :4].shape)
pred_boxes[..., 0] = x.data + grid_x
pred_boxes[..., 1] = y.data + grid_y
pred_boxes[..., 2] = torch.exp(w.data) * anchor_w
pred_boxes[..., 3] = torch.exp(h.data) * anchor_h
#----------------------------------------------------------#
# 将输出结果归一化成小数的形式
#----------------------------------------------------------#
_scale = torch.Tensor([input_width, input_height, input_width, input_height]).type(FloatTensor)
output = torch.cat((pred_boxes.view(batch_size, -1, 4) / _scale,
conf.view(batch_size, -1, 1), pred_cls.view(batch_size, -1, self.num_classes)), -1) #维度为:(bs,507,85) 3*13*13=507
outputs.append(output.data)
return outputs
def yolo_correct_boxes(self, box_xy, box_wh, input_shape, image_shape, letterbox_image):
#-----------------------------------------------------------------#
# 把y轴放前面是因为方便预测框和图像的宽高进行相乘
#-----------------------------------------------------------------#
box_yx = box_xy[..., ::-1] #把x和y的位置反过来
box_hw = box_wh[..., ::-1]
input_shape = np.array(input_shape)
image_shape = np.array(image_shape)
if letterbox_image: #这张图片不进入if
#-----------------------------------------------------------------#
# 这里求出来的offset是图像有效区域相对于图像左上角的偏移情况
# new_shape指的是宽高缩放情况
#-----------------------------------------------------------------#
new_shape = np.round(image_shape * np.min(input_shape/image_shape))
offset = (input_shape - new_shape)/2./input_shape
scale = input_shape/new_shape
box_yx = (box_yx - offset) * scale
box_hw *= scale
box_mins = box_yx - (box_hw / 2.) #框的左上角位置y1,x1
box_maxes = box_yx + (box_hw / 2.) #框的右下角位置y2,x2
boxes = np.concatenate([box_mins[..., 0:1], box_mins[..., 1:2], box_maxes[..., 0:1], box_maxes[..., 1:2]], axis=-1) #类似torch.cat的拼接操作
boxes *= np.concatenate([image_shape, image_shape], axis=-1) #先是按列拼接成四列,再执行乘法操作
return boxes #先验框左上角和右下角坐标,这张图片是12个框,所以维度是(12,4)
def non_max_suppression(self, prediction, num_classes, input_shape, image_shape, letterbox_image, conf_thres=0.5, nms_thres=0.4):
#----------------------------------------------------------#
# 将预测结果的格式转换成左上角右下角的格式。
# prediction [batch_size, num_anchors, 85]
#----------------------------------------------------------#
box_corner = prediction.new(prediction.shape) #传入的prediction是out中三个list按行拼接后的结果
box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2 #左上角x
box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2 #左上角y
box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2 #右下角x
box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2 #右下角y
prediction[:, :, :4] = box_corner[:, :, :4]
output = [None for _ in range(len(prediction))]
for i, image_pred in enumerate(prediction): #prediction维度是(bs,10647,85) 10647 = 3*13*13 + 3*26*26 + 3*52*52
#----------------------------------------------------------#
# 对种类预测部分取max。
# class_conf [num_anchors, 1] 种类置信度
# class_pred [num_anchors, 1] 种类
#----------------------------------------------------------#
class_conf, class_pred = torch.max(image_pred[:, 5:5 + num_classes], 1, keepdim=True) #第5列后才是80个种类的概率,输出一行中概率最大的种类 class_conf和class_pred维度是(10647,1)
#----------------------------------------------------------#
# 利用置信度进行第一轮筛选
#----------------------------------------------------------#
conf_mask = (image_pred[:, 4] * class_conf[:, 0] >= conf_thres).squeeze() #image_pred[:, 4]是判断是否有物体(结果为0或1) class_conf[:, 0]就是种类置信度class_conf(一个0至1的概率值) 相乘后结果大于等于阈值为实际值,否则为False,结果是一个10647的tensor
#----------------------------------------------------------#
# 根据置信度进行预测结果的筛选
#----------------------------------------------------------#
image_pred = image_pred[conf_mask] #一张图片最终的维度是(33,85),这说明这张图片在三个特征图中共有33个网格有物体(三个特征图共10647个网格)
class_conf = class_conf[conf_mask] #维度为:(33,1) 把有物体对应的种类置信度放置在一起
class_pred = class_pred[conf_mask] #维度为:(33,1) 种类名称(这里是索引)
if not image_pred.size(0):
continue
#-------------------------------------------------------------------------#
# detections [num_anchors, 7]
# 7的内容为:x1, y1, x2, y2, obj_conf, class_conf, class_pred
#-------------------------------------------------------------------------#
detections = torch.cat((image_pred[:, :5], class_conf.float(), class_pred.float()), 1) #维度为:(33,7)
#------------------------------------------#
# 获得预测结果中包含的所有种类
#------------------------------------------#
unique_labels = detections[:, -1].cpu().unique() #输出是一个tensor,是33个先验框的种类个数,这张图是3
if prediction.is_cuda:
unique_labels = unique_labels.cuda()
detections = detections.cuda()
for c in unique_labels:
#------------------------------------------#
# 获得某一类得分筛选后全部的预测结果
#------------------------------------------#
detections_class = detections[detections[:, -1] == c] #把c这个索引值对应的种类放在一起,比如在这个图片中33个框有22个框是0索引对应的这个种类,维度为:(22,7)
#------------------------------------------#
# 使用官方自带的非极大抑制会速度更快一些!
#------------------------------------------#
keep = nms( #keep返回的是一个tensor,非极大值抑制后的结果,这里是8个数据(对应行数)
detections_class[:, :4], #前0,1,2,3列
detections_class[:, 4] * detections_class[:, 5],
nms_thres
)
max_detections = detections_class[keep] #维度为:(8,7) 八个经过极大值抑制后的先验框的数据
# # 按照存在物体的置信度排序
# _, conf_sort_index = torch.sort(detections_class[:, 4]*detections_class[:, 5], descending=True)
# detections_class = detections_class[conf_sort_index]
# # 进行非极大抑制
# max_detections = []
# while detections_class.size(0):
# # 取出这一类置信度最高的,一步一步往下判断,判断重合程度是否大于nms_thres,如果是则去除掉
# max_detections.append(detections_class[0].unsqueeze(0))
# if len(detections_class) == 1:
# break
# ious = bbox_iou(max_detections[-1], detections_class[1:])
# detections_class = detections_class[1:][ious < nms_thres]
# # 堆叠
# max_detections = torch.cat(max_detections).data
# Add max detections to outputs
output[i] = max_detections if output[i] is None else torch.cat((output[i], max_detections)) #把三种物体经过非极大值抑制后的框放在一起,这里维度是:(12,7)
if output[i] is not None:
output[i] = output[i].cpu().numpy()
box_xy, box_wh = (output[i][:, 0:2] + output[i][:, 2:4])/2, output[i][:, 2:4] - output[i][:, 0:2] #维度都为:(12,2),输出先验框中心点位置和框的宽高 7的内容为:x1, y1, x2, y2, obj_conf, class_conf, class_pred
output[i][:, :4] = self.yolo_correct_boxes(box_xy, box_wh, input_shape, image_shape, letterbox_image)
return output