YOLO V3版本是一个强大和快速的物体检测模型，同时原理上也相对简单。我之前的博客中已经介绍了如何用Tensorflow来实现YOLO V1版本，之后我自己也用Tensorflow 1.X版本实现了YOLO V3，现在Tensorflow演进到了2.0版本，相比较１.X版本做了很大的改进，也更加易用了，因此我记录一下如何用Tensorflow 2.0版本来实现YOLO V3。网上能找到的很多Tensorflow YOLO V3的代码都没有完整的一个训练过程，基本上都是转换和加载YOLO的作者在Darknet上发布的训练好的权重数据，直接进行检测的。我的这个代码实现了完整的训练流程，包括了搭建基础架构网络Darknet53进行Imagenet预训练，以及增加YOLO V3网络模块进行物体检测训练，对模型的训练效果进行评测，以及用训练好的模型进行物体检测的过程。
训练数据的准备
需要准备两份训练数据，一个是Imagenent的物体分类数据，包括了1000种类别的物体的数据，共128万张图片。数据集需要预先处理为TFRECORD格式，具体过程可以参见我之前的博客https://blog.csdn.net/gzroy/article/details/85954329。 第二个训练数据是物体检测的数据，目前有很多个数据集可以采用，例如COCO数据集（包括80种物体的检测框），OpenImage，Pascal VOC等等，比较流行的是COCO数据集，大部分物体检测的论文都会基于这个数据集来提供性能指标。我也采用COCO数据集，同样也是预处理为TFRECORD格式，具体过程可以参见我的另一篇博客https://blog.csdn.net/gzroy/article/details/95027532
网络模型的搭建
按照YOLO V3论文的描述，基础网络架构是一个叫做Darknet53的网络模型，共有53个卷积层，其网络架构如下：
用Tensorflow可以很方便的构建一个Darknet53模型，代码如下：
import tensorflow as tf
from tensorflow.keras import Model
l=tf.keras.layers
def _conv(inputs, filters, kernel_size, strides, padding, bias=False, normalize=True, activation='relu', last=False):
output = inputs
padding_str = 'same'
if padding>0:
output = l.ZeroPadding2D(padding=padding, data_format='channels_first')(output)
padding_str = 'valid'
output = l.Conv2D(filters, kernel_size, strides, padding_str, \
'channels_first', use_bias=bias, \
kernel_initializer='he_normal', \
kernel_regularizer=tf.keras.regularizers.l2(l=5e-4))(output)
if normalize:
if not last:
output = l.BatchNormalization(axis=1)(output)
else:
output = l.BatchNormalization(axis=1, gamma_initializer='zeros')(output)
if activation=='relu':
output = l.ReLU()(output)
if activation=='relu6':
output = l.ReLU(max_value=6)(output)
if activation=='leaky_relu':
output = l.LeakyReLU(alpha=0.1)(output)
return output
def _residual(inputs, out_channels, activation='relu', name=None):
output1 = _conv(inputs, out_channels//2, 1, 1, 0, False, True, 'leaky_relu', False)
output2 = _conv(output1, out_channels, 3, 1, 1, False, True, 'leaky_relu', True)
output = l.Add(name=name)([inputs, output2])
return output
def darknet53_base():
image = tf.keras.Input(shape=(3,None,None))
net = _conv(image, 32, 3, 1, 1, False, True, 'leaky_relu') #32*H*W
net = _conv(net, 64, 3, 2, 1, False, True, 'leaky_relu') #64*H/2*W/2
net = _residual(net, 64, 'leaky_relu') #64*H/2*W/2
net = _conv(net, 128, 3, 2, 1, False, True, 'leaky_relu') #128*H/4*W/4
net = _residual(net, 128, 'leaky_relu') #128*H/4*W/4
net = _residual(net, 128, 'leaky_relu') #128*H/4*W/4
net = _conv(net, 256, 3, 2, 1, False, True, 'leaky_relu') #256*H/8*W/8
net = _residual(net, 256, 'leaky_relu') #256*H/8*W/8
net = _residual(net, 256, 'leaky_relu') #256*H/8*W/8
net = _residual(net, 256, 'leaky_relu') #256*H/8*W/8
net = _residual(net, 256, 'leaky_relu') #256*H/8*W/8
net = _residual(net, 256, 'leaky_relu') #256*H/8*W/8
net = _residual(net, 256, 'leaky_relu') #256*H/8*W/8
net = _residual(net, 256, 'leaky_relu') #256*H/8*W/8
net = _residual(net, 256, 'leaky_relu') #256*H/8*W/8
route1 = l.Activation('linear', dtype='float32', name='route1')(net)
net = _conv(net, 512, 3, 2, 1, False, True, 'leaky_relu') #512*H/16*W/16
net = _residual(net, 512, 'leaky_relu') #512*H/16*W/16
net = _residual(net, 512, 'leaky_relu') #512*H/16*W/16
net = _residual(net, 512, 'leaky_relu') #512*H/16*W/16
net = _residual(net, 512, 'leaky_relu') #512*H/16*W/16
net = _residual(net, 512, 'leaky_relu') #512*H/16*W/16
net = _residual(net, 512, 'leaky_relu') #512*H/16*W/16
net = _residual(net, 512, 'leaky_relu') #512*H/16*W/16
net = _residual(net, 512, 'leaky_relu') #512*H/16*W/16
route2 = l.Activation('linear', dtype='float32', name='route2')(net)
net = _conv(net, 1024, 3, 2, 1, False, True, 'leaky_relu') #1024*H/32*W/32
net = _residual(net, 1024, 'leaky_relu') #1024*H/32*W/32
net = _residual(net, 1024, 'leaky_relu') #1024*H/32*W/32
net = _residual(net, 1024, 'leaky_relu') #1024*H/32*W/32
net = _residual(net, 1024, 'leaky_relu') #1024*H/32*W/32
route3 = l.Activation('linear', dtype='float32', name='route3')(net)
net = tf.reduce_mean(net, axis=[2,3], keepdims=True)
net = _conv(net, 1000, 1, 1, 0, True, False, 'linear') #1000
net = l.Flatten(data_format='channels_first', name='logits')(net)
net = l.Activation('linear', dtype='float32', name='output')(net)
model = tf.keras.Model(inputs=image, outputs=[net, route1, route2, route3])
return model我们需要先基于这个骨干网络架构来进行Imagenet的预训练，以提取有效的图片内容的特征数据。我用这个网络训练了30个EPOCH，最终到达Top-1 71%，Top-5 91%的准确率。具体的训练过程可以见我的博客https://blog.csdn.net/gzroy/article/details/104170537
训练好了骨干网络之后，我们就可以在这个网络的基础上再增加相应的卷积层，实现图像特征金字塔(FPN)的架构，这里我们会用到骨干网络输出的route1, route2, route3这几个不同图像分辨率的特征值，最终构建一个可以对图片进行下采样8倍，16倍和32倍的基于网格的检测系统，例如训练图片的分辨率为416*416，那么将输出52*52, 26*26, 13*13这三个不同维度的检测结果。具体的原理可以参见网上的一些文章，例如：我这里参照Darknet的源代码来搭建了一个YOLO V3的网络，代码如下：
category_num = 80
vector_size = 3*(1+4+category_num)
def darknet53_yolov3():
route1 = tf.keras.Input(shape=(256,None,None), name='input1') #256*H/8*W/8
route2 = tf.keras.Input(shape=(512,None,None), name='input2') #256*H/16*W/16
route3 = tf.keras.Input(shape=(1024,None,None), name='input3') #256*H/32*W/32
net = _conv(route3, 512, 1, 1, 0, False, True, 'leaky_relu') #512*H/32*W/32
net = _conv(net, 1024, 3, 1, 1, False, True, 'leaky_relu') #1024*H/32*W/32
net = _conv(net, 512, 1, 1, 0, False, True, 'leaky_relu') #512*H/32*W/32
net = _conv(net, 1024, 3, 1, 1, False, True, 'leaky_relu') #1024*H/32*W/32
net = _conv(net, 512, 1, 1, 0, False, True, 'leaky_relu') #512*H/32*W/32
route4 = tf.identity(net, 'route4')
net = _conv(net, 1024, 3, 1, 1, False, True, 'leaky_relu') #1024*H/32*W/32
predict1 = _conv(net, vector_size, 1, 1, 0, True, False, 'linear') #vector_size*H/32*W/32
predict1 = l.Activation('linear', dtype='float32')(predict1)
predict1 = l.Reshape((vector_size, imageHeight//32*imageWidth//32))(predict1)
net = _conv(route4, 256, 1, 1, 0, False, True, 'leaky_relu') #256*H/32*W/32
net = l.UpSampling2D((2,2),"channels_first",'nearest')(net) #256*H/16*W/16
net = l.Concatenate(axis=1)([route2, net]) #768*H/16*W/16
net = _conv(net, 256, 1, 1, 0, False, True, 'leaky_relu') #256*H/16*W/16
net = _conv(net, 512, 3, 1, 1, False, True, 'leaky_relu') #512*H/16*W/16
net = _conv(net, 256, 1, 1, 0, False, True, 'leaky_relu') #256*H/16*W/16
net = _conv(net, 512, 3, 1, 1, False, True, 'leaky_relu') #512*H/16*W/16
net = _conv(net, 256, 1, 1, 0, False, True, 'leaky_relu') #256*H/16*W/16
route5 = tf.identity(net, 'route5')
net = _conv(net, 512, 3, 1, 1, False, True, 'leaky_relu') #512*H/16*W/16
predict2 = _conv(net, vector_size, 1, 1, 0, True, False, 'linear') #vector_size*H/16*W/16
predict2 = l.Activation('linear', dtype='float32')(predict2)
predict2 = l.Reshape((vector_size, imageHeight//16*imageWidth//16))(predict2)
net = _conv(route5, 128, 1, 1, 0, False, True, 'leaky_relu') #128*H/16*W/16
net = l.UpSampling2D((2,2),"channels_first",'nearest')(net) #128*H/8*W/8
net = l.Concatenate(axis=1)([route1, net]) #384*H/8*W/8
net = _conv(net, 128, 1, 1, 0, False, True, 'leaky_relu') #128*H/8*W/8
net = _conv(net, 256, 3, 1, 1, False, True, 'leaky_relu') #256*H/8*W/8
net = _conv(net, 128, 1, 1, 0, False, True, 'leaky_relu') #128*H/8*W/8
net = _conv(net, 256, 3, 1, 1, False, True, 'leaky_relu') #256*H/8*W/8
net = _conv(net, 128, 1, 1, 0, False, True, 'leaky_relu') #128*H/8*W/8
net = _conv(net, 256, 3, 1, 1, False, True, 'leaky_relu') #256*H/8*W/8
predict3 = _conv(net, vector_size, 1, 1, 0, True, False, 'linear') #vector_size*H/8*W/8
predict3 = l.Activation('linear', dtype='float32')(predict3)
predict3 = l.Reshape((vector_size, imageHeight//8*imageWidth//8))(predict3)
predict = l.Concatenate()([predict3, predict2, predict1])
predict = tf.transpose(predict, perm=[0, 2, 1], name='predict')
model = tf.keras.Model(inputs=[route1, route2, route3], outputs=predict, name='darknet53_yolo')
return model可以看到，这个网络模型是以骨干网络的三个输出route1, route2, route3作为输入的，最终输出的Predict是三个不同维度的预测结果。
YOLO V3训练过程
有了训练数据和搭建好网络模型之后，我们就可以开始训练了。整个训练过程分为如下几步：
1. 对骨干网络进行更高分辨率的训练
因为我们的骨干网络是基于224*224这个分辨率来进行训练和提取图片特征的，但是在物体检测中，这个分辨率太低，不利于检测小物体，因此我们需要基于更加高的分辨率，例如416*416来进行训练。我们可以把骨干网络基于这个高的分辨率再多训练一些次数，让网络适应高分辨率，这样可以最终提升物体检测的性能。为此我们可以重新加载之前训练好的骨干网络来进行训练。
2. 骨干网络和检测网络组合
把预训练完成后的骨干网络和检测网络组合起来，构成一个YOLO V3的网络模型。训练图片先通过骨干网络进行特征提取，输出route1,route2,route3这三个不同维度的图像特征数据，然后作为输入进到检测网络中进行训练，最终得到三个维度的预测结果。这个组合网络中，需要设置骨干网络的参数为不可训练，只训练检测网络的参数。代码如下：
#Load the pretrained backbone model
model_base = tf.keras.models.load_model('darknet53/epoch_60.h5')
model_base.trainable = False
image = tf.keras.Input(shape=(3,image_height,image_width))
_, route1, route2, route3 = model_base(image, training=False)
#The detect model will accept the backmodel output as input
predict = darknet53_yolov3(image_height,image_width)([route1, route2, route3])
#Construct the combined yolo model
model_yolo = tf.keras.Model(inputs=image, outputs=predict, name='model_yolo')3. 读取训练数据并进行预处理
读取COCO训练集的图片和检测框的数据，并进行数据增广，生成数据标签等预处理。这里的数据增广除了按照Darknet源码的处理方式之外，还参照论文https://arxiv.org/pdf/1902.04103.pdf中提出的数据增广处理流程，增加了Mixup的处理，即一次取两张图片，通过随机的透明度的处理之后，同时叠加在一起。
因为涉及到检测框的位置，在做数据增广时，需要相应调整检测框的位置。包括以下几个步骤：
- 图像缩放：随机缩放图像的宽和高(缩放系数为0.7-1.3之间的一个随机数)，并计算缩放后的宽高的比例，然后以缩放后的宽和高的长边为准，缩放为图像输入维度416，并按照比例来缩放短边。
- 图像的填充：因为上一步完成后，图像的长边为416像素，需要对短边进行填充使其也达到416像素。
- 检测框的调整：根据以上图像的变换，相应调整检测框的位置。
- 随机反转图像
- 再次调整检测框
- 随机调整图像的饱和度，明亮度等
- 添加PCA噪声
- 标准化图像的RGB通道值。
- 根据检测框的大小判断其应由哪个Anchor来负责预测
- 图像数据和检测框的数据作为Feature
- 根据检测框的数据生成Label，其维度为(1+1+4+1+80)*3=258，其中第1位为grid id，第2位表示是否存在Object，第3-6位表示如果Object的中央点的坐标和宽高，第7位表示Mixup的比例，最后的80位标识这个Object属于哪一类物体。
首先是定义模型的一些参数，代码如下：
mixup_flag = True
#Parameters for PCA noice
eigvec = tf.constant(
[
[-0.5675, 0.7192, 0.4009],
[-0.5808, -0.0045, -0.8140],
[-0.5836, -0.6948, 0.4203]
],
shape=[3,3],
dtype=tf.float32
)
eigval = tf.constant([55.46, 4.794, 1.148], shape=[3,1], dtype=tf.float32)
#Parameters for normalization
mean_RGB = tf.constant([123.68, 116.779, 109.939], dtype=tf.float32)
std_RGB = tf.constant([58.393, 57.12, 57.375], dtype=tf.float32)
#Train and valid batch size
batch_size = 16
val_batch_size = 10
epoch_size = 118287
epoch_batch = int(epoch_size/batch_size)
#Parameters for yolo loss scale
no_object_scale = 1.0
iou_threshold = 0.7
object_scale=3.0
class_scale=1.0
jitter = 0.3
#Label and prediction vector size
category_num = 80
label_vector_size = 1+1+4+1+category_num #index 0:grid_id,1:obj_conf,2-5:(x,y,w,h),6:mixup weight,7-86:category
vector_size = 1+4+category_num #index 0:obj_conf,1-4:(x,y,w,h),5-84:category
#Images parameter
image_size_list = [320, 352, 384, 416, 448, 480, 512, 544, 576, 608]
image_size = image_size_list[random.randint(0,9)]
val_image_size = 608
#Grids parameter
grid_wh_array = np.array([[8.,8.],[16.,16.],[32.,32.]])
grid_size = [8.,16.,32.]
#The Anchor size for image_size 416*416
anchors_base = [10,13, 16,30, 33,23, 30,61, 62,45, 59,119, 116,90, 156,198, 373,326]定义读取训练文件的函数：
def _parse_function(example_proto):
features = {
"image": tf.io.FixedLenFeature([], tf.string, default_value=""),
"height": tf.io.FixedLenFeature([1], tf.int64, default_value=[0]),
"width": tf.io.FixedLenFeature([1], tf.int64, default_value=[0]),
"channels": tf.io.FixedLenFeature([1], tf.int64, default_value=[3]),
"colorspace": tf.io.FixedLenFeature([], tf.string, default_value=""),
"img_format": tf.io.FixedLenFeature([], tf.string, default_value=""),
"label": tf.io.VarLenFeature(tf.int64),
"bbox_xmin": tf.io.VarLenFeature(tf.int64),
"bbox_xmax": tf.io.VarLenFeature(tf.int64),
"bbox_ymin": tf.io.VarLenFeature(tf.int64),
"bbox_ymax": tf.io.VarLenFeature(tf.int64),
"filename": tf.io.FixedLenFeature([], tf.string, default_value="")
}
parsed_features = tf.io.parse_single_example(example_proto, features)
label = tf.expand_dims(parsed_features["label"].values, 0)
label = tf.cast(label, tf.float32)
image_raw = tf.image.decode_jpeg(parsed_features["image"], channels=3)
image_decoded = tf.cast(image_raw, dtype=tf.float32)
filename = parsed_features["filename"]
#Get the coco image id as we need to use COCO API to evaluate
image_id = tf.strings.to_number(tf.strings.substr(filename, 0, 12), tf.int32)
image_id = tf.expand_dims(image_id, 0)
#Get the bbox
xmin = tf.cast(tf.expand_dims(parsed_features["bbox_xmin"].values, 0), tf.float32)
xmax = tf.cast(tf.expand_dims(parsed_features["bbox_xmax"].values, 0), tf.float32)
ymin = tf.cast(tf.expand_dims(parsed_features["bbox_ymin"].values, 0), tf.float32)
ymax = tf.cast(tf.expand_dims(parsed_features["bbox_ymax"].values, 0), tf.float32)
mixup_w = tf.ones_like(xmin)
boxes = tf.concat([xmin,ymin,xmax,ymax,label,mixup_w], axis=0)
boxes = tf.transpose(boxes, [1, 0])
return {'image':image_decoded, 'bbox':boxes, 'imageid':image_id}定义一个Flatmap函数，每次读取两张图片，通过Flatmap函数来把这两张图片组合在一起
def _flatmap_function(feature):
dataset_image = feature['image'].padded_batch(2, [-1,-1,3])
dataset_bbox = feature['bbox'].padded_batch(2, [-1,6])
dataset_combined = tf.data.Dataset.zip({'image':dataset_image, 'bbox':dataset_bbox})
return dataset_combinedMixup函数把组合后的两张图片进行数据增广处理，同时生成训练的Label
def _label_fn(bbox):
global image_size,grid_wh_array,anchors_base
grids_list = [image_size//8, image_size//16, image_size//32]
image_ratio = image_size/416
anchors = [round(a*image_ratio) for a in anchors_base]
labels_list = [np.zeros([a**2,label_vector_size]) for a in grids_list]
for i in range(3):
labels_list[i][:,0] = np.arange(grids_list[i]**2)
labels_list = [np.tile(a,3) for a in labels_list]
box_num, _ = bbox.shape
for i in range(box_num):
center_x = (bbox[i,0]+bbox[i,2])/2
center_y = (bbox[i,1]+bbox[i,3])/2
if (center_x==0 and center_y==0):
continue
box_width = bbox[i,2]-bbox[i,0]
box_height = bbox[i,3]-bbox[i,1]
label = np.int(bbox[i,4].numpy())
anchor_id = np.int(bbox[i,5].numpy())
featuremap_id = anchor_id//3
anchorid_offset = anchor_id%3
g_h = grid_wh_array[featuremap_id,1]
g_w = grid_wh_array[featuremap_id,0]
grid_id = np.int((center_y//g_h*grids_list[featuremap_id] + center_x//g_w).numpy())
index = anchorid_offset*label_vector_size
#set the object exist flag
labels_list[featuremap_id][grid_id, index+1] = 1.
#set the center_x_offset
labels_list[featuremap_id][grid_id, index+2]=(center_x%g_w)/g_w
#set the center_y_offset
labels_list[featuremap_id][grid_id, index+3]=(center_y%g_h)/g_h
#set the width
labels_list[featuremap_id][grid_id, index+4]=math.log(box_width/anchors[2*anchor_id])
#set the height
labels_list[featuremap_id][grid_id, index+5]=math.log(box_height/anchors[2*anchor_id+1])
#set the mixup weight
labels_list[featuremap_id][grid_id, index+6]=bbox[i,6]
#set the class label, using label smoothing
labels_list[featuremap_id][grid_id, (index+7):(index+label_vector_size)]=0.1/(category_num-1)
labels_list[featuremap_id][grid_id, index+7+label]=0.9
#labels_list[featuremap_id][grid_id, index+7+label]=1.0
return tf.concat(labels_list, axis=0)
def _mixup_function(features):
global anchors_base,image_size,mixup_flag,grid_size
image_ratio = image_size/416
anchors = [round(a*image_ratio) for a in anchors_base]
image_height = image_size
image_width = image_size
images = features['image']
bboxes = features['bbox']
#imageid = features['imageid']
if mixup_flag:
lam = np.random.beta(1.5,1.5,1)
lam_all = np.vstack([lam,1.-lam])
lam_all = np.expand_dims(lam_all, 1)
#bboxes = tf.cast(bboxes, tf.float32)
mixup_w = bboxes[...,-1:] + lam_all
bboxes_mixup = tf.concat([bboxes[...,:-1], mixup_w], axis=-1)
bboxes_mixup = tf.reshape(bboxes_mixup, [-1,6])
true_box_mask = tf.logical_or(
bboxes_mixup[:,1]>0,
bboxes_mixup[:,1]>0
)
bboxes_all = tf.boolean_mask(bboxes_mixup, true_box_mask)
image_mix = (images[0]*lam[0] + images[1]*(1.-lam[0]))
else:
image_mix = images
bboxes_all = bboxes
#Random jitter and resize the image
height = tf.shape(image_mix)[0]
width = tf.shape(image_mix)[1]
dw = jitter*tf.cast(width, tf.float32)
dh = jitter*tf.cast(height, tf.float32)
new_ar = tf.truediv(
tf.add(
tf.cast(width, tf.float32),
tf.random.uniform([1], minval=tf.math.negative(dw), maxval=dw)),
tf.add(
tf.cast(height, tf.float32),
tf.random.uniform([1], minval=tf.math.negative(dh), maxval=dh)))
nh, nw = tf.cond(
tf.less(new_ar[0],1), \
lambda:(image_height, tf.cast(tf.cast(image_height, tf.float32)*new_ar[0], tf.int32)), \
lambda:(tf.cast(tf.cast(image_width, tf.float32)/new_ar[0], tf.int32), image_width)
)
dx = tf.cond(
tf.equal(image_width, nw), \
lambda:tf.constant([0]), \
lambda:tf.random.uniform([1], minval=0, maxval=(image_width-nw), dtype=tf.int32)
)
dy = tf.cond(
tf.equal(image_height, nh), \
lambda:tf.constant([0]), \
lambda:tf.random.uniform([1], minval=0, maxval=(image_height-nh), dtype=tf.int32)
)
image_resize = tf.image.resize(image_mix, [nh, nw])
image_padded = tf.image.pad_to_bounding_box(image_resize, dy[0], dx[0], image_height, image_width)
#Adjust the boxes
xmin_new = tf.cast(tf.truediv(nw, width) * tf.cast(bboxes_all[:,0:1],tf.float64), tf.int32) + dx
xmax_new = tf.cast(tf.truediv(nw, width) * tf.cast(bboxes_all[:,2:3],tf.float64), tf.int32) + dx
ymin_new = tf.cast(tf.truediv(nh, height) * tf.cast(bboxes_all[:,1:2],tf.float64), tf.int32) + dy
ymax_new = tf.cast(tf.truediv(nh, height) * tf.cast(bboxes_all[:,3:4],tf.float64), tf.int32) + dy
# Random flip flag
random_flip_flag = tf.random.uniform([1], minval=0, maxval=1, dtype=tf.float32)
def flip_box():
xmax_flip = image_width - xmin_new
xmin_flip = image_width - xmax_new
image_flip = tf.image.flip_left_right(image_padded)
return xmin_flip, xmax_flip, image_flip
def notflip():
return xmin_new, xmax_new, image_padded
xmin_flip, xmax_flip, image_flip = tf.cond(tf.less(random_flip_flag[0], 0.5), notflip, flip_box)
boxes_width = xmax_flip-xmin_flip
boxes_height = ymax_new-ymin_new
boxes_area = boxes_width*boxes_height
# Determine the anchor
iou_list = []
for i in range(9):
intersect_area = tf.minimum(boxes_width, anchors[2*i])*tf.minimum(boxes_height, anchors[2*i+1])
union_area = boxes_area+anchors[2*i]*anchors[2*i+1]-intersect_area
iou_list.append(intersect_area/union_area)
iou = tf.concat(iou_list, axis=1)
anchor_id = tf.reshape(tf.argmax(iou, axis=1), [-1,1])
# Random distort the image
distorted = tf.image.random_hue(image_flip, max_delta=0.3)
distorted = tf.image.random_saturation(distorted, lower=0.6, upper=1.4)
distorted = tf.image.random_brightness(distorted, max_delta=0.3)
# Add PCA noice
alpha = tf.random.normal([3], mean=0.0, stddev=0.1)
pca_noice = tf.reshape(tf.matmul(tf.multiply(eigvec,alpha), eigval), [3])
distorted = tf.add(distorted, pca_noice)
# Normalize RGB
distorted = tf.subtract(distorted, mean_RGB)
distorted = tf.divide(distorted, std_RGB)
# Get the adjusted boxes
xmin_flip = tf.cast(xmin_flip, tf.float32)
xmax_flip = tf.cast(xmax_flip, tf.float32)
ymin_new = tf.cast(ymin_new, tf.float32)
ymax_new = tf.cast(ymax_new, tf.float32)
anchor_id = tf.cast(anchor_id, tf.float32)
boxes_new = tf.concat([xmin_flip,ymin_new,xmax_flip,ymax_new,bboxes_all[:,4:5],anchor_id,bboxes_all[:,-1:]], axis=1)
# Remove the boxes that height or width less than 5 pixels
boxes_mask = tf.math.logical_and(
tf.math.greater((boxes_new[:,2]-boxes_new[:,0]), 5),
tf.math.greater((boxes_new[:,3]-boxes_new[:,1]), 5))
boxes_new = tf.boolean_mask(boxes_new, boxes_mask)
boxes_new = tf.cast(boxes_new, tf.float32)
# Generate the labels
labels = tf.py_function(_label_fn, [boxes_new], [tf.float64])
labels = tf.cast(labels, tf.float32)
image_train = tf.transpose(distorted, perm=[2, 0, 1])
#features = {'images':image_train, 'bboxes':boxes_new, 'images_flip':image_flip, 'image_id':imageid}
features = {'images':image_train, 'bboxes':boxes_new}
return features, labels[0]然后就可以构造训练的数据集了
def train_input_fn():
global image_size
train_files = tf.data.Dataset.list_files("../dataset/coco/train2017_tf/*.tfrecord")
dataset_train = train_files.interleave(tf.data.TFRecordDataset, num_parallel_calls=tf.data.experimental.AUTOTUNE)
dataset_train = dataset_train.shuffle(buffer_size=1000, reshuffle_each_iteration=True)
dataset_train = dataset_train.repeat(8)
dataset_train = dataset_train.map(_parse_function, num_parallel_calls=tf.data.experimental.AUTOTUNE)
if mixup_flag:
dataset_train = dataset_train.window(2)
dataset_train = dataset_train.flat_map(_flatmap_function)
dataset_train = dataset_train.map(_mixup_function, num_parallel_calls=tf.data.experimental.AUTOTUNE)
dataset_train = dataset_train.padded_batch(batch_size, \
padded_shapes=(
{
'images':[3,image_size,image_size],
'bboxes':[None,7]
},
[None, label_vector_size*3]
)
)
dataset_train = dataset_train.prefetch(tf.data.experimental.AUTOTUNE)
return dataset_train数据增广后的处理效果可见下图：
4. 定义损失函数
这是整个训练中最具挑战性的部分。因为按照YOLO V3的源代码，损失函数由两部分组成：
- 没有对应检测物体的网格的损失函数，因为这部分网格没有对应的物体，只要计算其预测的物体存在的概率值与Label值的方差，对于其预测的物体的检测框的位置以及物体类别的概率不作惩罚。但是论文中也提到，这些网格虽然不负责预测，但是如果其预测的检测框与真实的检测框之间的IOU大于某个阈值(0.5)时，应忽略惩罚其预测物体存在概率与Label值的方差。由于每张图片的真实物体的检测框的数量不确定，因此如何计算每个真实物体检测框与这些网格的预测框的IOU是一个问题。这里我是采用广播的方式来进行匹配计算。我传入Feature的真实物体的BBOX的维度是[batch, V, 4]，其中V代表不定长度，取这个Batch中的最大值，4表示BBOX包括了xmin,ymin,xmax,ymax。例如这个Batch中，某一张图片拥有最多的物体检测框(12个), 那么V=12，其他图片的BBOX也填充为12个。然后把BBOX的维度扩展为[batch, 1, V, 6]，预测值计算出来的BBOX维度为[batch, 52**2+26**2+13**2, 4*3]，第2个维度是三种不同大小的网格的总数量，第三个维度是每种大小的网格预测3个BBOX。把预测的BBOX和真实BBOX进行IOU的计算，然后取最大值，并判断是否超过阈值，如超过则不惩罚其预测概率。这里的损失函数采用的是交叉熵。
- 对应预测物体的网格的损失函数。这部分比较简单，对于物体的中心点坐标，只要直接计算预测值（对中心点坐标需要先进行sigmoid函数激活）与Label的交叉熵，对于物体的宽高计算预测值与Label的方差，对于物体的存在概率，以及物体所属类别的概率，需要用交叉熵来计算。另外对于不同部分的方差还要与不同的系数进行相乘。
具体的代码如下：
# Predicts, combination of three dimention, [batch, 52*52+26*26+13*13, 85*3]
# Labels, combination of three dimention, [batch, 52*52+26*26+13*13, 87*3]
def new_loss_func(predict, label, gt_box, grids_property):
global image_size
predict = tf.reshape(predict, [batch_size,-1,vector_size]) #[batch, (52*52+26*26+13*13)*3, 85]
label = tf.reshape(label, [batch_size,-1,label_vector_size]) #[batch, (52*52+26*26+13*13)*3, 87]
noobj_mask = tf.cast(label[...,1:2]==0.0, tf.float32)
obj_mask = tf.cast(label[...,1:2]==1.0, tf.float32)
#Get the predict box center xy
predict_xy = (grids_property[...,0:2]+tf.nn.sigmoid(predict[...,1:3]))*grids_property[...,-2:]
#Get the predict box wh, only caluculate the noobj wh
predict_half_wh = tf.exp(predict[...,3:5])*grids_property[...,2:4]/2
predict_xmin = tf.clip_by_value((predict_xy[...,0:1]-predict_half_wh[...,0:1]), 0, image_size)
predict_xmax = tf.clip_by_value((predict_xy[...,0:1]+predict_half_wh[...,0:1]), 0, image_size)
predict_ymin = tf.clip_by_value((predict_xy[...,1:2]-predict_half_wh[...,1:2]), 0, image_size)
predict_ymax = tf.clip_by_value((predict_xy[...,1:2]+predict_half_wh[...,1:2]), 0, image_size)
predict_boxes_area = (predict_xmax-predict_xmin)*(predict_ymax-predict_ymin) #[-batch, (52*52+26*26+13*13)*3, 1]
#Assemble the predict box coords and expand dim, shape: [batch, (52*52+26*26+13*13)*3, 1, 4]
predict_boxes = tf.concat([predict_xmin,predict_ymin,predict_xmax,predict_ymax], axis=-1)
predict_boxes = tf.expand_dims(predict_boxes, 2)
#Expand ground boxes dim for broadcast, shape: [batch, 1, V, 4]
gt_box = tf.expand_dims(gt_box, 1)
gt_box = tf.cast(gt_box, tf.float32)
#gt_box_area = (gt_box[...,2:3]-gt_box[...,0:1])*(gt_box[...,3:4]-gt_box[...,1:2]) #[batch, 1, V, 1]
gt_box_area = (gt_box[...,2]-gt_box[...,0])*(gt_box[...,3]-gt_box[...,1]) #[batch, 1, V]
#Broadcast calculation, intersect_boxes_width shape [batch, noobjs_num, V, 1]
intersect_boxes_width = tf.minimum(predict_boxes[...,2:3], gt_box[...,2:3])-tf.maximum(predict_boxes[...,0:1], gt_box[...,0:1])
intersect_boxes_width = tf.clip_by_value(intersect_boxes_width, clip_value_min=0, clip_value_max=image_size)
intersect_boxes_height = tf.minimum(predict_boxes[...,3:4], gt_box[...,3:4])-tf.maximum(predict_boxes[...,1:2], gt_box[...,1:2])
intersect_boxes_height = tf.clip_by_value(intersect_boxes_height, clip_value_min=0, clip_value_max=image_size)
intersect_boxes_area = intersect_boxes_width * intersect_boxes_height # [batch, (52*52+26*26+13*13)*3, V, 1]
intersect_boxes_area = tf.squeeze(intersect_boxes_area) # [batch, (52*52+26*26+13*13)*3, V]
#Calculate the noobj predict box IOU with ground truth boxes, shape:[batch, (52*52+26*26+13*13)*3, V]
iou_boxes = intersect_boxes_area/(predict_boxes_area+gt_box_area-intersect_boxes_area) #
iou_max = tf.reduce_max(iou_boxes, axis=2, keepdims=True) #[batch, (52*52+26*26+13*13)*3, 1]
#iou_max = tf.expand_dims(iou_max, 2)
#Ignore the noobj loss for the IOU larger than threshold
no_ignore_mask = tf.cast(iou_max[...,0:1]<iou_threshold, tf.float32)
noobj_loss_mask = tf.cast(noobj_mask*no_ignore_mask, tf.bool)
noobj_loss_mask = tf.reshape(noobj_loss_mask, [batch_size, -1])
noobj_predict = tf.boolean_mask(predict, noobj_loss_mask)[...,0:1]
#noobj_predict_topk, _ = tf.nn.top_k(noobj_predict[...,0], k=500)
#Calculate the noobj predict loss
#loss_noobj = tf.reduce_sum(tf.nn.sigmoid_cross_entropy_with_logits(tf.zeros_like(noobj_predict_topk), noobj_predict_topk))
loss_noobj = tf.reduce_sum(
tf.nn.sigmoid_cross_entropy_with_logits(tf.zeros_like(noobj_predict), noobj_predict))
loss_noobj = loss_noobj/batch_size
#Calculate the scale for object coord loss
coord_scale = 2.0-\
(tf.exp(label[...,4:5])*grids_property[...,2:3])*\
(tf.exp(label[...,5:6])*grids_property[...,3:4])/\
image_size**2
loss_xy = tf.reduce_sum(
0.5*coord_scale*label[...,6:7]*obj_mask*
tf.nn.sigmoid_cross_entropy_with_logits(
labels=label[...,2:4],
logits=predict[...,1:3]
)
)
loss_wh = tf.reduce_sum(
0.5*coord_scale*label[...,6:7]*obj_mask*
tf.square(
label[...,4:6]-
predict[...,3:5]
)
)
#Calculate the conf loss
loss_conf = tf.reduce_sum(
object_scale*label[...,6:7]*obj_mask*
tf.nn.sigmoid_cross_entropy_with_logits(
tf.ones_like(predict[...,0:1]),
predict[...,0:1]
)
)
#Calculate the predict class loss
loss_class = tf.reduce_sum(
class_scale*label[...,6:7]*obj_mask*
tf.nn.sigmoid_cross_entropy_with_logits(
labels=label[...,7:],
logits=predict[...,5:]
)
)
loss_obj = (loss_xy+loss_wh+loss_conf+loss_class)/batch_size
final_loss = loss_obj + loss_noobj
return final_loss
tf_new_loss_func = tf.function(new_loss_func, experimental_relax_shapes=True)5. 模型的训练
模型的训练过程，我是采用了自定义训练的方式来做的。YOLO论文提到训练时可以随机采用多种图片尺度，例如416*416, 608*608，352*352等，这样的好处是模型能够更好的适应不同尺寸大小的图片的检测。
随机采用多种图片尺度的代码如下：
def random_image():
global image_size_list,image_size
global grid_wh_array
global anchors_base
image_size = image_size_list[random.randint(0,9)]
#image_size = 608
image_ratio = image_size/416
grids_list = [image_size//8, image_size//16, image_size//32]
anchors = [round(a*image_ratio) for a in anchors_base]
grids_x_list = [np.reshape(np.arange(a**2)%a,[-1,1]) for a in grids_list]
grids_x = np.vstack(grids_x_list)
grids_x = np.reshape(np.hstack([grids_x,grids_x,grids_x]),[-1,1])
grids_y_list = [np.reshape(np.arange(a**2)//a,[-1,1]) for a in grids_list]
grids_y = np.vstack(grids_y_list)
grids_y = np.reshape(np.hstack([grids_y,grids_y,grids_y]),[-1,1])
anchors_all = np.vstack(
[
np.reshape(np.tile(np.reshape(np.array(anchors[:6]),[-1,6]),[grids_list[0]**2,1]),[-1,2]),
np.reshape(np.tile(np.reshape(np.array(anchors[6:12]),[-1,6]),[grids_list[1]**2,1]),[-1,2]),
np.reshape(np.tile(np.reshape(np.array(anchors[12:]),[-1,6]),[grids_list[2]**2,1]),[-1,2])
]
)
grid_wh_all = np.vstack(
[
np.tile(grid_wh_array[:1,:], (grids_list[0]**2*3,1)),
np.tile(grid_wh_array[1:2,:], (grids_list[1]**2*3,1)),
np.tile(grid_wh_array[2:3,:], (grids_list[2]**2*3,1))
]
)
grids_property = np.concatenate([grids_x, grids_y, anchors_all, grid_wh_all], axis=-1)
grids_property_all = tf.constant(grids_property, dtype=tf.float32)
grids_property_all = tf.expand_dims(grids_property_all, 0)
grids_property_all = tf.tile(grids_property_all, [batch_size,1,1])
return grids_property_all自定义训练过程的代码如下：
model_base = tf.keras.models.load_model('darknet53_20200228/epoch_42.h5')
model_base.trainable = False
image = tf.keras.Input(shape=(3,None,None))
_, route1, route2, route3 = model_base(image, training=False)
predict = darknet53_yolov3()([route1, route2, route3])
model_yolo = tf.keras.Model(inputs=image, outputs=predict, name='model_yolo')
START_EPOCH = 0
NUM_EPOCH = 1
STEPS_EPOCH = epoch_batch
STEPS_OFFSET = STEPS_EPOCH*START_EPOCH
initial_warmup_steps = 1000
initial_lr = 0.0005
optimizer=tf.keras.optimizers.SGD(learning_rate=0.00001, momentum=0.9)
mp_opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(optimizer)
def train_step(images, bbox, labels, grids_property_all):
with tf.GradientTape() as tape:
predict = model_yolo(images, training=True)
regularization_loss = tf.math.add_n(model_yolo.losses)
pred_loss = tf_new_loss_func(predict, labels, bbox, grids_property_all)
total_loss = pred_loss + regularization_loss
gradients = tape.gradient(total_loss, model_yolo.trainable_variables)
mp_opt.apply_gradients(zip(gradients, model_yolo.trainable_variables))
return total_loss, predict
tf_train_step = tf.function(train_step, experimental_relax_shapes=True)
#Loss rate step decay
boundaries = [STEPS_EPOCH*4, STEPS_EPOCH*10, STEPS_EPOCH*13, STEPS_EPOCH*16]
values = [0.0005, 0.0001, 0.00005, 0.00001, 0.00005]
learning_rate_fn = tf.keras.optimizers.schedules.PiecewiseConstantDecay(boundaries, values)
steps = STEPS_OFFSET
for epoch in range(NUM_EPOCH):
loss_sum = 0
start_time = time.time()
grids_property_all = new_random_image()
train_data = iter(train_input_fn())
#for features, labels in train_data:
while(True):
if steps < initial_warmup_steps:
newlr = (initial_lr/initial_warmup_steps)*steps
tf.keras.backend.set_value(optimizer.lr, newlr)
features, data_labels = train_data.next()
loss_temp, predict_temp = tf_train_step(features['images'], features['bboxes'], data_labels, grids_property_all)
loss_sum += loss_temp
steps += 1
if steps%100 == 0:
elasp_time = time.time()-start_time
lr = tf.keras.backend.get_value(optimizer.lr)
print("Step:{}, Image_size:{:d}, Loss:{:4.2f}, LR:{:5f}, Time:{:3.1f}s".format(steps, image_size, loss_sum/100, lr, elasp_time))
loss_sum = 0
if steps > initial_warmup_steps:
tf.keras.backend.set_value(optimizer.lr, learning_rate_fn(steps))
start_time = time.time()
if steps%STEPS_EPOCH == 0:
START_EPOCH += 1
model_yolo.save('model_yolov3/yolo_v10_'+str(START_EPOCH)+'.h5')
break模型的训练非常耗时，在我的电脑（2080Ti）的配置下，训练一个Epoch大概要花2个小时，我训练了14个EPOCH，mAP .50的准确度大概为32%，和论文提到的57.9%还有比较大的差距。不过Darknet的源码是训练了200多个EPOCH的，可能继续训练会进一步提高准确度。这个有待以后继续验证。
6. 评价模型的性能指标
目标检测一般采用mAP来评价性能，这个指标的计算比较复杂，我是直接采用了COCO API来进行计算，这个也是和论文中的计算方法保持一致。
首先是构造COCO测试集，代码如下：
def _parse_val_function(example_proto):
global val_image_size
features = {
"image": tf.io.FixedLenFeature([], tf.string, default_value=""),
"height": tf.io.FixedLenFeature([1], tf.int64, default_value=[0]),
"width": tf.io.FixedLenFeature([1], tf.int64, default_value=[0]),
"channels": tf.io.FixedLenFeature([1], tf.int64, default_value=[3]),
"colorspace": tf.io.FixedLenFeature([], tf.string, default_value=""),
"img_format": tf.io.FixedLenFeature([], tf.string, default_value=""),
"label": tf.io.VarLenFeature(tf.int64),
"bbox_xmin": tf.io.VarLenFeature(tf.int64),
"bbox_xmax": tf.io.VarLenFeature(tf.int64),
"bbox_ymin": tf.io.VarLenFeature(tf.int64),
"bbox_ymax": tf.io.VarLenFeature(tf.int64),
"filename": tf.io.FixedLenFeature([], tf.string, default_value="")
}
parsed_features = tf.io.parse_single_example(example_proto, features)
label = tf.expand_dims(parsed_features["label"].values, 0)
label = tf.cast(label, tf.int32)
channels = parsed_features["channels"]
filename = parsed_features["filename"]
#Get the coco image id as we need to use COCO API to evaluate
image_id = tf.strings.to_number(
tf.strings.substr(filename, 0, 12),
tf.int32
)
#Decode the image
image_raw = tf.image.decode_jpeg(parsed_features["image"], channels=3)
image_decoded = tf.cast(image_raw, dtype=tf.float32)
image_h = tf.constant(val_image_size)
image_w = tf.constant(val_image_size)
height = tf.shape(image_decoded)[0]
width = tf.shape(image_decoded)[1]
original_size = tf.stack([height, width], axis=0)
original_size = tf.cast(original_size, tf.float32)
ratio = tf.truediv(tf.cast(height, tf.float32), tf.cast(width, tf.float32))
nh, nw = tf.cond(
tf.less(ratio,1),
lambda:(tf.cast(tf.cast(image_h, tf.float32)*ratio, tf.int32), image_w),
lambda:(image_h, tf.cast(tf.cast(image_w, tf.float32)/ratio, tf.int32)))
dx = tf.cond(
tf.equal(image_w, nw), \
lambda:tf.constant(0), \
lambda:tf.cast((image_w-nw)/2, tf.int32))
dy = tf.cond(
tf.equal(image_h, nh), \
lambda:0, \
lambda:tf.cast((image_h-nh)/2, tf.int32))
image_resize = tf.image.resize(image_decoded, [nh, nw])
image_padded = tf.image.pad_to_bounding_box(image_resize, dy, dx, image_h, image_w)
image_normalize = tf.subtract(image_padded, mean_RGB)
image_normalize = tf.divide(image_normalize, std_RGB)
image_val = tf.transpose(image_normalize, perm=[2, 0, 1])
features = {'images':image_val, 'image_id':image_id, 'original_size':original_size}
return features
def val_input_fn():
val_files = tf.data.Dataset.list_files("../dataset/coco/val2017_tf/*.tfrecord")
dataset_val = val_files.interleave(tf.data.TFRecordDataset, cycle_length=12, num_parallel_calls=12)
dataset_val = dataset_val.map(_parse_val_function, num_parallel_calls=12)
dataset_val = dataset_val.batch(val_batch_size)
dataset_val = dataset_val.prefetch(1)
return dataset_val解码预测的结果，转换为相应的BBOX，如以下代码：
def predict_func(predict, image_id, original_size):
global val_image_size, anchors_base
val_grids_list = [val_image_size//8, val_image_size//16, val_image_size//32]
image_ratio = val_image_size/416
val_anchors = [round(a*image_ratio) for a in anchors_base]
val_grids_x_list = [np.reshape(np.arange(a**2)%a,[-1,1]) for a in val_grids_list]
val_grids_x = np.vstack(val_grids_x_list)
val_grids_y_list = [np.reshape(np.arange(a**2)//a,[-1,1]) for a in val_grids_list]
val_grids_y = np.vstack(val_grids_y_list)
val_anchors_all = np.vstack(
[
np.tile(np.reshape(np.array(val_anchors[:6]),[-1,6]),[val_grids_list[0]**2,1]),
np.tile(np.reshape(np.array(val_anchors[6:12]),[-1,6]),[val_grids_list[1]**2,1]),
np.tile(np.reshape(np.array(val_anchors[12:]),[-1,6]),[val_grids_list[2]**2,1])
]
)
grid_wh_all = np.vstack(
[
np.tile(grid_wh_array[:1,:], (val_grids_list[0]**2,1)),
np.tile(grid_wh_array[1:2,:], (val_grids_list[1]**2,1)),
np.tile(grid_wh_array[2:3,:], (val_grids_list[2]**2,1))
]
)
val_grids_property = np.concatenate([val_grids_x, val_grids_y, val_anchors_all, grid_wh_all], axis=-1)
val_grids_property_all = tf.constant(val_grids_property, dtype=tf.float32)
val_grids_property_all = tf.expand_dims(val_grids_property_all, 0)
val_grids_property_all = tf.tile(val_grids_property_all, [predict.shape[0],1,1])
result_json = []
original_height = original_size[...,0]
original_width = original_size[...,1]
hw_ratio = original_height/original_width
hw_ratio_mask = tf.cast(tf.less(hw_ratio, 1.), tf.float32)
ratio = \
hw_ratio_mask*(original_width/val_image_size) + \
(1.-hw_ratio_mask)*(original_height/val_image_size)
dx = (1.-hw_ratio_mask)*((original_height-original_width)//2)
dy = hw_ratio_mask*((original_width-original_height)//2)
confidence_threshold = 0.2
probabilty_threshold = 0.5
predict_boxes_list = []
for i in range(3):
predict_conf = tf.nn.sigmoid(predict[...,i*vector_size:(i*vector_size+1)])
predict_xy = tf.nn.sigmoid(predict[...,(i*vector_size+1):(i*vector_size+3)])
predict_xy = predict_xy + val_grids_property_all[...,0:2]
predict_x = predict_xy[...,0:1] * val_grids_property_all[...,-2:-1]
predict_y = predict_xy[...,1:] * val_grids_property_all[...,-1:]
predict_w = tf.exp(predict[...,(i*vector_size+3):(i*vector_size+4)])
predict_w = predict_w * val_grids_property_all[...,(2+i*2):(2+i*2+1)]
predict_h = tf.exp(predict[...,(i*vector_size+4):(i*vector_size+5)])
predict_h = predict_h * val_grids_property_all[...,(2+i*2+1):(2+i*2+2)]
min_x = tf.clip_by_value((predict_x-predict_w/2), 0, val_image_size)
max_x = tf.clip_by_value((predict_x + predict_w/2), 0, val_image_size)
min_y = tf.clip_by_value((predict_y - predict_h/2), 0, val_image_size)
max_y = tf.clip_by_value((predict_y + predict_h/2), 0, val_image_size)
predict_class = tf.argmax(predict[...,(i*vector_size+5):((i+1)*vector_size)], axis=-1)
predict_class = tf.cast(predict_class, tf.float32)
predict_class = tf.expand_dims(predict_class, 2)
predict_proba = tf.nn.sigmoid(
tf.reduce_max(
predict[...,(i*vector_size+5):((i+1)*vector_size)], axis=-1, keepdims=True
)
)
predict_box = tf.concat([predict_conf, min_x, min_y, max_x, max_y, predict_class, predict_proba], axis=-1)
predict_boxes_list.append(predict_box)
predict_boxes = tf.concat(predict_boxes_list, axis=1)
for i in range(predict.shape[0]):
obj_mask = tf.logical_and(
predict_boxes[i,:,0]>=confidence_threshold,
predict_boxes[i,:,-1]>=probabilty_threshold)
predict_true_box = tf.boolean_mask(predict_boxes[i], obj_mask)
predict_classes, _ = tf.unique(predict_true_box[:,5])
predict_classes_list = tf.unstack(predict_classes)
for class_id in predict_classes_list:
class_mask = tf.math.equal(predict_true_box[:, 5], class_id)
predict_true_box_class = tf.boolean_mask(predict_true_box, class_mask)
predict_true_box_xy = predict_true_box_class[:, 1:5]
predict_true_box_score = predict_true_box_class[:, 6]*predict_true_box_class[:, 0]
#predict_true_box_score = predict_true_box_class[:, 0]
selected_indices = tf.image.non_max_suppression(
predict_true_box_xy,
predict_true_box_score,
100,
iou_threshold=0.2
#score_threshold=confidence_threshold
)
#Shape [box_num, 7]
selected_boxes = tf.gather(predict_true_box_class, selected_indices)
original_bbox_xmin = tf.clip_by_value(
selected_boxes[:,1:2]*ratio[i]-dx[i], 0, original_width[i])
original_bbox_xmax = tf.clip_by_value(
selected_boxes[:,3:4]*ratio[i]-dx[i], 0, original_width[i])
original_bbox_ymin = tf.clip_by_value(
selected_boxes[:,2:3]*ratio[i]-dy[i], 0, original_height[i])
original_bbox_ymax = tf.clip_by_value(
selected_boxes[:,4:5]*ratio[i]-dy[i], 0, original_height[i])
original_bbox_width = original_bbox_xmax - original_bbox_xmin
original_bbox_height = original_bbox_ymax - original_bbox_ymin
original_bbox = tf.concat(
[
selected_boxes[:,0:1],
original_bbox_xmin,
original_bbox_ymin,
original_bbox_width,
original_bbox_height,
selected_boxes[:,5:]
], axis=-1
)
original_bbox_list = tf.unstack(original_bbox)
for item in original_bbox_list:
result = {}
result['image_id'] = int(image_id.numpy()[i])
result['category_id'] = cocoid_mapping_labels[int(class_id.numpy())]
result['bbox'] = item[1:5].numpy().tolist()
result['bbox'] = [int(a*10)/10 for a in result['bbox']]
result['score'] = int((item[0]*item[6]).numpy()*1000)/1000
result['conf'] = str(int(item[0].numpy()*1000)/1000)
result['prop'] = str(int(item[6].numpy()*1000)/1000)
result_json.append(result)
return result_json利用COCO API来计算mAP：
START_EPOCH = 14
val_image_size = 608
dataset_val = val_input_fn()
all_result_json = []
i = 0
for val_features in dataset_val:
predict = model_yolo(val_features['images'], training=False)
result_json = predict_func(
predict, val_features['image_id'], val_features['original_size']
)
all_result_json.extend(result_json)
i +=1
all_result_str = ','.join([json.dumps(item) for item in all_result_json])
all_result_str = '['+all_result_str+']'
result_filename = 'test_v11_epoch_'+str(START_EPOCH)+'_result.json'
result_file = open(result_filename, 'w')
result_file.write(all_result_str)
result_file.close()
cocodt = coco.loadRes(result_filename)
annType = 'bbox'
imgIds=sorted(coco.getImgIds())
cocoEval = COCOeval(coco,cocodt,annType)
cocoEval.params.imgIds = imgIds
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()模型的预测效果
最后我们来看一下模型的预测效果如何，首先是Kite图片：
再看看Darknet YOLOV3官方模型的预测结果：
看样子我的模型的预测效果还好一些，出乎意料：），官方模型有几个风筝没有检测到。不过我的模型则错误的把一个浪花检测为人。总体好像还是我的模型预测的准确一些。
再看看另外一张图片dog，以下是我的检测效果：
官方模型的检测结果：
总体来看检测结果基本一致，官方模型的IOU检测的更精确一些。