睿智的目标检测48——Tensorflow2 搭建自己的Centernet目标检测平台
睿智的目标检测48——Tensorflow2 搭建自己的Centernet目标检测平台
- 学习前言
- 什么是Centernet目标检测算法
- 源码下载
- Centernet实现思路
-
- 一、预测部分
-
- 1、主干网络介绍
- 2、利用初步特征获得高分辨率特征图
- 3、Center Head从特征获取预测结果
- 4、预测结果的解码
- 5、在原图上进行绘制
- 二、训练部分
-
- 1、真实框的处理
- 2、利用处理完的真实框与对应图片的预测结果计算loss
- 训练自己的Centernet模型
学习前言
Tensorflow2版本的实现也要做一下。
什么是Centernet目标检测算法
如今常见的目标检测算法通常使用先验框的设定，即先在图片上设定大量的先验框，网络的预测结果会对先验框进行调整获得预测框，先验框很大程度上提高了网络的检测能力，但是也会收到物体尺寸的限制。
Centernet采用不同的方法，构建模型时将目标作为一个点——即目标BBox的中心点。
Centernet的检测器采用关键点估计来找到中心点，并回归到其他目标属性。
论文中提到：模型是端到端可微的，更简单，更快，更精确。Centernet的模型实现了速度和精确的很好权衡。
源码下载
https://github.com/bubbliiiing/centernet-tf2
喜欢的可以点个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。
#-------------------------------------------------------------#
# ResNet50的网络部分
#-------------------------------------------------------------#
import numpy as np
import tensorflow.keras.backend as K
from tensorflow.keras import layers
from tensorflow.keras.layers import (Activation, AveragePooling2D,
BatchNormalization, Conv2D,
Conv2DTranspose, Dense, Dropout, Flatten,
Input, MaxPooling2D, ZeroPadding2D)
from tensorflow.keras.models import Model
from tensorflow.keras.preprocessing import image
from tensorflow.keras.regularizers import l2
def identity_block(input_tensor, kernel_size, filters, stage, block):
filters1, filters2, filters3 = filters
conv_name_base = 'res' + str(stage) + block + '_branch'
bn_name_base = 'bn' + str(stage) + block + '_branch'
x = Conv2D(filters1, (1, 1), name=conv_name_base + '2a', use_bias=False)(input_tensor)
x = BatchNormalization(name=bn_name_base + '2a')(x)
x = Activation('relu')(x)
x = Conv2D(filters2, kernel_size,padding='same', name=conv_name_base + '2b', use_bias=False)(x)
x = BatchNormalization(name=bn_name_base + '2b')(x)
x = Activation('relu')(x)
x = Conv2D(filters3, (1, 1), name=conv_name_base + '2c', use_bias=False)(x)
x = BatchNormalization(name=bn_name_base + '2c')(x)
x = layers.add([x, input_tensor])
x = Activation('relu')(x)
return x
def conv_block(input_tensor, kernel_size, filters, stage, block, strides=(2, 2)):
filters1, filters2, filters3 = filters
conv_name_base = 'res' + str(stage) + block + '_branch'
bn_name_base = 'bn' + str(stage) + block + '_branch'
x = Conv2D(filters1, (1, 1), strides=strides,
name=conv_name_base + '2a', use_bias=False)(input_tensor)
x = BatchNormalization(name=bn_name_base + '2a')(x)
x = Activation('relu')(x)
x = Conv2D(filters2, kernel_size, padding='same',
name=conv_name_base + '2b', use_bias=False)(x)
x = BatchNormalization(name=bn_name_base + '2b')(x)
x = Activation('relu')(x)
x = Conv2D(filters3, (1, 1), name=conv_name_base + '2c', use_bias=False)(x)
x = BatchNormalization(name=bn_name_base + '2c')(x)
shortcut = Conv2D(filters3, (1, 1), strides=strides,
name=conv_name_base + '1', use_bias=False)(input_tensor)
shortcut = BatchNormalization(name=bn_name_base + '1')(shortcut)
x = layers.add([x, shortcut])
x = Activation('relu')(x)
return x
def ResNet50(inputs):
# 512x512x3
x = ZeroPadding2D((3, 3))(inputs)
# 256,256,64
x = Conv2D(64, (7, 7), strides=(2, 2), name='conv1', use_bias=False)(x)
x = BatchNormalization(name='bn_conv1')(x)
x = Activation('relu')(x)
# 256,256,64 -> 128,128,64
x = MaxPooling2D((3, 3), strides=(2, 2), padding="same")(x)
# 128,128,64 -> 128,128,256
x = conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1))
x = identity_block(x, 3, [64, 64, 256], stage=2, block='b')
x = identity_block(x, 3, [64, 64, 256], stage=2, block='c')
# 128,128,256 -> 64,64,512
x = conv_block(x, 3, [128, 128, 512], stage=3, block='a')
x = identity_block(x, 3, [128, 128, 512], stage=3, block='b')
x = identity_block(x, 3, [128, 128, 512], stage=3, block='c')
x = identity_block(x, 3, [128, 128, 512], stage=3, block='d')
# 64,64,512 -> 32,32,1024
x = conv_block(x, 3, [256, 256, 1024], stage=4, block='a')
x = identity_block(x, 3, [256, 256, 1024], stage=4, block='b')
x = identity_block(x, 3, [256, 256, 1024], stage=4, block='c')
x = identity_block(x, 3, [256, 256, 1024], stage=4, block='d')
x = identity_block(x, 3, [256, 256, 1024], stage=4, block='e')
x = identity_block(x, 3, [256, 256, 1024], stage=4, block='f')
# 32,32,1024 -> 16,16,2048
x = conv_block(x, 3, [512, 512, 2048], stage=5, block='a')
x = identity_block(x, 3, [512, 512, 2048], stage=5, block='b')
x = identity_block(x, 3, [512, 512, 2048], stage=5, block='c')
return x
2、利用初步特征获得高分辨率特征图
利用上一步获得到的resnet50的最后一个特征层的shape为(16,16,2048)。
对于该特征层，centernet利用三次反卷积进行上采样，从而更高的分辨率输出。为了节省计算量，这3个反卷积的输出通道数分别为256,128,64。
每一次反卷积，特征层的高和宽会变为原来的两倍，因此，在进行三次反卷积上采样后，我们获得的特征层的高和宽变为原来的8倍，此时特征层的高和宽为128x128，通道数为64。
此时我们获得了一个128x128x64的有效特征层（高分辨率特征图），我们会利用该有效特征层获得最终的预测结果。
实现代码如下：
x = Dropout(rate=0.5)(x)
#-------------------------------#
# 解码器
#-------------------------------#
num_filters = 256
# 16, 16, 2048 -> 32, 32, 256 -> 64, 64, 128 -> 128, 128, 64
for i in range(3):
# 进行上采样
x = Conv2DTranspose(num_filters // pow(2, i), (4, 4), strides=2, use_bias=False, padding='same',
kernel_initializer='he_normal',
kernel_regularizer=l2(5e-4))(x)
x = BatchNormalization()(x)
x = Activation('relu')(x)
3、Center Head从特征获取预测结果
通过上一步我们可以获得一个128x128x64的高分辨率特征图。
这个特征层相当于将整个图片划分成128x128个区域，每个区域存在一个特征点，如果某个物体的中心落在这个区域，那么就由这个特征点来确定。
（某个物体的中心落在这个区域，则由这个区域左上角的特征点来约定）
我们可以利用这个特征层进行三个卷积，分别是：
1、热力图预测，此时卷积的通道数为num_classes，最终结果为(128,128,num_classes)，代表每一个热力点是否有物体存在，以及物体的种类；
2、中心点预测，此时卷积的通道数为2，最终结果为(128,128,2)，代表每一个物体中心距离热力点偏移的情况；
3、宽高预测，此时卷积的通道数为2，最终结果为(128,128,2)，代表每一个物体宽高的预测情况；
实现代码为：
def centernet_head(x,num_classes):
x = Dropout(rate=0.5)(x)
#-------------------------------#
# 解码器
#-------------------------------#
num_filters = 256
# 16, 16, 2048 -> 32, 32, 256 -> 64, 64, 128 -> 128, 128, 64
for i in range(3):
# 进行上采样
x = Conv2DTranspose(num_filters // pow(2, i), (4, 4), strides=2, use_bias=False, padding='same',
kernel_initializer='he_normal',
kernel_regularizer=l2(5e-4))(x)
x = BatchNormalization()(x)
x = Activation('relu')(x)
# 最终获得128,128,64的特征层
# hm header
y1 = Conv2D(64, 3, padding='same', use_bias=False, kernel_initializer='he_normal', kernel_regularizer=l2(5e-4))(x)
y1 = BatchNormalization()(y1)
y1 = Activation('relu')(y1)
y1 = Conv2D(num_classes, 1, kernel_initializer='he_normal', kernel_regularizer=l2(5e-4), activation='sigmoid')(y1)
# wh header
y2 = Conv2D(64, 3, padding='same', use_bias=False, kernel_initializer='he_normal', kernel_regularizer=l2(5e-4))(x)
y2 = BatchNormalization()(y2)
y2 = Activation('relu')(y2)
y2 = Conv2D(2, 1, kernel_initializer='he_normal', kernel_regularizer=l2(5e-4))(y2)
# reg header
y3 = Conv2D(64, 3, padding='same', use_bias=False, kernel_initializer='he_normal', kernel_regularizer=l2(5e-4))(x)
y3 = BatchNormalization()(y3)
y3 = Activation('relu')(y3)
y3 = Conv2D(2, 1, kernel_initializer='he_normal', kernel_regularizer=l2(5e-4))(y3)
return y1, y2, y3
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 nms(heat, kernel=3):
hmax = MaxPooling2D((kernel, kernel), strides=1, padding='SAME')(heat)
heat = tf.where(tf.equal(hmax, heat), heat, tf.zeros_like(heat))
return heat
def topk(hm, max_objects=100):
# hm -> Hot map热力图
# 进行热力图的非极大抑制，利用3x3的卷积对热力图进行Max筛选，找出值最大的
hm = nms(hm)
# (b, h * w * c)
b, h, w, c = tf.shape(hm)[0], tf.shape(hm)[1], tf.shape(hm)[2], tf.shape(hm)[3]
# 将所有结果平铺，获得(b, h * w * c)
hm = tf.reshape(hm, (b, -1))
# (b, k), (b, k)
scores, indices = tf.math.top_k(hm, k=max_objects, sorted=True)
# 计算求出网格点，类别
class_ids = indices % c
xs = indices // c % w
ys = indices // c // w
indices = ys * w + xs
return scores, indices, class_ids, xs, ys
def decode(hm, wh, reg, max_objects=100,num_classes=20):
scores, indices, class_ids, xs, ys = topk(hm, max_objects=max_objects)
# 获得batch_size
b = tf.shape(hm)[0]
# (b, h * w, 2)
reg = tf.reshape(reg, [b, -1, 2])
# (b, h * w, 2)
wh = tf.reshape(wh, [b, -1, 2])
length = tf.shape(wh)[1]
# 找到其在1维上的索引
batch_idx = tf.expand_dims(tf.range(0, b), 1)
batch_idx = tf.tile(batch_idx, (1, max_objects))
full_indices = tf.reshape(batch_idx, [-1]) * tf.to_int32(length) + tf.reshape(indices, [-1])
# 取出top_k个框对应的参数
topk_reg = tf.gather(tf.reshape(reg, [-1,2]), full_indices)
topk_reg = tf.reshape(topk_reg, [b, -1, 2])
topk_wh = tf.gather(tf.reshape(wh, [-1,2]), full_indices)
topk_wh = tf.reshape(topk_wh, [b, -1, 2])
# 计算调整后的中心
topk_cx = tf.cast(tf.expand_dims(xs, axis=-1), tf.float32) + topk_reg[..., 0:1]
topk_cy = tf.cast(tf.expand_dims(ys, axis=-1), tf.float32) + topk_reg[..., 1:2]
# (b,k,1) (b,k,1)
topk_x1, topk_y1 = topk_cx - topk_wh[..., 0:1] / 2, topk_cy - topk_wh[..., 1:2] / 2
# (b,k,1) (b,k,1)
topk_x2, topk_y2 = topk_cx + topk_wh[..., 0:1] / 2, topk_cy + topk_wh[..., 1:2] / 2
# (b,k,1)
scores = tf.expand_dims(scores, axis=-1)
# (b,k,1)
class_ids = tf.cast(tf.expand_dims(class_ids, axis=-1), tf.float32)
# (b,k,6)
detections = tf.concat([topk_x1, topk_y1, topk_x2, topk_y2, scores, class_ids], axis=-1)
return detections
5、在原图上进行绘制
通过第三步，我们可以获得预测框在原图上的位置，而且这些预测框都是经过筛选的。这些筛选后的框可以直接绘制在图片上，就可以获得结果了。
二、训练部分
1、真实框的处理
既然在centernet中，物体的中心落在哪个特征点的右下角就由哪个特征点来负责预测，那么在训练的时候我们就需要找到真实框和特征点之间的关系。
真实框和特征点之间的关系，对应方式如下：
1、找到真实框的中心，通过真实框的中心找到其对应的特征点。
2、根据该真实框的种类，对网络应该有的热力图进行设置，即heatmap热力图。其实就是将对应的特征点里面的对应的种类，它的中心值设置为1，然后这个特征点附近的其它特征点中该种类对应的值按照高斯分布不断下降。
3、除去heatmap热力图外，还需要设置特征点对应的reg中心点和wh宽高，在找到真实框对应的特征点后，还需要设置该特征点对应的reg中心和wh宽高。这里的reg中心和wh宽高都是对于128x128的特征层的。
4、在获得网络应该有的预测结果后，就可以将预测结果和应该有的预测结果进行对比，对网络进行反向梯度调整了。
实现代码如下：
class Generator(object):
def __init__(self,batch_size,train_lines,val_lines,
input_size,num_classes,max_objects=100):
self.batch_size = batch_size
self.train_lines = train_lines
self.val_lines = val_lines
self.input_size = input_size
self.output_size = (int(input_size[0]/4) , int(input_size[1]/4))
self.num_classes = num_classes
self.max_objects = max_objects
def get_random_data(self, annotation_line, input_shape, random=True, jitter=.3, hue=.1, sat=1.5, val=1.5, proc_img=True):
'''r实时数据增强的随机预处理'''
line = annotation_line.split()
image = Image.open(line[0])
iw, ih = image.size
h, w = input_shape
box = np.array([np.array(list(map(int,box.split(',')))) for box in line[1:]])
# resize image
new_ar = w/h * rand(1-jitter,1+jitter)/rand(1-jitter,1+jitter)
scale = rand(0.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)
# place image
dx = int(rand(0, w-nw))
dy = int(rand(0, h-nh))
new_image = Image.new('RGB', (w,h), (128,128,128))
new_image.paste(image, (dx, dy))
image = new_image
# flip image or not
flip = rand()<.5
if flip: image = image.transpose(Image.FLIP_LEFT_RIGHT)
# distort image
hue = rand(-hue, hue)
sat = rand(1, sat) if rand()<.5 else 1/rand(1, sat)
val = rand(1, val) if rand()<.5 else 1/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
# correct boxes
box_data = np.zeros((len(box),5))
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)] # discard invalid box
box_data = np.zeros((len(box),5))
box_data[:len(box)] = box
if len(box) == 0:
return image_data, []
if (box_data[:,:4]>0).any():
return image_data, box_data
else:
return image_data, []
def generate(self, train=True):
while True:
if train:
# 打乱
shuffle(self.train_lines)
lines = self.train_lines
else:
shuffle(self.val_lines)
lines = self.val_lines
batch_images = np.zeros((self.batch_size, self.input_size[0], self.input_size[1], self.input_size[2]), dtype=np.float32)
batch_hms = np.zeros((self.batch_size, self.output_size[0], self.output_size[1], self.num_classes), dtype=np.float32)
batch_whs = np.zeros((self.batch_size, self.max_objects, 2), dtype=np.float32)
batch_regs = np.zeros((self.batch_size, self.max_objects, 2), dtype=np.float32)
batch_reg_masks = np.zeros((self.batch_size, self.max_objects), dtype=np.float32)
batch_indices = np.zeros((self.batch_size, self.max_objects), dtype=np.float32)
b = 0
for annotation_line in lines:
img,y=self.get_random_data(annotation_line,self.input_size[0:2])
if len(y)!=0:
boxes = np.array(y[:,:4],dtype=np.float32)
boxes[:,0] = boxes[:,0]/self.input_size[1]*self.output_size[1]
boxes[:,1] = boxes[:,1]/self.input_size[0]*self.output_size[0]
boxes[:,2] = boxes[:,2]/self.input_size[1]*self.output_size[1]
boxes[:,3] = boxes[:,3]/self.input_size[0]*self.output_size[0]
for i in range(len(y)):
bbox = boxes[i].copy()
bbox = np.array(bbox)
bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, self.output_size[1] - 1)
bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, self.output_size[0] - 1)
cls_id = int(y[i,-1])
h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]
if h > 0 and w > 0:
ct = np.array([(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2], dtype=np.float32)
ct_int = ct.astype(np.int32)
# 获得热力图
radius = gaussian_radius((math.ceil(h), math.ceil(w)))
radius = max(0, int(radius))
batch_hms[b, :, :, cls_id] = draw_gaussian(batch_hms[b, :, :, cls_id], ct_int, radius)
batch_whs[b, i] = 1. * w, 1. * h
# 计算中心偏移量
batch_regs[b, i] = ct - ct_int
# 将对应的mask设置为1，用于排除多余的0
batch_reg_masks[b, i] = 1
# 表示第ct_int[1]行的第ct_int[0]个。
batch_indices[b, i] = ct_int[1] * self.output_size[0] + ct_int[0]
batch_images[b] = preprocess_image(img)
b = b + 1
if b == self.batch_size:
b = 0
yield [batch_images, batch_hms, batch_whs, batch_regs, batch_reg_masks, batch_indices], np.zeros((self.batch_size,))
batch_images = np.zeros((self.batch_size, self.input_size[0], self.input_size[1], 3), dtype=np.float32)
batch_hms = np.zeros((self.batch_size, self.output_size[0], self.output_size[1], self.num_classes),
dtype=np.float32)
batch_whs = np.zeros((self.batch_size, self.max_objects, 2), dtype=np.float32)
batch_regs = np.zeros((self.batch_size, self.max_objects, 2), dtype=np.float32)
batch_reg_masks = np.zeros((self.batch_size, self.max_objects), dtype=np.float32)
batch_indices = np.zeros((self.batch_size, self.max_objects), dtype=np.float32)
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(hm_pred, hm_true):
# 找到正样本和负样本
pos_mask = tf.cast(tf.equal(hm_true, 1), tf.float32)
# 小于1的都是负样本
neg_mask = tf.cast(tf.less(hm_true, 1), tf.float32)
neg_weights = tf.pow(1 - hm_true, 4)
pos_loss = -tf.log(tf.clip_by_value(hm_pred, 1e-7, 1.)) * tf.pow(1 - hm_pred, 2) * pos_mask
neg_loss = -tf.log(tf.clip_by_value(1 - hm_pred, 1e-7, 1.)) * tf.pow(hm_pred, 2) * neg_weights * neg_mask
num_pos = tf.reduce_sum(pos_mask)
pos_loss = tf.reduce_sum(pos_loss)
neg_loss = tf.reduce_sum(neg_loss)
cls_loss = tf.cond(tf.greater(num_pos, 0), lambda: (pos_loss + neg_loss) / num_pos, lambda: neg_loss)
return cls_loss
def reg_l1_loss(y_pred, y_true, indices, mask):
b, c = tf.shape(y_pred)[0], tf.shape(y_pred)[-1]
k = tf.shape(indices)[1]
y_pred = tf.reshape(y_pred, (b, -1, c))
length = tf.shape(y_pred)[1]
indices = tf.cast(indices, tf.int32)
# 找到其在1维上的索引
batch_idx = tf.expand_dims(tf.range(0, b), 1)
batch_idx = tf.tile(batch_idx, (1, k))
full_indices = (tf.reshape(batch_idx, [-1]) * tf.to_int32(length) +
tf.reshape(indices, [-1]))
# 取出对应的预测值
y_pred = tf.gather(tf.reshape(y_pred, [-1,c]),full_indices)
y_pred = tf.reshape(y_pred, [b, -1, c])
mask = tf.tile(tf.expand_dims(mask, axis=-1), (1, 1, 2))
# 求取l1损失值
total_loss = tf.reduce_sum(tf.abs(y_true * mask - y_pred * mask))
reg_loss = total_loss / (tf.reduce_sum(mask) + 1e-4)
return reg_loss
def loss(args):
#-----------------------------------------------------------------------------------------------------------------#
# hm_pred：热力图的预测值 (self.batch_size, self.output_size[0], self.output_size[1], self.num_classes)
# wh_pred：宽高的预测值 (self.batch_size, self.output_size[0], self.output_size[1], 2)
# reg_pred：中心坐标偏移预测值 (self.batch_size, self.output_size[0], self.output_size[1], 2)
# hm_true：热力图的真实值 (self.batch_size, self.output_size[0], self.output_size[1], self.num_classes)
# wh_true：宽高的真实值 (self.batch_size, self.max_objects, 2)
# reg_true：中心坐标偏移真实值 (self.batch_size, self.max_objects, 2)
# reg_mask：真实值的mask (self.batch_size, self.max_objects)
# indices：真实值对应的坐标 (self.batch_size, self.max_objects)
#-----------------------------------------------------------------------------------------------------------------#
hm_pred, wh_pred, reg_pred, hm_true, wh_true, reg_true, reg_mask, indices = args
hm_loss = focal_loss(hm_pred, hm_true)
wh_loss = 0.1 * reg_l1_loss(wh_pred, wh_true, indices, reg_mask)
reg_loss = reg_l1_loss(reg_pred, reg_true, indices, reg_mask)
total_loss = hm_loss + wh_loss + reg_loss
# total_loss = tf.Print(total_loss,[hm_loss,wh_loss,reg_loss])
return total_loss
训练自己的Centernet模型
Centernet整体的文件夹构架如下：
本文使用VOC格式进行训练。
训练前将标签文件放在VOCdevkit文件夹下的VOC2007文件夹下的Annotation中。
训练前将图片文件放在VOCdevkit文件夹下的VOC2007文件夹下的JPEGImages中。
在训练前利用voc2centernet.py文件生成对应的txt。
再运行根目录下的voc_annotation.py，运行前需要将classes改成你自己的classes。
classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
就会生成对应的2007_train.txt，每一行对应其图片位置及其真实框的位置。
在训练前需要修改model_data里面的voc_classes.txt文件，需要将classes改成你自己的classes。
然后运行train.py就可以开始训练啦。
运行train.py即可开始训练。
