【深度学习实战03】——YOLOtensorflow运行及源码解析

import time
import datetimeclass Timer(object):&＃39;&＃39;&＃39;A simple timer.&＃39;&＃39;&＃39;def __init__(self):self.init_time &＃61; time.time()self.total_time &＃61; 0.self.calls &＃61; 0self.start_time &＃61; 0.self.diff &＃61; 0.self.average_time &＃61; 0.self.remain_time &＃61; 0.def tic(self):# using time.time instead of time.clock because time time.clock# does not normalize for multithreadingself.start_time &＃61; time.time() #获取当前系统时间def toc(self, average&＃61;True): self.diff &＃61; time.time() - self.start_time #获取当前系统时间-之前获取的系统时间&＃61;时间差self.total_time &＃43;&＃61; self.diff #获取总的时间差self.calls &＃43;&＃61; 1 #调用次数self.average_time &＃61; self.total_time / self.calls #多次时间调用&＃xff0c;计算平均时间差if average:return self.average_timeelse:return self.diffdef remain(self, iters, max_iters): #用于计算完成剩余迭代次数预计所费时间if iters &＃61;&＃61; 0:self.remain_time &＃61; 0else:self.remain_time &＃61; (time.time() - self.init_time) * \(max_iters - iters) / itersreturn str(datetime.timedelta(seconds&＃61;int(self.remain_time)))


import os#
# path and dataset parameter
#DATA_PATH &＃61; &＃39;data&＃39;PASCAL_PATH &＃61; os.path.join(DATA_PATH, &＃39;pascal_voc&＃39;) #pascal的路径是;当前工作路径/data/pascal_vocCACHE_PATH &＃61; os.path.join(PASCAL_PATH, &＃39;cache&＃39;) #cache的路径是;当前工作路径/data/pascal_voc/cacheOUTPUT_DIR &＃61; os.path.join(PASCAL_PATH, &＃39;output&＃39;) #output的路径是;当前工作路径/data/pascal_voc/outputWEIGHTS_DIR &＃61; os.path.join(PASCAL_PATH, &＃39;weights&＃39;) #weights的路径是;当前工作路径/data/pascal_voc/weightsWEIGHTS_FILE &＃61; None
# WEIGHTS_FILE &＃61; os.path.join(DATA_PATH, &＃39;weights&＃39;, &＃39;YOLO_small.ckpt&＃39;)CLASSES &＃61; [&＃39;aeroplane&＃39;, &＃39;bicycle&＃39;, &＃39;bird&＃39;, &＃39;boat&＃39;, &＃39;bottle&＃39;, &＃39;bus&＃39;, #目标类别&＃39;car&＃39;, &＃39;cat&＃39;, &＃39;chair&＃39;, &＃39;cow&＃39;, &＃39;diningtable&＃39;, &＃39;dog&＃39;, &＃39;horse&＃39;,&＃39;motorbike&＃39;, &＃39;person&＃39;, &＃39;pottedplant&＃39;, &＃39;sheep&＃39;, &＃39;sofa&＃39;,&＃39;train&＃39;, &＃39;tvmonitor&＃39;]FLIPPED &＃61; True #是否flipped#
# model parameter
#IMAGE_SIZE &＃61; 448CELL_SIZE &＃61; 7BOXES_PER_CELL &＃61; 2ALPHA &＃61; 0.1DISP_CONSOLE &＃61; FalseOBJECT_SCALE &＃61; 1.0 #这四个损失函数系数
NOOBJECT_SCALE &＃61; 1.0
CLASS_SCALE &＃61; 2.0
COORD_SCALE &＃61; 5.0#
# solver parameter
#GPU &＃61; &＃39;&＃39;LEARNING_RATE &＃61; 0.0001DECAY_STEPS &＃61; 30000DECAY_RATE &＃61; 0.1STAIRCASE &＃61; TrueBATCH_SIZE &＃61; 45MAX_ITER &＃61; 15000SUMMARY_ITER &＃61; 10SAVE_ITER &＃61; 1000#
# test parameter
#THRESHOLD &＃61; 0.2IOU_THRESHOLD &＃61; 0.5


import numpy as np
import tensorflow as tf
import yolo.config as cfgslim &＃61; tf.contrib.slimclass YOLONet(object): #定义一个YOLONet类def __init__(self, is_training&＃61;True):self.classes &＃61; cfg.CLASSES #目标类别self.num_class &＃61; len(self.classes) #目标类别数量&＃xff0c;值为20self.image_size &＃61; cfg.IMAGE_SIZE #图像尺寸&＃xff0c;为448self.cell_size &＃61; cfg.CELL_SIZE # cell尺寸&＃xff0c;为7self.boxes_per_cell &＃61; cfg.BOXES_PER_CELL #每个grid cell负责的boxes数量&＃xff0c;为2self.output_size &＃61; (self.cell_size * self.cell_size) *\ #输出特征维度&＃xff0c;7X7X(20&＃43;2X5)(self.num_class &＃43; self.boxes_per_cell * 5)self.scale &＃61; 1.0 * self.image_size / self.cell_size #尺寸缩放系数&＃xff0c; 448/7&＃61;64self.boundary1 &＃61; self.cell_size * self.cell_size * self.num_class # 7X7X20self.boundary2 &＃61; self.boundary1 &＃43;\ # 7X7X20 &＃43; 7X7X2 49个所属20个物体类别的概率&＃43;98个bboxself.cell_size * self.cell_size * self.boxes_per_cellself.object_scale &＃61; cfg.OBJECT_SCALE #值为1&＃xff0c;有目标存在的系数self.noobject_scale &＃61; cfg.NOOBJECT_SCALE #值为1&＃xff0c;没有目标存在的系数&＃xff08;论文貌似为0.5&＃xff09;self.class_scale &＃61; cfg.CLASS_SCALE #值为2.0&＃xff0c; 类别损失函数的系数self.coord_scale &＃61; cfg.COORD_SCALE #值为5.0&＃xff0c;坐标损失函数的系数self.learning_rate &＃61; cfg.LEARNING_RATE #学习率&＃61;0.0001self.batch_size &＃61; cfg.BATCH_SIZE #batch_size&＃61;45self.alpha &＃61; cfg.ALPHA #alpha&＃61;0.1self.offset &＃61; np.transpose(np.reshape(np.array([np.arange(self.cell_size)] * self.cell_size * self.boxes_per_cell), #将2X7X7的三维矩阵,转为7X7X2的三维矩阵(self.boxes_per_cell, self.cell_size, self.cell_size)), (1, 2, 0))self.images &＃61; tf.placeholder(tf.float32, [None, self.image_size, self.image_size, 3], #创建输入图像占位符 448X448 3通道 name&＃61;&＃39;images&＃39;)self.logits &＃61; self.build_network( #输出logits值&＃xff08;预测值&＃xff09;self.images, num_outputs&＃61;self.output_size, alpha&＃61;self.alpha,is_training&＃61;is_training)if is_training:self.labels &＃61; tf.placeholder(tf.float32,[None, self.cell_size, self.cell_size, 5 &＃43; self.num_class]) #为label&＃xff08;真实值&＃xff09;穿件占位符self.loss_layer(self.logits, self.labels) #求lossself.total_loss &＃61; tf.losses.get_total_loss() #求所有的losstf.summary.scalar(&＃39;total_loss&＃39;, self.total_loss)def build_network(self, #建立网络&＃xff08;卷积层&＃43;池化层&＃43;全连接层&＃xff09;images, #输入的图像 [None&＃xff0c;448,448,3]num_outputs, #输出特征维度[None,7X7X30]alpha,keep_prob&＃61;0.5, #dropoutis_training&＃61;True,scope&＃61;&＃39;yolo&＃39;): #命个名字with tf.variable_scope(scope):with slim.arg_scope([slim.conv2d, slim.fully_connected], activation_fn&＃61;leaky_relu(alpha), #激活函数用的是leaky_reluweights_regularizer&＃61;slim.l2_regularizer(0.0005), #权重正则化用的是l2weights_initializer&＃61;tf.truncated_normal_initializer(0.0, 0.01) #权重初始化用的是正态分布&＃xff08;0.0,0.01&＃xff09;):net &＃61; tf.pad( #为输入图像进行填充&＃xff0c;单张图上下左右各用0填充3行/列images, np.array([[0, 0], [3, 3], [3, 3], [0, 0]]), #BatchSize维度不填充&＃xff0c;行维度上下填充3行0&＃xff0c;列维度左右填充3列0&＃xff0c;channel维度不填充name&＃61;&＃39;pad_1&＃39;)net &＃61; slim.conv2d( # input&＃61;net; num_outputs&＃61;64个特征图;kernel_size:7X7; strides&＃61;2;net, 64, 7, 2, padding&＃61;&＃39;VALID&＃39;, scope&＃61;&＃39;conv_2&＃39;) # 上面已经pad了&＃xff0c;所以选padding&＃61;VALID&＃xff0c;即不停留在图像边缘net &＃61; slim.max_pool2d(net, 2, padding&＃61;&＃39;SAME&＃39;, scope&＃61;&＃39;pool_3&＃39;) #最大池化 2X2的核结构&＃xff0c;stride&＃61;2&＃xff1b;输出net 224X224X64net &＃61; slim.conv2d(net, 192, 3, scope&＃61;&＃39;conv_4&＃39;) #卷积&＃xff0c;输出特征图192个&＃xff0c;kernel_size:3X3; 输出net: 224X224X192net &＃61; slim.max_pool2d(net, 2, padding&＃61;&＃39;SAME&＃39;, scope&＃61;&＃39;pool_5&＃39;) #最大池化 2X2&＃xff0c; stride&＃61;2&＃xff1b; 输出net:112X112X192 OK net &＃61; slim.conv2d(net, 128, 1, scope&＃61;&＃39;conv_6&＃39;) #卷积&＃xff0c; kernel&＃61;1X1; 输出net: 112X112X128net &＃61; slim.conv2d(net, 256, 3, scope&＃61;&＃39;conv_7&＃39;) #卷积&＃xff0c; kernel&＃61;3X3&＃xff1b;输出net: 112X112X256net &＃61; slim.conv2d(net, 256, 1, scope&＃61;&＃39;conv_8&＃39;) #卷积&＃xff0c; kernel&＃61;1X1; 输出net: 112X112X256net &＃61; slim.conv2d(net, 512, 3, scope&＃61;&＃39;conv_9&＃39;) #卷积&＃xff0c; kernel&＃61;3X3&＃xff1b;输出net: 112X112X512net &＃61; slim.max_pool2d(net, 2, padding&＃61;&＃39;SAME&＃39;, scope&＃61;&＃39;pool_10&＃39;) #最大池化 2X2&＃xff0c;stride&＃61;2; 输出net: 56x56x256net &＃61; slim.conv2d(net, 256, 1, scope&＃61;&＃39;conv_11&＃39;) #连续4组 卷积输出特征数256和512的组合&＃xff1b;net &＃61; slim.conv2d(net, 512, 3, scope&＃61;&＃39;conv_12&＃39;)net &＃61; slim.conv2d(net, 256, 1, scope&＃61;&＃39;conv_13&＃39;)net &＃61; slim.conv2d(net, 512, 3, scope&＃61;&＃39;conv_14&＃39;)net &＃61; slim.conv2d(net, 256, 1, scope&＃61;&＃39;conv_15&＃39;)net &＃61; slim.conv2d(net, 512, 3, scope&＃61;&＃39;conv_16&＃39;)net &＃61; slim.conv2d(net, 256, 1, scope&＃61;&＃39;conv_17&＃39;)net &＃61; slim.conv2d(net, 512, 3, scope&＃61;&＃39;conv_18&＃39;)net &＃61; slim.conv2d(net, 512, 1, scope&＃61;&＃39;conv_19&＃39;) #卷积&＃xff0c;kernel&＃61;1X1&＃xff1b;输出net: 56x56x512net &＃61; slim.conv2d(net, 1024, 3, scope&＃61;&＃39;conv_20&＃39;) #卷积&＃xff0c;kernel&＃61;3X3; 输出net: 56x56x1024 &＃xff1f;&＃xff1f;&＃xff1f;net &＃61; slim.max_pool2d(net, 2, padding&＃61;&＃39;SAME&＃39;, scope&＃61;&＃39;pool_21&＃39;) #最大池化 2X2&＃xff0c;stride&＃61;2&＃xff1b;输出net:28x28x512 &＃xff1f;&＃xff1f;net &＃61; slim.conv2d(net, 512, 1, scope&＃61;&＃39;conv_22&＃39;) #连续两组 卷积输出特征数512和1024的组合net &＃61; slim.conv2d(net, 1024, 3, scope&＃61;&＃39;conv_23&＃39;)net &＃61; slim.conv2d(net, 512, 1, scope&＃61;&＃39;conv_24&＃39;)net &＃61; slim.conv2d(net, 1024, 3, scope&＃61;&＃39;conv_25&＃39;) net &＃61; slim.conv2d(net, 1024, 3, scope&＃61;&＃39;conv_26&＃39;) #卷积&＃xff0c;kernel&＃61;3X3;输出net:28X28X1024net &＃61; tf.pad( #对net进行填充net, np.array([[0, 0], [1, 1], [1, 1], [0, 0]]), #batch维度不填充&＃xff1b;28的行维度上下填充1行&＃xff08;值为0&＃xff09;&＃xff1b;28的列维度左右填充1列&＃xff08;值为0&＃xff09;&＃xff0c;channel维度不填充&＃xff1b;name&＃61;&＃39;pad_27&＃39;)net &＃61; slim.conv2d( net, 1024, 3, 2, padding&＃61;&＃39;VALID&＃39;, scope&＃61;&＃39;conv_28&＃39;) #上面已经pad了&＃xff0c;所以选padding&＃61;VALID&＃xff0c;kernel&＃61;3X3&＃xff0c;stride&＃61;2,输出net:14x14x1024 &＃xff1f;&＃xff1f;&＃xff1f;net &＃61; slim.conv2d(net, 1024, 3, scope&＃61;&＃39;conv_29&＃39;) #连续两个卷积&＃xff0c;特征数为1024&＃xff0c;kernel&＃61;3x3net &＃61; slim.conv2d(net, 1024, 3, scope&＃61;&＃39;conv_30&＃39;) #输出net: 7x7x1024 &＃xff1f;&＃xff1f;&＃xff1f;net &＃61; tf.transpose(net, [0, 3, 1, 2], name&＃61;&＃39;trans_31&＃39;) #输出net:[batchsize,channel,28,28]net &＃61; slim.flatten(net, scope&＃61;&＃39;flat_32&＃39;) #输出net: (1,batchsize x channel x w x h)net &＃61; slim.fully_connected(net, 512, scope&＃61;&＃39;fc_33&＃39;) #全连接层 输出net:1x512net &＃61; slim.fully_connected(net, 4096, scope&＃61;&＃39;fc_34&＃39;) #全连接层 输出net:1x4096net &＃61; slim.dropout( #dropout层&＃xff0c;防止过拟合net, keep_prob&＃61;keep_prob, is_training&＃61;is_training,scope&＃61;&＃39;dropout_35&＃39;)net &＃61; slim.fully_connected( #全连接层&＃xff0c;输出net:7x7x30特征net, num_outputs, activation_fn&＃61;None, scope&＃61;&＃39;fc_36&＃39;)return net #返回net: 7x7x30def calc_iou(self, boxes1, boxes2, scope&＃61;&＃39;iou&＃39;): #计算box和groundtruth的IOU值"""calculate iousArgs:boxes1: 5-D tensor [BATCH_SIZE, CELL_SIZE, CELL_SIZE, BOXES_PER_CELL, 4] &＃61;&＃61;&＃61;&＃61;> (x_center, y_center, w, h)boxes2: 5-D tensor [BATCH_SIZE, CELL_SIZE, CELL_SIZE, BOXES_PER_CELL, 4] &＃61;&＃61;&＃61;> (x_center, y_center, w, h)Return:iou: 4-D tensor [BATCH_SIZE, CELL_SIZE, CELL_SIZE, BOXES_PER_CELL]"""with tf.variable_scope(scope):# transform (x_center, y_center, w, h) to (x1, y1, x2, y2)boxes1_t &＃61; tf.stack([boxes1[..., 0] - boxes1[..., 2] / 2.0, #x-w/2&＃61;x1(左上)boxes1[..., 1] - boxes1[..., 3] / 2.0, #y-h/2&＃61;y1(左上)boxes1[..., 0] &＃43; boxes1[..., 2] / 2.0, #x&＃43;w/2&＃61;x2(右下)boxes1[..., 1] &＃43; boxes1[..., 3] / 2.0], #y&＃43;h/2&＃61;y2(右下)axis&＃61;-1) #替换最后那个维度boxes2_t &＃61; tf.stack([boxes2[..., 0] - boxes2[..., 2] / 2.0,boxes2[..., 1] - boxes2[..., 3] / 2.0,boxes2[..., 0] &＃43; boxes2[..., 2] / 2.0,boxes2[..., 1] &＃43; boxes2[..., 3] / 2.0],axis&＃61;-1)# calculate the left up point & right down point #计算重叠区域最左上和最右下点lu &＃61; tf.maximum(boxes1_t[..., :2], boxes2_t[..., :2])rd &＃61; tf.minimum(boxes1_t[..., 2:], boxes2_t[..., 2:])# intersectionintersection &＃61; tf.maximum(0.0, rd - lu) #重叠区域inter_square &＃61; intersection[..., 0] * intersection[..., 1] #重叠区域面积# calculate the boxs1 square and boxs2 squaresquare1 &＃61; boxes1[..., 2] * boxes1[..., 3] #box1.w * box1.hsquare2 &＃61; boxes2[..., 2] * boxes2[..., 3] #box2.w * box2.hunion_square &＃61; tf.maximum(square1 &＃43; square2 - inter_square, 1e-10) return tf.clip_by_value(inter_square / union_square, 0.0, 1.0) #将IOU计算得到的值归一化到&＃xff08;0,1&＃xff09;def loss_layer(self, predicts, labels, scope&＃61;&＃39;loss_layer&＃39;): #定义损失函数with tf.variable_scope(scope):predict_classes &＃61; tf.reshape( #预测的类别 batchsize x 7x7x20predicts[:, :self.boundary1], [self.batch_size, self.cell_size, self.cell_size, self.num_class])predict_scales &＃61; tf.reshape( #预测的scale batchsize x 7x7x2predicts[:, self.boundary1:self.boundary2],[self.batch_size, self.cell_size, self.cell_size, self.boxes_per_cell])predict_boxes &＃61; tf.reshape( #预测的框 batchsize x 7x7x2,每个box四个位置坐标信息predicts[:, self.boundary2:],[self.batch_size, self.cell_size, self.cell_size, self.boxes_per_cell, 4])response &＃61; tf.reshape( #label后0位置&＃xff1a;有无目标 labels[..., 0],[self.batch_size, self.cell_size, self.cell_size, 1])boxes &＃61; tf.reshape( #label后&＃xff08;1,2,3,4&＃xff09;位置&＃xff1a;目标坐标labels[..., 1:5],[self.batch_size, self.cell_size, self.cell_size, 1, 4])boxes &＃61; tf.tile( #由于单个cell预测boxes_per_cell个box信息&＃xff0c;先对box进行该维度上的拼贴一份相同尺度的&＃xff1b;后将坐标尺度归一化到整幅图boxes, [1, 1, 1, self.boxes_per_cell, 1]) / self.image_sizeclasses &＃61; labels[..., 5:] #label后[5:25]位置&＃xff1a;目标类别信息offset &＃61; tf.reshape(tf.constant(self.offset, dtype&＃61;tf.float32), #将offset维度由7x7x2 reshape成 1x7x7x2[1, self.cell_size, self.cell_size, self.boxes_per_cell])offset &＃61; tf.tile(offset, [self.batch_size, 1, 1, 1]) #将offset的第一个维度拼贴为batchsize大小&＃xff0c;即offset变为&＃xff1a;batchsize x 7x7x2offset_tran &＃61; tf.transpose(offset, (0, 2, 1, 3)) #作者是否考虑非AXA情况&＃xff1f;&＃xff1f;如7x8predict_boxes_tran &＃61; tf.stack([(predict_boxes[..., 0] &＃43; offset) / self.cell_size, #&＃xff08;预测box的x坐标&＃43;偏移量&＃xff09;/7(predict_boxes[..., 1] &＃43; offset_tran) / self.cell_size, #&＃xff08;预测box的y坐标&＃43;偏移量&＃xff09;/7tf.square(predict_boxes[..., 2]), #对w求平方tf.square(predict_boxes[..., 3])], axis&＃61;-1) #对h求平方iou_predict_truth &＃61; self.calc_iou(predict_boxes_tran, boxes) #计算IOU的值# calculate I tensor [BATCH_SIZE, CELL_SIZE, CELL_SIZE, BOXES_PER_CELL] #计算有目标object_mask object_mask &＃61; tf.reduce_max(iou_predict_truth, 3, keep_dims&＃61;True) #找出iou_predict_truth 第 3维度&＃xff08;即box_per_cell&＃xff09;维度计算得到的最大值构成一个tensorobject_mask &＃61; tf.cast((iou_predict_truth >&＃61; object_mask), tf.float32) * response #object_mask:表示有目标 以及 目标与gt的IOU # calculate no_I tensor [CELL_SIZE, CELL_SIZE, BOXES_PER_CELL] #计算无目标noobject_masknoobject_mask &＃61; tf.ones_like( #新建一个与给定tensor&＃xff08;object_mask&＃xff09;大小一致的tensor&＃xff0c;其所有元素都为1object_mask, dtype&＃61;tf.float32) - object_mask boxes_tran &＃61; tf.stack([boxes[..., 0] * self.cell_size - offset,boxes[..., 1] * self.cell_size - offset_tran,tf.sqrt(boxes[..., 2]),tf.sqrt(boxes[..., 3])], axis&＃61;-1)# class_loss #类别损失函数class_delta &＃61; response * (predict_classes - classes) #有目标情况下 类别误差class_loss &＃61; tf.reduce_mean( tf.reduce_sum(tf.square(class_delta), axis&＃61;[1, 2, 3]), #对7x7x20每个维度上预测的类别做误差平方求和后&＃xff0c;乘以损失函数系数class_scalename&＃61;&＃39;class_loss&＃39;) * self.class_scale# object_loss #含有object的box的confidence预测object_delta &＃61; object_mask * (predict_scales - iou_predict_truth)object_loss &＃61; tf.reduce_mean(tf.reduce_sum(tf.square(object_delta), axis&＃61;[1, 2, 3]),name&＃61;&＃39;object_loss&＃39;) * self.object_scale# noobject_loss #不含object的box的confidence预测noobject_delta &＃61; noobject_mask * predict_scalesnoobject_loss &＃61; tf.reduce_mean(tf.reduce_sum(tf.square(noobject_delta), axis&＃61;[1, 2, 3]),name&＃61;&＃39;noobject_loss&＃39;) * self.noobject_scale# coord_loss #坐标损失函数coord_mask &＃61; tf.expand_dims(object_mask, 4) #先扩维boxes_delta &＃61; coord_mask * (predict_boxes - boxes_tran) #需要判断第i个cell中第j个box会否负责这个objectcoord_loss &＃61; tf.reduce_mean( tf.reduce_sum(tf.square(boxes_delta), axis&＃61;[1, 2, 3, 4]), #坐标四个维度对应求差&＃xff0c;平方和name&＃61;&＃39;coord_loss&＃39;) * self.coord_scaletf.losses.add_loss(class_loss)tf.losses.add_loss(object_loss)tf.losses.add_loss(noobject_loss)tf.losses.add_loss(coord_loss)tf.summary.scalar(&＃39;class_loss&＃39;, class_loss) #以下方便tensorboard显示用tf.summary.scalar(&＃39;object_loss&＃39;, object_loss)tf.summary.scalar(&＃39;noobject_loss&＃39;, noobject_loss)tf.summary.scalar(&＃39;coord_loss&＃39;, coord_loss)tf.summary.histogram(&＃39;boxes_delta_x&＃39;, boxes_delta[..., 0])tf.summary.histogram(&＃39;boxes_delta_y&＃39;, boxes_delta[..., 1])tf.summary.histogram(&＃39;boxes_delta_w&＃39;, boxes_delta[..., 2])tf.summary.histogram(&＃39;boxes_delta_h&＃39;, boxes_delta[..., 3])tf.summary.histogram(&＃39;iou&＃39;, iou_predict_truth)def leaky_relu(alpha): #leaky_relu激活函数def op(inputs):return tf.nn.leaky_relu(inputs, alpha&＃61;alpha, name&＃61;&＃39;leaky_relu&＃39;)return op


import os
import argparse
import datetime
import tensorflow as tf
import yolo.config as cfg
from yolo.yolo_net import YOLONet
from utils.timer import Timer
from utils.pascal_voc import pascal_voc slim &＃61; tf.contrib.slim #tensorflow 16年推出的瘦身版代码模块#这部分主要是用pascal_voc2007数据训练自己的网络权重数据
class Solver(object): def __init__(self, net, data):self.net &＃61; netself.data &＃61; dataself.weights_file &＃61; cfg.WEIGHTS_FILE #权重文件&＃xff0c;默认无self.max_iter &＃61; cfg.MAX_ITER #默认15000self.initial_learning_rate &＃61; cfg.LEARNING_RATE #初始学习率0.0001self.decay_steps &＃61; cfg.DECAY_STEPS #衰减步长&＃xff1a;30000self.decay_rate &＃61; cfg.DECAY_RATE #衰减率&＃xff1a;0.1self.staircase &＃61; cfg.STAIRCASEself.summary_iter &＃61; cfg.SUMMARY_ITER #日志记录迭代步数&＃xff1a;10self.save_iter &＃61; cfg.SAVE_ITER #保存迭代步长&＃xff1a;1000self.output_dir &＃61; os.path.join(cfg.OUTPUT_DIR, datetime.datetime.now().strftime(&＃39;%Y_%m_%d_%H_%M&＃39;)) #保存路径&＃xff1a;output/年_月_日_时_分if not os.path.exists(self.output_dir):os.makedirs(self.output_dir)self.save_cfg()self.variable_to_restore &＃61; tf.global_variables()self.saver &＃61; tf.train.Saver(self.variable_to_restore, max_to_keep&＃61;None)self.ckpt_file &＃61; os.path.join(self.output_dir, &＃39;yolo&＃39;) #模型文件路径&＃xff1a; 输出目录/yoloself.summary_op &＃61; tf.summary.merge_all()self.writer &＃61; tf.summary.FileWriter(self.output_dir, flush_secs&＃61;60)self.global_step &＃61; tf.train.create_global_step()self.learning_rate &＃61; tf.train.exponential_decay( #产生一个指数衰减的学习速率&＃xff0c;learning_rate&＃61;initial_learning_rate*decay_rate^(global_step/decay_steps)self.initial_learning_rate, self.global_step, self.decay_steps,self.decay_rate, self.staircase, name&＃61;&＃39;learning_rate&＃39;)self.optimizer &＃61; tf.train.GradientDescentOptimizer(learning_rate&＃61;self.learning_rate)self.train_op &＃61; slim.learning.create_train_op(self.net.total_loss, self.optimizer, global_step&＃61;self.global_step)gpu_options &＃61; tf.GPUOptions()config &＃61; tf.ConfigProto(gpu_options&＃61;gpu_options)self.sess &＃61; tf.Session(config&＃61;config)self.sess.run(tf.global_variables_initializer())if self.weights_file is not None: #如果权重文件空&＃xff0c;则打印“恢复权重文件从&＃xff1a;”print(&＃39;Restoring weights from: &＃39; &＃43; self.weights_file)self.saver.restore(self.sess, self.weights_file)self.writer.add_graph(self.sess.graph)def train(self):train_timer &＃61; Timer() #定义类对象load_timer &＃61; Timer()for step in range(1, self.max_iter &＃43; 1): #最大迭代&＃xff1a;15000load_timer.tic() #开始计时images, labels &＃61; self.data.get() #从pascal_voc数据集读取图像和实际标签信息load_timer.toc() #终止该步&＃xff08;数据加载&＃xff09;计时feed_dict &＃61; {self.net.images: images, #生成一个图像和label对应的字典self.net.labels: labels}if step % self.summary_iter &＃61;&＃61; 0: #迭代每10步时执行如下&＃xff1a;日志记录步长if step % (self.summary_iter * 10) &＃61;&＃61; 0: #迭代每100步时执行如下&＃xff1a;训练模型&＃xff0c;生成报文并打印&＃xff08;主要是打印报文&＃xff09;train_timer.tic() #训练开始计时summary_str, loss, _ &＃61; self.sess.run([self.summary_op, self.net.total_loss, self.train_op], #模型训练&＃xff0c;返回 lossfeed_dict&＃61;feed_dict)train_timer.toc() #训练结束计时log_str &＃61; &＃39;&＃39;&＃39;{} Epoch: {}, Step: {}, Learning rate: {},&＃39;&＃39;&＃39; #报文字符串内容&＃39;&＃39;&＃39; Loss: {:5.3f}\nSpeed: {:.3f}s/iter,&＃39;&＃39;&＃39;&＃39;&＃39;&＃39;&＃39; Load: {:.3f}s/iter, Remain: {}&＃39;&＃39;&＃39;.format(datetime.datetime.now().strftime(&＃39;%m-%d %H:%M:%S&＃39;),self.data.epoch,int(step),round(self.learning_rate.eval(session&＃61;self.sess), 6),loss,train_timer.average_time,load_timer.average_time,train_timer.remain(step, self.max_iter))print(log_str) else: #训练模型&＃xff0c;并计时train_timer.tic()summary_str, _ &＃61; self.sess.run([self.summary_op, self.train_op],feed_dict&＃61;feed_dict)train_timer.toc()self.writer.add_summary(summary_str, step) #每训练10步&＃xff0c;记录日志文件else: #其他训练步长时&＃xff0c;不记录日志&＃xff0c;只计时train_timer.tic()self.sess.run(self.train_op, feed_dict&＃61;feed_dict)train_timer.toc()if step % self.save_iter &＃61;&＃61; 0: #模型每训练1000步保存一次print(&＃39;{} Saving checkpoint file to: {}&＃39;.format(datetime.datetime.now().strftime(&＃39;%m-%d %H:%M:%S&＃39;),self.output_dir))self.saver.save(self.sess, self.ckpt_file, global_step&＃61;self.global_step)def save_cfg(self): #保存当前的模型配置信息with open(os.path.join(self.output_dir, &＃39;config.txt&＃39;), &＃39;w&＃39;) as f: #往output/config.txt中写配置信息cfg_dict &＃61; cfg.__dict__for key in sorted(cfg_dict.keys()):if key[0].isupper():cfg_str &＃61; &＃39;{}: {}\n&＃39;.format(key, cfg_dict[key])f.write(cfg_str)def update_config_paths(data_dir, weights_file):cfg.DATA_PATH &＃61; data_dircfg.PASCAL_PATH &＃61; os.path.join(data_dir, &＃39;pascal_voc&＃39;)cfg.CACHE_PATH &＃61; os.path.join(cfg.PASCAL_PATH, &＃39;cache&＃39;)cfg.OUTPUT_DIR &＃61; os.path.join(cfg.PASCAL_PATH, &＃39;output&＃39;)cfg.WEIGHTS_DIR &＃61; os.path.join(cfg.PASCAL_PATH, &＃39;weights&＃39;) #权重文件在pascal_voc/weights中cfg.WEIGHTS_FILE &＃61; os.path.join(cfg.WEIGHTS_DIR, weights_file)def main():parser &＃61; argparse.ArgumentParser()parser.add_argument(&＃39;--weights&＃39;, default&＃61;"YOLO_small.ckpt", type&＃61;str)parser.add_argument(&＃39;--data_dir&＃39;, default&＃61;"data", type&＃61;str)parser.add_argument(&＃39;--threshold&＃39;, default&＃61;0.2, type&＃61;float)parser.add_argument(&＃39;--iou_threshold&＃39;, default&＃61;0.5, type&＃61;float)parser.add_argument(&＃39;--gpu&＃39;, default&＃61;&＃39;&＃39;, type&＃61;str)args &＃61; parser.parse_args()if args.gpu is not None: #如果训练传进来的gpu参数非空&＃xff0c;则将传进来的gpu信息赋值给配置文件中cfg.GPU &＃61; args.gpuif args.data_dir !&＃61; cfg.DATA_PATH: #如果传经来的数据路径与当前配置文件数据路径不一致&＃xff0c;则更新配置信息update_config_paths(args.data_dir, args.weights)os.environ[&＃39;CUDA_VISIBLE_DEVICES&＃39;] &＃61; cfg.GPUyolo &＃61; YOLONet() #声明类对象yolopascal &＃61; pascal_voc(&＃39;train&＃39;) #定义类别solver &＃61; Solver(yolo, pascal) #利用yolo网络结构&＃xff0c;对传进的数据&＃xff0c;生成solverprint(&＃39;Start training ...&＃39;) #开始训练solver.train()print(&＃39;Done training.&＃39;) #完成训练if __name__ &＃61;&＃61; &＃39;__main__&＃39;:# python train.py --weights YOLO_small.ckpt --gpu 0 #示例&＃xff0c;默认使用第0个GPUmain()


import os
import cv2
import argparse
import numpy as np
import tensorflow as tf
import yolo.config as cfg
from yolo.yolo_net import YOLONet
from utils.timer import Timer#这部分主要是加载训练好的权重文件做测试&＃xff0c;这个权重文件可以是下载的YOLO_small.ckpt&＃xff0c;也可以是自己训练的。
class Detector(object):def __init__(self, net, weight_file):self.net &＃61; netself.weights_file &＃61; weight_fileself.classes &＃61; cfg.CLASSESself.num_class &＃61; len(self.classes)self.image_size &＃61; cfg.IMAGE_SIZEself.cell_size &＃61; cfg.CELL_SIZEself.boxes_per_cell &＃61; cfg.BOXES_PER_CELLself.threshold &＃61; cfg.THRESHOLDself.iou_threshold &＃61; cfg.IOU_THRESHOLDself.boundary1 &＃61; self.cell_size * self.cell_size * self.num_classself.boundary2 &＃61; self.boundary1 &＃43;\self.cell_size * self.cell_size * self.boxes_per_cellself.sess &＃61; tf.Session()self.sess.run(tf.global_variables_initializer())print(&＃39;Restoring weights from: &＃39; &＃43; self.weights_file) self.saver &＃61; tf.train.Saver()self.saver.restore(self.sess, self.weights_file) #加载权重文件def draw_result(self, img, result): #在输入图像img上对检测到的result进行绘制框并标注类别概率信息for i in range(len(result)): #目标个数遍历绘图x &＃61; int(result[i][1]) #目标中心xy &＃61; int(result[i][2]) #目标中心yw &＃61; int(result[i][3] / 2) #目标宽取一半h &＃61; int(result[i][4] / 2) #目标高取一半cv2.rectangle(img, (x - w, y - h), (x &＃43; w, y &＃43; h), (0, 255, 0), 2) #目标框cv2.rectangle(img, (x - w, y - h - 20), #显示目标类别和概率值的灰色填充框(x &＃43; w, y - h), (125, 125, 125), -1)lineType &＃61; cv2.LINE_AA if cv2.__version__ > &＃39;3&＃39; else cv2.CV_AA #根据opencv版本&＃xff0c;作者已经做了考虑了cv2.putText(img, result[i][0] &＃43; &＃39; : %.2f&＃39; % result[i][5], #概率是两位小数的浮点数(x - w &＃43; 5, y - h - 7), cv2.FONT_HERSHEY_SIMPLEX, 0.5,(0, 0, 0), 1, lineType)def detect(self, img): #对输入图像做目标检测img_h, img_w, _ &＃61; img.shapeinputs &＃61; cv2.resize(img, (self.image_size, self.image_size)) #尺寸缩放到448x448的图像&＃xff1a;inputsinputs &＃61; cv2.cvtColor(inputs, cv2.COLOR_BGR2RGB).astype(np.float32) #opencv读取图像格式是bgr,需要转换为rgb格式&＃xff1b;inputs &＃61; (inputs / 255.0) * 2.0 - 1.0 #读取图像归一化到【-1,1】inputs &＃61; np.reshape(inputs, (1, self.image_size, self.image_size, 3)) #维度变化为[1,448,448,3]result &＃61; self.detect_from_cvmat(inputs)[0] for i in range(len(result)):result[i][1] *&＃61; (1.0 * img_w / self.image_size) #检测到目标中心坐标x是448下的坐标&＃xff0c;需要变化到原图像尺寸result[i][2] *&＃61; (1.0 * img_h / self.image_size)result[i][3] *&＃61; (1.0 * img_w / self.image_size)result[i][4] *&＃61; (1.0 * img_h / self.image_size)return result #返回原图像上检测到的目标坐标尺寸信息def detect_from_cvmat(self, inputs): #输入的inputs&＃xff1b;[1,448,448,3]net_output &＃61; self.sess.run(self.net.logits, #网络回归输出目标feed_dict&＃61;{self.net.images: inputs})results &＃61; []for i in range(net_output.shape[0]): #遍历目标个数&＃xff0c;将结果放进results中results.append(self.interpret_output(net_output[i]))return results #在448x448大小图像上检测到的目标信息def interpret_output(self, output):probs &＃61; np.zeros((self.cell_size, self.cell_size, #所有box (98个)对应每个类别的概率&＃xff0c;[7,7,2,20]self.boxes_per_cell, self.num_class))class_probs &＃61; np.reshape(output[0:self.boundary1], #输出的[0:7x7x20]这980个数代表每个cell预测的每个类别的概率值(self.cell_size, self.cell_size, self.num_class)) #最后输出时&＃xff0c;每个cell只返回一个类别&＃xff0c;因此类别概率维度变为[7,7,20]scales &＃61; np.reshape(output[self.boundary1:self.boundary2], #输出的[7x7x20:7x7x22]这98个数reshape成[7,7,2]&＃xff0c;个人理解是有无目标落在这98个box中(self.cell_size, self.cell_size, self.boxes_per_cell))boxes &＃61; np.reshape( #输出的[7x7x22:]这些数记录的是每个box对应的目标坐标信息&＃xff0c;reshape为[7,7,2,4]output[self.boundary2:],(self.cell_size, self.cell_size, self.boxes_per_cell, 4))offset &＃61; np.array([np.arange(self.cell_size)] * self.cell_size * self.boxes_per_cell)offset &＃61; np.transpose(np.reshape(offset,[self.boxes_per_cell, self.cell_size, self.cell_size]),#offset&＃xff1b;[2,7,7]->[7,7,2](1, 2, 0))boxes[:, :, :, 0] &＃43;&＃61; offset boxes[:, :, :, 1] &＃43;&＃61; np.transpose(offset, (1, 0, 2)) boxes[:, :, :, :2] &＃61; 1.0 * boxes[:, :, :, 0:2] / self.cell_sizeboxes[:, :, :, 2:] &＃61; np.square(boxes[:, :, :, 2:])boxes *&＃61; self.image_size #将目标坐标相对cell的偏移量反映到448图像上for i in range(self.boxes_per_cell):for j in range(self.num_class):probs[:, :, i, j] &＃61; np.multiply( #某cell中第i个box中含目标的概率*该cell中数据第j个类别概率class_probs[:, :, j], scales[:, :, i])filter_mat_probs &＃61; np.array(probs >&＃61; self.threshold, dtype&＃61;&＃39;bool&＃39;) #若概率大于0.2&＃xff0c;filter_mat_probs&＃61;1filter_mat_boxes &＃61; np.nonzero(filter_mat_probs) #过滤掉一个cell中的两个box的其中一个&＃xff0c;返回filter_mat_probs中不为0的下标boxes_filtered &＃61; boxes[filter_mat_boxes[0],filter_mat_boxes[1], filter_mat_boxes[2]]probs_filtered &＃61; probs[filter_mat_probs]classes_num_filtered &＃61; np.argmax(filter_mat_probs, axis&＃61;3)[filter_mat_boxes[0], filter_mat_boxes[1], filter_mat_boxes[2]]argsort &＃61; np.array(np.argsort(probs_filtered))[::-1]boxes_filtered &＃61; boxes_filtered[argsort] #过滤刷选出boxprobs_filtered &＃61; probs_filtered[argsort] #过滤刷选出probs高的classes_num_filtered &＃61; classes_num_filtered[argsort] #过滤刷选出类别for i in range(len(boxes_filtered)):if probs_filtered[i] &＃61;&＃61; 0:continuefor j in range(i &＃43; 1, len(boxes_filtered)):if self.iou(boxes_filtered[i], boxes_filtered[j]) > self.iou_threshold:probs_filtered[j] &＃61; 0.0filter_iou &＃61; np.array(probs_filtered > 0.0, dtype&＃61;&＃39;bool&＃39;)boxes_filtered &＃61; boxes_filtered[filter_iou]probs_filtered &＃61; probs_filtered[filter_iou]classes_num_filtered &＃61; classes_num_filtered[filter_iou]result &＃61; []for i in range(len(boxes_filtered)):result.append([self.classes[classes_num_filtered[i]],boxes_filtered[i][0],boxes_filtered[i][1],boxes_filtered[i][2],boxes_filtered[i][3],probs_filtered[i]])return result #输出过滤后的类别&＃xff0c;以及对应box的坐标def iou(self, box1, box2):tb &＃61; min(box1[0] &＃43; 0.5 * box1[2], box2[0] &＃43; 0.5 * box2[2]) - \ #得到的tb为重叠区域的宽max(box1[0] - 0.5 * box1[2], box2[0] - 0.5 * box2[2])lr &＃61; min(box1[1] &＃43; 0.5 * box1[3], box2[1] &＃43; 0.5 * box2[3]) - \ #得到的lr为重叠区域的高max(box1[1] - 0.5 * box1[3], box2[1] - 0.5 * box2[3])inter &＃61; 0 if tb <0 or lr <0 else tb * lr #重叠区域面积inter&＃61;tb*lrreturn inter / (box1[2] * box1[3] &＃43; box2[2] * box2[3] - inter) #IOU&＃61;inter/(box1面积&＃43;box2面积)def camera_detector(self, cap, wait&＃61;10): #读取摄像头&＃xff0c;延迟10msdetect_timer &＃61; Timer()ret, _ &＃61; cap.read()while ret:ret, frame &＃61; cap.read()detect_timer.tic()result &＃61; self.detect(frame)detect_timer.toc()print(&＃39;Average detecting time: {:.3f}s&＃39;.format( #统计平均检测时间detect_timer.average_time))self.draw_result(frame, result) #绘制结果cv2.imshow(&＃39;Camera&＃39;, frame)cv2.waitKey(wait)ret, frame &＃61; cap.read()def image_detector(self, imname, wait&＃61;0): #读取图像&＃xff0c;一直显示detect_timer &＃61; Timer()image &＃61; cv2.imread(imname)detect_timer.tic()result &＃61; self.detect(image)detect_timer.toc()print(&＃39;Average detecting time: {:.3f}s&＃39;.format(detect_timer.average_time))self.draw_result(image, result)cv2.imshow(&＃39;Image&＃39;, image)cv2.waitKey(wait)def main():parser &＃61; argparse.ArgumentParser()parser.add_argument(&＃39;--weights&＃39;, default&＃61;"YOLO_small.ckpt", type&＃61;str)parser.add_argument(&＃39;--weight_dir&＃39;, default&＃61;&＃39;weights&＃39;, type&＃61;str)parser.add_argument(&＃39;--data_dir&＃39;, default&＃61;"data", type&＃61;str)parser.add_argument(&＃39;--gpu&＃39;, default&＃61;&＃39;&＃39;, type&＃61;str)args &＃61; parser.parse_args()os.environ[&＃39;CUDA_VISIBLE_DEVICES&＃39;] &＃61; args.gpuyolo &＃61; YOLONet(False)weight_file &＃61; os.path.join(args.data_dir, args.weight_dir, args.weights) #权重文件目录detector &＃61; Detector(yolo, weight_file) # detect from camera #以下是用摄像头做检测输入源# cap &＃61; cv2.VideoCapture(-1)# detector.camera_detector(cap)# detect from image file #以下是用图像做检测输入源imname &＃61; &＃39;test/person.jpg&＃39;detector.image_detector(imname)if __name__ &＃61;&＃61; &＃39;__main__&＃39;:main()


我的运行结果如下&＃xff08;环境是win10,用Spyder运行的&＃xff0c;其中tensorflow版本建议更换到1.4以上&＃xff09;&＃xff1a;

我下载的YOLO_small.ckpt放在weights目录下。单张图像检测用时2.527s。

运行视频时如果报错&＃xff0c;建议将

with tf.variable_scope(scope):

中添加reuse&＃61;True就OK了&＃xff0c;关于摄像头参数传0&＃xff0c;-1&＃xff0c;还是1看你具体设备了。

YOLO对卡通识别效果不错&＃xff0c;只是当前模型识别种类太少。

参考文章&＃xff1a;

https://blog.csdn.net/qq1483661204/article/details/79681926

https://blog.csdn.net/qq_34784753/article/details/78803423