睿智的目标检测46——Pytorch搭建自己的Centernet目标检测平台
学习前言
Pytorch版本的实现也要做一下。
什么是Centernet目标检测算法
如今常见的目标检测算法通常使用先验框的设定，即先在图片上设定大量的先验框，网络的预测结果会对先验框进行调整获得预测框，先验框很大程度上提高了网络的检测能力，但是也会收到物体尺寸的限制。
Centernet采用不同的方法，构建模型时将目标作为一个点——即目标BBox的中心点。
Centernet的检测器采用关键点估计来找到中心点，并回归到其他目标属性。
论文中提到：模型是端到端可微的，更简单，更快，更精确。Centernet的模型实现了速度和精确的很好权衡。
源码下载
https://github.com/bubbliiiing/centernet-pytorch
喜欢的可以点个star噢。
Centernet实现思路
一、预测部分
1、主干网络介绍
Centernet用到的主干特征网络有多种，一般是以Hourglass Network、DLANet或者Resnet为主干特征提取网络，由于centernet所用到的Hourglass Network参数量太大，有19000W参数，DLANet并没有keras资源，本文以Resnet为例子进行解析。
ResNet50有两个基本的块，分别名为Conv Block和Identity Block，其中Conv Block输入和输出的维度是不一样的，所以不能连续串联，它的作用是改变网络的维度；Identity Block输入维度和输出维度相同，可以串联，用于加深网络的。
Conv Block的结构如下：
Identity Block的结构如下：
这两个都是残差网络结构。
当我们输入的图片是512x512x3的时候，整体的特征层shape变化为：
我们取出最终一个block的输出进行下一步的处理。也就是图上的C5，它的shape为16x16x2048。利用主干特征提取网络，我们获取到了一个初步的特征层，其shape为16x16x2048。
from __future__ import absolute_import, division, print_function
import math
import torch.nn as nn
from torchvision.models.utils import load_state_dict_from_url
model_urls = {
'resnet18': 'https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth',
'resnet34': 'https://s3.amazonaws.com/pytorch/models/resnet34-333f7ec4.pth',
'resnet50': 'https://s3.amazonaws.com/pytorch/models/resnet50-19c8e357.pth',
'resnet101': 'https://s3.amazonaws.com/pytorch/models/resnet101-5d3b4d8f.pth',
'resnet152': 'https://s3.amazonaws.com/pytorch/models/resnet152-b121ed2d.pth',
}
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(Bottleneck, self).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride, bias=False) # change
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, # change
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):
residual = 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:
residual = self.downsample(x)
out += residual
out = self.relu(out)
return out
#-----------------------------------------------------------------#
# 使用Renset50作为主干特征提取网络，最终会获得一个
# 16x16x2048的有效特征层
#-----------------------------------------------------------------#
class ResNet(nn.Module):
def __init__(self, block, layers, num_classes=1000):
self.inplanes = 64
super(ResNet, self).__init__()
# 512,512,3 -> 256,256,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)
# 256x256x64 -> 128x128x64
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=0, ceil_mode=True) # change
# 128x128x64 -> 128x128x256
self.layer1 = self._make_layer(block, 64, layers[0])
# 128x128x256 -> 64x64x512
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
# 64x64x512 -> 32x32x1024
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
# 32x32x1024 -> 16x16x2048
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
self.avgpool = nn.AvgPool2d(7)
self.fc = nn.Linear(512 * block.expansion, num_classes)
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_()
def _make_layer(self, block, planes, blocks, stride=1):
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 = []
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 forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
def resnet50(pretrained = True):
model = ResNet(Bottleneck, [3, 4, 6, 3])
if pretrained:
state_dict = load_state_dict_from_url(model_urls['resnet50'], model_dir = 'model_data/')
model.load_state_dict(state_dict)
#----------------------------------------------------------#
# 获取特征提取部分
#----------------------------------------------------------#
features = list([model.conv1, model.bn1, model.relu, model.maxpool, model.layer1, model.layer2, model.layer3, model.layer4])
features = nn.Sequential(*features)
return features
2、利用初步特征获得高分辨率特征图
利用上一步获得到的resnet50的最后一个特征层的shape为(16,16,2048)。
对于该特征层，centernet利用三次反卷积进行上采样，从而更高的分辨率输出。为了节省计算量，这3个反卷积的输出通道数分别为256,128,64。
每一次反卷积，特征层的高和宽会变为原来的两倍，因此，在进行三次反卷积上采样后，我们获得的特征层的高和宽变为原来的8倍，此时特征层的高和宽为128x128，通道数为64。
此时我们获得了一个128x128x64的有效特征层（高分辨率特征图），我们会利用该有效特征层获得最终的预测结果。
实现代码如下：
class resnet50_Decoder(nn.Module):
def __init__(self, inplanes, bn_momentum=0.1):
super(resnet50_Decoder, self).__init__()
self.bn_momentum = bn_momentum
self.inplanes = inplanes
self.deconv_with_bias = False
#----------------------------------------------------------#
# 16,16,2048 -> 32,32,256 -> 64,64,128 -> 128,128,64
# 利用ConvTranspose2d进行上采样。
# 每次特征层的宽高变为原来的两倍。
#----------------------------------------------------------#
self.deconv_layers = self._make_deconv_layer(
num_layers=3,
num_filters=[256, 128, 64],
num_kernels=[4, 4, 4],
)
def _make_deconv_layer(self, num_layers, num_filters, num_kernels):
layers = []
for i in range(num_layers):
kernel = num_kernels[i]
planes = num_filters[i]
layers.append(
nn.ConvTranspose2d(
in_channels=self.inplanes,
out_channels=planes,
kernel_size=kernel,
stride=2,
padding=1,
output_padding=0,
bias=self.deconv_with_bias))
layers.append(nn.BatchNorm2d(planes, momentum=self.bn_momentum))
layers.append(nn.ReLU(inplace=True))
self.inplanes = planes
return nn.Sequential(*layers)
def forward(self, x):
return self.deconv_layers(x)
3、Center Head从特征获取预测结果
通过上一步我们可以获得一个128x128x64的高分辨率特征图。
这个特征层相当于将整个图片划分成128x128个区域，每个区域存在一个特征点，如果某个物体的中心落在这个区域，那么就由这个特征点来确定。
（某个物体的中心落在这个区域，则由这个区域左上角的特征点来约定）
我们可以利用这个特征层进行三个卷积，分别是：
1、热力图预测，此时卷积的通道数为num_classes，最终结果为(128,128,num_classes)，代表每一个热力点是否有物体存在，以及物体的种类；
2、中心点预测，此时卷积的通道数为2，最终结果为(128,128,2)，代表每一个物体中心距离热力点偏移的情况；
3、宽高预测，此时卷积的通道数为2，最终结果为(128,128,2)，代表每一个物体宽高的预测情况；
实现代码为：
class resnet50_Head(nn.Module):
def __init__(self, num_classes=80, channel=64, bn_momentum=0.1):
super(resnet50_Head, self).__init__()
#-----------------------------------------------------------------#
# 对获取到的特征进行上采样，进行分类预测和回归预测
# 128, 128, 64 -> 128, 128, 64 -> 128, 128, num_classes
# -> 128, 128, 64 -> 128, 128, 2
# -> 128, 128, 64 -> 128, 128, 2
#-----------------------------------------------------------------#
# 热力图预测部分
self.cls_head = nn.Sequential(
nn.Conv2d(64, channel,
kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(64, momentum=bn_momentum),
nn.ReLU(inplace=True),
nn.Conv2d(channel, num_classes,
kernel_size=1, stride=1, padding=0))
# 宽高预测的部分
self.wh_head = nn.Sequential(
nn.Conv2d(64, channel,
kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(64, momentum=bn_momentum),
nn.ReLU(inplace=True),
nn.Conv2d(channel, 2,
kernel_size=1, stride=1, padding=0))
# 中心点预测的部分
self.reg_head = nn.Sequential(
nn.Conv2d(64, channel,
kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(64, momentum=bn_momentum),
nn.ReLU(inplace=True),
nn.Conv2d(channel, 2,
kernel_size=1, stride=1, padding=0))
def forward(self, x):
hm = self.cls_head(x).sigmoid_()
wh = self.wh_head(x)
offset = self.reg_head(x)
return hm, wh, offset
4、预测结果的解码
在对预测结果进行解码之前，我们再来看看预测结果代表了什么，预测结果可以分为3个部分：
1、heatmap热力图预测，此时卷积的通道数为num_classes，最终结果为(128,128,num_classes)，代表每一个热力点是否有物体存在，以及物体的种类，最后一维度num_classes中的预测值代表属于每一个类的概率；
2、reg中心点预测，此时卷积的通道数为2，最终结果为(128,128,2)，代表每一个物体中心距离热力点偏移的情况，最后一维度2中的预测值代表当前这个特征点向右下角偏移的情况；
3、wh宽高预测，此时卷积的通道数为2，最终结果为(128,128,2)，代表每一个物体宽高的预测情况，最后一维度2中的预测值代表当前这个特征点对应的预测框的宽高；
特征层相当于将图像划分成128x128个特征点每个特征点负责预测中心落在其右下角一片区域的物体。
如图所示，蓝色的点为128x128的特征点，此时我们对左图红色的三个点进行解码操作演示：
1、进行中心点偏移，利用reg中心点预测对特征点坐标进行偏移，左图红色的三个特征点偏移后是右图橙色的三个点；
2、利用中心点加上和减去 wh宽高预测结果除以2，获得预测框的左上角和右下角。
3、此时获得的预测框就可以绘制在图片上了。
除去这样的解码操作，还有非极大抑制的操作需要进行，防止同一种类的框的堆积。
在论文中所说，centernet不像其它目标检测算法，在解码之后需要进行非极大抑制，centernet的非极大抑制在解码之前进行。采用的方法是最大池化，利用3x3的池化核在热力图上搜索，然后只保留一定区域内得分最大的框。
在实际使用时发现，当目标为小目标时，确实可以不在解码之后进行非极大抑制的后处理，如果目标为大目标，网络无法正确判断目标的中心时，还是需要进行非击大抑制的后处理的。
def pool_nms(heat, kernel = 3):
pad = (kernel - 1) // 2
hmax = nn.functional.max_pool2d(heat, (kernel, kernel), stride=1, padding=pad)
keep = (hmax == heat).float()
return heat * keep
def decode_bbox(pred_hms, pred_whs, pred_offsets, confidence, cuda):
#-------------------------------------------------------------------------#
# 当利用512x512x3图片进行coco数据集预测的时候
# h = w = 128 num_classes = 80
# Hot map热力图 -> b, 80, 128, 128,
# 进行热力图的非极大抑制，利用3x3的卷积对热力图进行最大值筛选
# 找出一定区域内，得分最大的特征点。
#-------------------------------------------------------------------------#
pred_hms = pool_nms(pred_hms)
b, c, output_h, output_w = pred_hms.shape
detects = []
#-------------------------------------------------------------------------#
# 只传入一张图片，循环只进行一次
#-------------------------------------------------------------------------#
for batch in range(b):
#-------------------------------------------------------------------------#
# heat_map 128*128, num_classes 热力图
# pred_wh 128*128, 2 特征点的预测宽高
# pred_offset 128*128, 2 特征点的xy轴偏移情况
#-------------------------------------------------------------------------#
heat_map = pred_hms[batch].permute(1, 2, 0).view([-1, c])
pred_wh = pred_whs[batch].permute(1, 2, 0).view([-1, 2])
pred_offset = pred_offsets[batch].permute(1, 2, 0).view([-1, 2])
yv, xv = torch.meshgrid(torch.arange(0, output_h), torch.arange(0, output_w))
#-------------------------------------------------------------------------#
# xv 128*128, 特征点的x轴坐标
# yv 128*128, 特征点的y轴坐标
#-------------------------------------------------------------------------#
xv, yv = xv.flatten().float(), yv.flatten().float()
if cuda:
xv = xv.cuda()
yv = yv.cuda()
#-------------------------------------------------------------------------#
# class_conf 128*128, 特征点的种类置信度
# class_pred 128*128, 特征点的种类
#-------------------------------------------------------------------------#
class_conf, class_pred = torch.max(heat_map, dim = -1)
mask = class_conf > confidence
#-----------------------------------------#
# 取出得分筛选后对应的结果
#-----------------------------------------#
pred_wh_mask = pred_wh[mask]
pred_offset_mask = pred_offset[mask]
if len(pred_wh_mask) == 0:
detects.append([])
continue
#----------------------------------------#
# 计算调整后预测框的中心
#----------------------------------------#
xv_mask = torch.unsqueeze(xv[mask] + pred_offset_mask[..., 0], -1)
yv_mask = torch.unsqueeze(yv[mask] + pred_offset_mask[..., 1], -1)
#----------------------------------------#
# 计算预测框的宽高
#----------------------------------------#
half_w, half_h = pred_wh_mask[..., 0:1] / 2, pred_wh_mask[..., 1:2] / 2
#----------------------------------------#
# 获得预测框的左上角和右下角
#----------------------------------------#
bboxes = torch.cat([xv_mask - half_w, yv_mask - half_h, xv_mask + half_w, yv_mask + half_h], dim=1)
bboxes[:, [0, 2]] /= output_w
bboxes[:, [1, 3]] /= output_h
detect = torch.cat([bboxes, torch.unsqueeze(class_conf[mask],-1), torch.unsqueeze(class_pred[mask],-1).float()], dim=-1)
detects.append(detect)
return detects
5、在原图上进行绘制
通过第三步，我们可以获得预测框在原图上的位置，而且这些预测框都是经过筛选的。这些筛选后的框可以直接绘制在图片上，就可以获得结果了。
二、训练部分
1、真实框的处理
既然在centernet中，物体的中心落在哪个特征点的右下角就由哪个特征点来负责预测，那么在训练的时候我们就需要找到真实框和特征点之间的关系。
真实框和特征点之间的关系，对应方式如下：
1、找到真实框的中心，通过真实框的中心找到其对应的特征点。
2、根据该真实框的种类，对网络应该有的热力图进行设置，即heatmap热力图。其实就是将对应的特征点里面的对应的种类，它的中心值设置为1，然后这个特征点附近的其它特征点中该种类对应的值按照高斯分布不断下降。
3、除去heatmap热力图外，还需要设置特征点对应的reg中心点和wh宽高，在找到真实框对应的特征点后，还需要设置该特征点对应的reg中心和wh宽高。这里的reg中心和wh宽高都是对于128x128的特征层的。
4、在获得网络应该有的预测结果后，就可以将预测结果和应该有的预测结果进行对比，对网络进行反向梯度调整了。
实现代码如下：
import math
import cv2
import numpy as np
from PIL import Image
from torch.utils.data.dataset import Dataset
from utils.utils import cvtColor, preprocess_input
def draw_gaussian(heatmap, center, radius, k=1):
diameter = 2 * radius + 1
gaussian = gaussian2D((diameter, diameter), sigma=diameter / 6)
x, y = int(center[0]), int(center[1])
height, width = heatmap.shape[0:2]
left, right = min(x, radius), min(width - x, radius + 1)
top, bottom = min(y, radius), min(height - y, radius + 1)
masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right]
masked_gaussian = gaussian[radius - top:radius + bottom, radius - left:radius + right]
if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0: # TODO debug
np.maximum(masked_heatmap, masked_gaussian * k, out=masked_heatmap)
return heatmap
def gaussian2D(shape, sigma=1):
m, n = [(ss - 1.) / 2. for ss in shape]
y, x = np.ogrid[-m:m + 1, -n:n + 1]
h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
h[h < np.finfo(h.dtype).eps * h.max()] = 0
return h
def gaussian_radius(det_size, min_overlap=0.7):
height, width = det_size
a1 = 1
b1 = (height + width)
c1 = width * height * (1 - min_overlap) / (1 + min_overlap)
sq1 = np.sqrt(b1 ** 2 - 4 * a1 * c1)
r1 = (b1 + sq1) / 2
a2 = 4
b2 = 2 * (height + width)
c2 = (1 - min_overlap) * width * height
sq2 = np.sqrt(b2 ** 2 - 4 * a2 * c2)
r2 = (b2 + sq2) / 2
a3 = 4 * min_overlap
b3 = -2 * min_overlap * (height + width)
c3 = (min_overlap - 1) * width * height
sq3 = np.sqrt(b3 ** 2 - 4 * a3 * c3)
r3 = (b3 + sq3) / 2
return min(r1, r2, r3)
class CenternetDataset(Dataset):
def __init__(self, annotation_lines, input_shape, num_classes, train):
super(CenternetDataset, self).__init__()
self.annotation_lines = annotation_lines
self.length = len(self.annotation_lines)
self.input_shape = input_shape
self.output_shape = (int(input_shape[0]/4) , int(input_shape[1]/4))
self.num_classes = num_classes
self.train = train
def __len__(self):
return self.length
def __getitem__(self, index):
index = index % self.length
#-------------------------------------------------#
# 进行数据增强
#-------------------------------------------------#
image, box = self.get_random_data(self.annotation_lines[index], self.input_shape, random = self.train)
batch_hm = np.zeros((self.output_shape[0], self.output_shape[1], self.num_classes), dtype=np.float32)
batch_wh = np.zeros((self.output_shape[0], self.output_shape[1], 2), dtype=np.float32)
batch_reg = np.zeros((self.output_shape[0], self.output_shape[1], 2), dtype=np.float32)
batch_reg_mask = np.zeros((self.output_shape[0], self.output_shape[1]), dtype=np.float32)
if len(box) != 0:
boxes = np.array(box[:, :4],dtype=np.float32)
boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]] / self.input_shape[1] * self.output_shape[1], 0, self.output_shape[1] - 1)
boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]] / self.input_shape[0] * self.output_shape[0], 0, self.output_shape[0] - 1)
for i in range(len(box)):
bbox = boxes[i].copy()
cls_id = int(box[i, -1])
h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]
if h > 0 and w > 0:
radius = gaussian_radius((math.ceil(h), math.ceil(w)))
radius = max(0, int(radius))
#-------------------------------------------------#
# 计算真实框所属的特征点
#-------------------------------------------------#
ct = np.array([(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2], dtype=np.float32)
ct_int = ct.astype(np.int32)
#----------------------------#
# 绘制高斯热力图
#----------------------------#
batch_hm[:, :, cls_id] = draw_gaussian(batch_hm[:, :, cls_id], ct_int, radius)
#---------------------------------------------------#
# 计算宽高真实值
#---------------------------------------------------#
batch_wh[ct_int[1], ct_int[0]] = 1. * w, 1. * h
#---------------------------------------------------#
# 计算中心偏移量
#---------------------------------------------------#
batch_reg[ct_int[1], ct_int[0]] = ct - ct_int
#---------------------------------------------------#
# 将对应的mask设置为1
#---------------------------------------------------#
batch_reg_mask[ct_int[1], ct_int[0]] = 1
image = np.transpose(preprocess_input(image), (2, 0, 1))
return image, batch_hm, batch_wh, batch_reg, batch_reg_mask
def rand(self, a=0, b=1):
return np.random.rand()*(b-a) + a
def get_random_data(self, annotation_line, input_shape, jitter=.3, hue=.1, sat=1.5, val=1.5, random=True):
line = annotation_line.split()
#------------------------------#
# 读取图像并转换成RGB图像
#------------------------------#
image = Image.open(line[0])
image = cvtColor(image)
#------------------------------#
# 获得图像的高宽与目标高宽
#------------------------------#
iw, ih = image.size
h, w = input_shape
#------------------------------#
# 获得预测框
#------------------------------#
box = np.array([np.array(list(map(int,box.split(',')))) for box in line[1:]])
if not random:
scale = min(w/iw, h/ih)
nw = int(iw*scale)
nh = int(ih*scale)
dx = (w-nw)//2
dy = (h-nh)//2
#---------------------------------#
# 将图像多余的部分加上灰条
#---------------------------------#
image = image.resize((nw,nh), Image.BICUBIC)
new_image = Image.new('RGB', (w,h), (128,128,128))
new_image.paste(image, (dx, dy))
image_data = np.array(new_image, np.float32)
#---------------------------------#
# 对真实框进行调整
#---------------------------------#
if len(box)>0:
np.random.shuffle(box)
box[:, [0,2]] = box[:, [0,2]]*nw/iw + dx
box[:, [1,3]] = box[:, [1,3]]*nh/ih + dy
box[:, 0:2][box[:, 0:2]<0] = 0
box[:, 2][box[:, 2]>w] = w
box[:, 3][box[:, 3]>h] = h
box_w = box[:, 2] - box[:, 0]
box_h = box[:, 3] - box[:, 1]
box = box[np.logical_and(box_w>1, box_h>1)] # discard invalid box
return image_data, box
#------------------------------------------#
# 对图像进行缩放并且进行长和宽的扭曲
#------------------------------------------#
new_ar = w/h * self.rand(1-jitter,1+jitter) / self.rand(1-jitter,1+jitter)
scale = self.rand(.25, 2)
if new_ar < 1:
nh = int(scale*h)
nw = int(nh*new_ar)
else:
nw = int(scale*w)
nh = int(nw/new_ar)
image = image.resize((nw,nh), Image.BICUBIC)
#------------------------------------------#
# 将图像多余的部分加上灰条
#------------------------------------------#
dx = int(self.rand(0, w-nw))
dy = int(self.rand(0, h-nh))
new_image = Image.new('RGB', (w,h), (128,128,128))
new_image.paste(image, (dx, dy))
image = new_image
#------------------------------------------#
# 翻转图像
#------------------------------------------#
flip = self.rand()<.5
if flip: image = image.transpose(Image.FLIP_LEFT_RIGHT)
#------------------------------------------#
# 色域扭曲
#------------------------------------------#
hue = self.rand(-hue, hue)
sat = self.rand(1, sat) if self.rand()<.5 else 1/self.rand(1, sat)
val = self.rand(1, val) if self.rand()<.5 else 1/self.rand(1, val)
x = cv2.cvtColor(np.array(image,np.float32)/255, cv2.COLOR_RGB2HSV)
x[..., 0] += hue*360
x[..., 0][x[..., 0]>1] -= 1
x[..., 0][x[..., 0]<0] += 1
x[..., 1] *= sat
x[..., 2] *= val
x[x[:,:, 0]>360, 0] = 360
x[:, :, 1:][x[:, :, 1:]>1] = 1
x[x<0] = 0
image_data = cv2.cvtColor(x, cv2.COLOR_HSV2RGB)*255
#---------------------------------#
# 对真实框进行调整
#---------------------------------#
if len(box)>0:
np.random.shuffle(box)
box[:, [0,2]] = box[:, [0,2]]*nw/iw + dx
box[:, [1,3]] = box[:, [1,3]]*nh/ih + dy
if flip: box[:, [0,2]] = w - box[:, [2,0]]
box[:, 0:2][box[:, 0:2]<0] = 0
box[:, 2][box[:, 2]>w] = w
box[:, 3][box[:, 3]>h] = h
box_w = box[:, 2] - box[:, 0]
box_h = box[:, 3] - box[:, 1]
box = box[np.logical_and(box_w>1, box_h>1)]
return image_data, box
# DataLoader中collate_fn使用
def centernet_dataset_collate(batch):
imgs, batch_hms, batch_whs, batch_regs, batch_reg_masks = [], [], [], [], []
for img, batch_hm, batch_wh, batch_reg, batch_reg_mask in batch:
imgs.append(img)
batch_hms.append(batch_hm)
batch_whs.append(batch_wh)
batch_regs.append(batch_reg)
batch_reg_masks.append(batch_reg_mask)
imgs = np.array(imgs)
batch_hms = np.array(batch_hms)
batch_whs = np.array(batch_whs)
batch_regs = np.array(batch_regs)
batch_reg_masks = np.array(batch_reg_masks)
return imgs, batch_hms, batch_whs, batch_regs, batch_reg_masks
2、利用处理完的真实框与对应图片的预测结果计算loss
loss计算分为三个部分，分别是：
1、热力图的loss
2、reg中心点的loss
3、wh宽高的loss
具体情况如下：
1、热力图的loss采用focal loss的思想进行运算，其中 α 和 β 是Focal Loss的超参数，而其中的N是正样本的数量，用于进行标准化。 α 和 β在这篇论文中和分别是2和4。
整体思想和Focal Loss类似，对于容易分类的样本，适当减少其训练比重也就是loss值。
在公式中，带帽子的Y是预测值，不戴帽子的Y是真实值。
2、reg中心点的loss和wh宽高的loss使用的是普通L1损失函数
reg中心点预测和wh宽高预测都直接采用了特征层坐标的尺寸，也就是在0到128之内。
由于wh宽高预测的loss会比较大，其loss乘上了一个系数，论文是0.1。
reg中心点预测的系数则为1。
总的loss就变成了：
实现代码如下：
def focal_loss(pred, target):
pred = pred.permute(0,2,3,1)
pos_inds = target.eq(1).float()
neg_inds = target.lt(1).float()
neg_weights = torch.pow(1 - target, 4)
pred = torch.clamp(pred, 1e-12)
pos_loss = torch.log(pred) * torch.pow(1 - pred, 2) * pos_inds
neg_loss = torch.log(1 - pred) * torch.pow(pred, 2) * neg_weights * neg_inds
num_pos = pos_inds.float().sum()
pos_loss = pos_loss.sum()
neg_loss = neg_loss.sum()
if num_pos == 0:
loss = -neg_loss
else:
loss = -(pos_loss + neg_loss) / num_pos
return loss
def reg_l1_loss(pred, target, mask):
pred = pred.permute(0,2,3,1)
expand_mask = torch.unsqueeze(mask,-1).repeat(1,1,1,2)
loss = F.l1_loss(pred * expand_mask, target * expand_mask, reduction='sum')
loss = loss / (mask.sum() + 1e-4)
return loss
c_loss = focal_loss(hm, batch_hms)
wh_loss = 0.1*reg_l1_loss(wh, batch_whs, batch_reg_masks)
off_loss = reg_l1_loss(offset, batch_regs, batch_reg_masks)
loss = c_loss + wh_loss + off_loss
loss.backward()
optimizer.step()
训练自己的Centernet模型
首先前往Github下载对应的仓库，下载完后利用解压软件解压，之后用编程软件打开文件夹。
注意打开的根目录必须正确，否则相对目录不正确的情况下，代码将无法运行。
一定要注意打开后的根目录是文件存放的目录。
一、数据集的准备
本文使用VOC格式进行训练，训练前需要自己制作好数据集，如果没有自己的数据集，可以通过Github连接下载VOC12+07的数据集尝试下。
训练前将标签文件放在VOCdevkit文件夹下的VOC2007文件夹下的Annotation中。
训练前将图片文件放在VOCdevkit文件夹下的VOC2007文件夹下的JPEGImages中。
此时数据集的摆放已经结束。
二、数据集的处理
在完成数据集的摆放之后，我们需要对数据集进行下一步的处理，目的是获得训练用的2007_train.txt以及2007_val.txt，需要用到根目录下的voc_annotation.py。
voc_annotation.py里面有一些参数需要设置。
分别是annotation_mode、classes_path、trainval_percent、train_percent、VOCdevkit_path，第一次训练可以仅修改classes_path
'''
annotation_mode用于指定该文件运行时计算的内容
annotation_mode为0代表整个标签处理过程，包括获得VOCdevkit/VOC2007/ImageSets里面的txt以及训练用的2007_train.txt、2007_val.txt
annotation_mode为1代表获得VOCdevkit/VOC2007/ImageSets里面的txt
annotation_mode为2代表获得训练用的2007_train.txt、2007_val.txt
'''
annotation_mode = 0
'''
必须要修改，用于生成2007_train.txt、2007_val.txt的目标信息
与训练和预测所用的classes_path一致即可
如果生成的2007_train.txt里面没有目标信息
那么就是因为classes没有设定正确
仅在annotation_mode为0和2的时候有效
'''
classes_path = 'model_data/voc_classes.txt'
'''
trainval_percent用于指定(训练集+验证集)与测试集的比例，默认情况下 (训练集+验证集):测试集 = 9:1
train_percent用于指定(训练集+验证集)中训练集与验证集的比例，默认情况下 训练集:验证集 = 9:1
仅在annotation_mode为0和1的时候有效
'''
trainval_percent = 0.9
train_percent = 0.9
'''
指向VOC数据集所在的文件夹
默认指向根目录下的VOC数据集
'''
VOCdevkit_path = 'VOCdevkit'
classes_path用于指向检测类别所对应的txt，以voc数据集为例，我们用的txt为：
训练自己的数据集时，可以自己建立一个cls_classes.txt，里面写自己所需要区分的类别。
三、开始网络训练
通过voc_annotation.py我们已经生成了2007_train.txt以及2007_val.txt，此时我们可以开始训练了。
训练的参数较多，大家可以在下载库后仔细看注释，其中最重要的部分依然是train.py里的classes_path。
classes_path用于指向检测类别所对应的txt，这个txt和voc_annotation.py里面的txt一样！训练自己的数据集必须要修改！
修改完classes_path后就可以运行train.py开始训练了，在训练多个epoch后，权值会生成在logs文件夹中。
其它参数的作用如下：
#-------------------------------#
# 是否使用Cuda
# 没有GPU可以设置成False
#-------------------------------#
Cuda = True
#--------------------------------------------------------#
# 训练前一定要修改classes_path，使其对应自己的数据集
#--------------------------------------------------------#
classes_path = 'model_data/voc_classes.txt'
#----------------------------------------------------------------------------------------------------------------------------#
# 权值文件请看README，百度网盘下载。数据的预训练权重对不同数据集是通用的，因为特征是通用的。
# 预训练权重对于99%的情况都必须要用，不用的话权值太过随机，特征提取效果不明显，网络训练的结果也不会好。
#
# 如果想要断点续练就将model_path设置成logs文件夹下已经训练的权值文件。
# 当model_path = ''的时候不加载整个模型的权值。
#
# 此处使用的是整个模型的权重，因此是在train.py进行加载的，pretrain不影响此处的权值加载。
# 如果想要让模型从主干的预训练权值开始训练，则设置model_path = ''，pretrain = True，此时仅加载主干。
# 如果想要让模型从0开始训练，则设置model_path = ''，pretrain = Fasle，Freeze_Train = Fasle，此时从0开始训练，且没有冻结主干的过程。
#----------------------------------------------------------------------------------------------------------------------------#
model_path = 'model_data/centernet_resnet50_voc.pth'
#------------------------------------------------------#
# 输入的shape大小，32的倍数
#------------------------------------------------------#
input_shape = [512, 512]
#-------------------------------------------#
# 主干特征提取网络的选择
# resnet50和hourglass
#-------------------------------------------#
backbone = "resnet50"
#----------------------------------------------------------------------------------------------------------------------------#
# 是否使用主干网络的预训练权重，此处使用的是主干的权重，因此是在模型构建的时候进行加载的。
# 如果设置了model_path，则主干的权值无需加载，pretrained的值无意义。
# 如果不设置model_path，pretrained = True，此时仅加载主干开始训练。
# 如果不设置model_path，pretrained = False，Freeze_Train = Fasle，此时从0开始训练，且没有冻结主干的过程。
#----------------------------------------------------------------------------------------------------------------------------#
pretrained = False
#----------------------------------------------------#
# 训练分为两个阶段，分别是冻结阶段和解冻阶段。
# 显存不足与数据集大小无关，提示显存不足请调小batch_size。
# 受到BatchNorm层影响，batch_size最小为2，不能为1。
#----------------------------------------------------#
#----------------------------------------------------#
# 冻结阶段训练参数
# 此时模型的主干被冻结了，特征提取网络不发生改变
# 占用的显存较小，仅对网络进行微调
#----------------------------------------------------#
Init_Epoch = 0
Freeze_Epoch = 50
Freeze_batch_size = 16
Freeze_lr = 1e-3
#----------------------------------------------------#
# 解冻阶段训练参数
# 此时模型的主干不被冻结了，特征提取网络会发生改变
# 占用的显存较大，网络所有的参数都会发生改变
#----------------------------------------------------#
UnFreeze_Epoch = 100
Unfreeze_batch_size = 8
Unfreeze_lr = 1e-4
#------------------------------------------------------#
# 是否进行冻结训练，默认先冻结主干训练后解冻训练。
#------------------------------------------------------#
Freeze_Train = True
#------------------------------------------------------#
# 用于设置是否使用多线程读取数据
# 开启后会加快数据读取速度，但是会占用更多内存
# 内存较小的电脑可以设置为2或者0
#------------------------------------------------------#
num_workers = 4
#----------------------------------------------------#
# 获得图片路径和标签
#----------------------------------------------------#
train_annotation_path = '2007_train.txt'
val_annotation_path = '2007_val.txt'
四、训练结果预测
训练结果预测需要用到两个文件，分别是centernet.py和predict.py。
我们首先需要去centernet.py里面修改model_path以及classes_path，这两个参数必须要修改。
model_path指向训练好的权值文件，在logs文件夹里。
classes_path指向检测类别所对应的txt。
完成修改后就可以运行predict.py进行检测了。运行后输入图片路径即可检测。
657
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
