yolo源码解析（二）

五 读取数据pascal_voc.py文件解析

我们在YOLENet类中定义了两个占位符，一个是输入图片占位符，一个是图片对应的标签占位符，如下：

        #输入图片占位符 [NONE,image_size,image_size,3]
        self.images = tf.placeholder(
            tf.float32, [None, self.image_size, self.image_size, 3],
            name='images')
        #设置标签占位符 [None,S,S,5+C]  即[None,7,7,25]
        self.labels = tf.placeholder(
            tf.float32,
            [None, self.cell_size, self.cell_size, 5 + self.num_class])        

而pascal_voc.py文件的目的就是为了准备数据，赋值给占位符。在pascal_voc.py文件中定义了一个pascal_voc，该类包含了类初始化函数(__init__())，准备数据函数(prepare()),读取batch大小的图片以及图片对应的标签(get())等函数。

import os
import xml.etree.ElementTree as ET
import numpy as np
import cv2
import pickle
import copy
import yolo.config as cfg


'''
VOC2012数据集处理
'''

class pascal_voc(object):

1、类初始化函数

'''
VOC2012数据集处理
'''

class pascal_voc(object):
    '''
    VOC2012数据集处理的类，主要用来获取训练集图片文件，以及生成对应的标签文件
    '''
    def __init__(self, phase, rebuild=False):
        '''
        准备训练或者测试的数据
        
        args:
            phase：传入字符串 'train'：表示训练
                              'test'：测试
            rebuild：是否重新创建数据集的标签文件，保存在缓存文件夹下
        '''
        #VOCdevkit文件夹路径
        self.devkil_path = os.path.join(cfg.PASCAL_PATH, 'VOCdevkit')
        #VOC2012文件夹路径
        self.data_path = os.path.join(self.devkil_path, 'VOC2012')
        #catch文件所在路径
        self.cache_path = cfg.CACHE_PATH
        #批大小
        self.batch_size = cfg.BATCH_SIZE
        #图像大小
        self.image_size = cfg.IMAGE_SIZE
        #单元格大小S
        self.cell_size = cfg.CELL_SIZE
        #VOC 2012数据集类别名
        self.classes = cfg.CLASSES
        #类别名->索引的dict
        self.class_to_ind = dict(zip(self.classes, range(len(self.classes))))
        ##图片是否采用水平镜像扩充训练集？
        self.flipped = cfg.FLIPPED
        #训练或测试？
        self.phase = phase
        #是否重新创建数据集标签文件
        self.rebuild = rebuild
        #从gt_labels加载数据，cursor表明当前读取到第几个
        self.cursor = 0
        #存放当前训练的轮数
        self.epoch = 1
        #存放数据集的标签 是一个list 每一个元素都是一个dict，对应一个图片 
        #如果我们在配置文件中指定flipped=True，则数据集会扩充一倍，每一张原始图片都有一个水平对称的镜像文件
        #      imname：图片路径 
        #      label：图片标签
        #      flipped：图片水平镜像？
        self.gt_labels = None
        #加载数据集标签  初始化gt_labels
        self.prepare()

2、prepare()所有数据准备函数

prepare()函数调用load_labels()函数，加载所有数据集的标签，保存在遍历gt_labels集合中，如果在配置文件中指定了水平镜像，则追加一倍的训练数据集。

    def prepare(self):
        '''
        初始化数据集的标签，保存在变量gt_labels中
        
        return:
            gt_labels:返回数据集的标签 是一个list  每一个元素对应一张图片，是一个dict                       
                                     imname：图片文件路径
                                     label：图片文件对应的标签 [7,7,25]的矩阵
                                     flipped：是否使用水平镜像？ 设置为False
        '''
        #加载数据集的标签
        gt_labels = self.load_labels()
        #如果水平镜像，则追加一倍的训练数据集
        if self.flipped:
            print('Appending horizontally-flipped training examples ...')
            #深度拷贝
            gt_labels_cp = copy.deepcopy(gt_labels)
            #遍历每一个图片标签
            for idx in range(len(gt_labels_cp)):
                #设置flipped属性为True
                gt_labels_cp[idx]['flipped'] = True
                #目标所在格子也进行水平镜像 [7,7,25]
                gt_labels_cp[idx]['label'] =
                    gt_labels_cp[idx]['label'][:, ::-1, :]
                for i in range(self.cell_size):
                    for j in range(self.cell_size):
                        #置信度==1，表示这个格子有目标
                        if gt_labels_cp[idx]['label'][i, j, 0] == 1:
                            #中心的x坐标水平镜像
                            gt_labels_cp[idx]['label'][i, j, 1] = 
                                self.image_size - 1 -
                                gt_labels_cp[idx]['label'][i, j, 1]
            #追加数据集的标签   后面的是由原数据集标签扩充的水平镜像数据集标签
            gt_labels += gt_labels_cp
        #打乱数据集的标签
        np.random.shuffle(gt_labels)
        self.gt_labels = gt_labels
        return gt_labels

3、get()批量数据读取函数

get()函数用在训练的时候，每次从gt_labels集合随机读取batch大小的图片以及图片对应的标签。

    def get(self):
        '''
        加载数据集 每次读取batch大小的图片以及图片对应的标签
        
        return:
            images:读取到的图片数据 [45,448,448,3]
            labels:对应的图片标签 [45,7,7,25]
        '''
        #[45,448,448,3]
        images = np.zeros(
            (self.batch_size, self.image_size, self.image_size, 3))
        #[45,7,7,25]
        labels = np.zeros(
            (self.batch_size, self.cell_size, self.cell_size, 25))
        count = 0
        #一次加载batch_size个图片数据
        while count < self.batch_size:
            #获取图片路径
            imname = self.gt_labels[self.cursor]['imname']
            #是否使用水平镜像？
            flipped = self.gt_labels[self.cursor]['flipped']
            #读取图片数据
            images[count, :, :, :] = self.image_read(imname, flipped)
            #读取图片标签
            labels[count, :, :, :] = self.gt_labels[self.cursor]['label']
            count += 1
            self.cursor += 1
            #如果读取完一轮数据，则当前cursor置为0，当前训练轮数+1
            if self.cursor >= len(self.gt_labels):
                #打乱数据集
                np.random.shuffle(self.gt_labels)
                self.cursor = 0                
                self.epoch += 1
        return images, labels

4、image_read()函数读取图片

图片读取函数，先读取图片，然后缩放，转换为RGB格式，再对数据进行归一化处理。

    def image_read(self, imname, flipped=False):
        '''
        读取图片
        
        args:
            imname：图片路径
            flipped：图片是否水平镜像处理？ 
            
        return：
            image：图片数据 [448,448,3]
        '''
        #读取图片数据
        image = cv2.imread(imname)
        #缩放处理
        image = cv2.resize(image, (self.image_size, self.image_size))
        #BGR->RGB  uint->float32
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
        #归一化处理 [-1.0,1.0]
        image = (image / 255.0) * 2.0 - 1.0
        #宽倒序  即水平镜像
        if flipped:
            image = image[:, ::-1, :]
        return image

5、load_labels()加载标签函数

    def load_labels(self):
        '''
        加载数据集标签
        
        return:
            gt_labels：是一个list  每一个元素对应一张图片，是一个dict                       
                                     imname：图片文件路径
                                     label：图片文件对应的标签 [7,7,25]的矩阵
                                     flipped：是否使用水平镜像？ 设置为False   
        '''
        #缓冲文件名：即用来保存数据集标签的文件
        cache_file = os.path.join(
            self.cache_path, 'pascal_' + self.phase + '_gt_labels.pkl')

        #文件存在，且不重新创建则直接读取
        if os.path.isfile(cache_file) and not self.rebuild:
            print('Loading gt_labels from: ' + cache_file)
            with open(cache_file, 'rb') as f:
                gt_labels = pickle.load(f)
            return gt_labels

        print('Processing gt_labels from: ' + self.data_path)

        #如果缓冲文件目录不存在，创建
        if not os.path.exists(self.cache_path):
            os.makedirs(self.cache_path)
            
        #获取训练测试集的数据文件名
        if self.phase == 'train':
            txtname = os.path.join(
                self.data_path, 'ImageSets', 'Main', 'trainval.txt')
        #获取测试集的数据文件名
        else:
            txtname = os.path.join(
                self.data_path, 'ImageSets', 'Main', 'test.txt')
        with open(txtname, 'r') as f:
            self.image_index = [x.strip() for x in f.readlines()]

        #存放图片的标签，图片路径，是否使用水平镜像？
        gt_labels = []
        #遍历每一张图片的信息
        for index in self.image_index:
            #读取每一张图片的标签label [7,7,25]
            label, num = self.load_pascal_annotation(index)
            if num == 0:
                continue
            #图片文件路径
            imname = os.path.join(self.data_path, 'JPEGImages', index + '.jpg')
            #保存该图片的信息
            gt_labels.append({'imname': imname,
                              'label': label,
                              'flipped': False})
        print('Saving gt_labels to: ' + cache_file)
        #保存
        with open(cache_file, 'wb') as f:
            pickle.dump(gt_labels, f)
        return gt_labels

6、load_pascal_annotation()函数

    def load_pascal_annotation(self, index):
        """
        Load image and bounding boxes info from XML file in the PASCAL VOC
        format.
        
        args:
            index：图片文件的index
            
        return :
            label：标签 [7,7,25] 
                      0:1：置信度，表示这个地方是否有目标
                      1:5：目标边界框  目标中心，宽度和高度(这里是实际值，没有归一化)
                      5:25：目标的类别
            len(objs)：objs对象长度
        """
        #获取图片文件名路径
        imname = os.path.join(self.data_path, 'JPEGImages', index + '.jpg')
        #读取数据
        im = cv2.imread(imname)
        #宽和高缩放比例
        h_ratio = 1.0 * self.image_size / im.shape[0]
        w_ratio = 1.0 * self.image_size / im.shape[1]
        # im = cv2.resize(im, [self.image_size, self.image_size])
        #用于保存图片文件的标签
        label = np.zeros((self.cell_size, self.cell_size, 25))
        #图片文件的标注xml文件
        filename = os.path.join(self.data_path, 'Annotations', index + '.xml')
        tree = ET.parse(filename)
        objs = tree.findall('object')

        for obj in objs:
            bbox = obj.find('bndbox')
            # Make pixel indexes 0-based  当图片缩放到image_size时，边界框也进行同比例缩放
            x1 = max(min((float(bbox.find('xmin').text) - 1) * w_ratio, self.image_size - 1), 0)
            y1 = max(min((float(bbox.find('ymin').text) - 1) * h_ratio, self.image_size - 1), 0)
            x2 = max(min((float(bbox.find('xmax').text) - 1) * w_ratio, self.image_size - 1), 0)
            y2 = max(min((float(bbox.find('ymax').text) - 1) * h_ratio, self.image_size - 1), 0)
            #根据图片的分类名 ->类别index 转换
            cls_ind = self.class_to_ind[obj.find('name').text.lower().strip()]
            #计算边框中心点x,y,w,h(没有归一化)
            boxes = [(x2 + x1) / 2.0, (y2 + y1) / 2.0, x2 - x1, y2 - y1]
            #计算当前物体的中心在哪个格子中
            x_ind = int(boxes[0] * self.cell_size / self.image_size)
            y_ind = int(boxes[1] * self.cell_size / self.image_size)
            #表明该图片已经初始化过了
            if label[y_ind, x_ind, 0] == 1:
                continue
            #置信度，表示这个地方有物体
            label[y_ind, x_ind, 0] = 1
            #物体边界框
            label[y_ind, x_ind, 1:5] = boxes
            #物体的类别
            label[y_ind, x_ind, 5 + cls_ind] = 1

        return label, len(objs)

六 训练网络

模型训练包含于train.py文件，Solver类的train()方法之中，训练部分只需要看懂了初始化参数，整个结构就很清晰了。

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 = tf.contrib.slim

'''
用来训练YOLO网络模型
'''

class Solver(object):
    '''
    求解器的类，用于训练YOLO网络
    '''

1、类初始化函数

   def __init__(self, net, data):
        '''
        构造函数，加载训练参数
        
        args：
            net：YOLONet对象
            data：pascal_voc对象
        '''
        #yolo网络
        self.net = net
        #voc2012数据处理
        self.data = data
        #检查点文件路径
        self.weights_file = cfg.WEIGHTS_FILE
        #训练最大迭代次数
        self.max_iter = cfg.MAX_ITER
        #初始学习率
        self.initial_learning_rate = cfg.LEARNING_RATE
        ##退化学习率衰减步数
        self.decay_steps = cfg.DECAY_STEPS
        #衰减率
        self.decay_rate = cfg.DECAY_RATE
        self.staircase = cfg.STAIRCASE
        ##日志文件保存间隔步
        self.summary_iter = cfg.SUMMARY_ITER
        ##模型保存间隔步
        self.save_iter = cfg.SAVE_ITER
        
        #输出文件夹路径
        self.output_dir = os.path.join(
            cfg.OUTPUT_DIR, datetime.datetime.now().strftime('%Y_%m_%d_%H_%M'))
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)
        #保存配置信息
        self.save_cfg()
        #指定保存的张量 这里指定所有变量
        self.variable_to_restore = tf.global_variables()        
        self.saver = tf.train.Saver(self.variable_to_restore, max_to_keep=None)
        #指定保存的模型名称
        self.ckpt_file = os.path.join(self.output_dir, 'yolo.cpkt')
        #合并所有的summary
        self.summary_op = tf.summary.merge_all()
        #创建writer，指定日志文件路径，用于写日志文件
        self.writer = tf.summary.FileWriter(self.output_dir, flush_secs=60)

        #创建变量，保存当前迭代次数
        self.global_step = tf.train.create_global_step()
        #退化学习率
        self.learning_rate = tf.train.exponential_decay(
            self.initial_learning_rate, self.global_step, self.decay_steps,
            self.decay_rate, self.staircase, name='learning_rate')
        #创建求解器
        self.optimizer = tf.train.GradientDescentOptimizer(
            learning_rate=self.learning_rate)
        # create_train_op that ensures that when we evaluate it to get the loss,
        # the update_ops are done and the gradient updates are computed.
        self.train_op = slim.learning.create_train_op(
            self.net.total_loss, self.optimizer, global_step=self.global_step)

        #设置GPU使用资源
        gpu_options = tf.GPUOptions()
        #按需分配GPU使用的资源
        config = tf.ConfigProto(gpu_options=gpu_options)
        self.sess = tf.Session(config=config)
        
        #运行图之前，初始化变量
        self.sess.run(tf.global_variables_initializer())

        #恢复模型
        if self.weights_file is not None:
            print('Restoring weights from: ' + self.weights_file)
            self.saver.restore(self.sess, self.weights_file)

        #将图写入日志文件
        self.writer.add_graph(self.sess.graph)

 2、train()训练函数

 def train(self):
        '''
        开始训练
        '''
        #训练时间
        train_timer = Timer()
        #数据集加载时间
        load_timer = Timer()

        #开始迭代
        for step in range(1, self.max_iter + 1):
            #计算每次迭代加载数据的起始时间
            load_timer.tic()
            #加载数据集 每次读取batch大小的图片以及图片对应的标签
            images, labels = self.data.get()
            #计算这次迭代加载数据集所使用的时间
            load_timer.toc()
            
            feed_dict = {self.net.images: images,
                         self.net.labels: labels}

            #迭代summary_iter次，保存一次日志文件，迭代summary_iter*10次，输出一次的迭代信息
            if step % self.summary_iter == 0:
                if step % (self.summary_iter * 10) == 0:
                    #计算每次迭代训练的起始时间
                    train_timer.tic()
                    loss = 0.0001  
                    #开始迭代训练，每一次迭代后global_step自加1
                    summary_str, loss, _ = self.sess.run(
                        [self.summary_op, self.net.total_loss, self.train_op],
                        feed_dict=feed_dict)
                    #输出信息
                    log_str = '{} Epoch: {}, Step: {}, Learning rate: {}, Loss: {:5.3f}
Speed: {:.3f}s/iter,Load: {:.3f}s/iter, Remain: {}'.format(
                        datetime.datetime.now().strftime('%m-%d %H:%M:%S'),
                        self.data.epoch,
                        int(step),
                        round(self.learning_rate.eval(session=self.sess), 6),
                        loss,
                        train_timer.average_time,
                        load_timer.average_time,
                        train_timer.remain(step, self.max_iter))
                    print(log_str)

                else:
                    #计算每次迭代训练的起始时间
                    train_timer.tic()           
                    #开始迭代训练，每一次迭代后global_step自加1
                    summary_str, _ = self.sess.run(
                        [self.summary_op, self.train_op],
                        feed_dict=feed_dict)  
                    #计算这次迭代训练所使用的时间
                    train_timer.toc()
                    
                #将summary写入文件
                self.writer.add_summary(summary_str, step)

            else:
                #计算每次迭代训练的起始时间
                train_timer.tic()
                #开始迭代训练，每一次迭代后global_step自加1
                self.sess.run(self.train_op, feed_dict=feed_dict)
                #计算这次迭代训练所使用的时间
                train_timer.toc()

            #没迭代save_iter次，保存一次模型
            if step % self.save_iter == 0:
                print('{} Saving checkpoint file to: {}'.format(
                    datetime.datetime.now().strftime('%m-%d %H:%M:%S'),
                    self.output_dir))
                self.saver.save(
                    self.sess, self.ckpt_file, global_step=self.global_step)

3、保存配置参数

    def save_cfg(self):
        '''
        保存配置信息
        '''
        with open(os.path.join(self.output_dir, 'config.txt'), 'w') as f:
            cfg_dict = cfg.__dict__
            for key in sorted(cfg_dict.keys()):
                if key[0].isupper():
                    cfg_str = '{}: {}
'.format(key, cfg_dict[key])
                    f.write(cfg_str)

train.py文件除了上面介绍的求解器Solver这个类外，还包含了两个函数，一个是update_config_paths()函数，这个函数主要使用了设定数据集路径，以及检查点文件路径。

def update_config_paths(data_dir, weights_file):
    '''
    数据集路径，和模型检查点文件路径
    
    args:
        data_dir:数据文件夹  数据集放在pascal_voc目录下  
        weights_file：检查点文件名 该文件放在数据集目录下的weights文件夹下
        
    '''
    cfg.DATA_PATH = data_dir                                                   #数据所在文件夹
    cfg.PASCAL_PATH = os.path.join(data_dir, 'pascal_voc')                     #VOC2012数据所在文件夹
    cfg.CACHE_PATH = os.path.join(cfg.PASCAL_PATH, 'cache')                    #保存生成的数据集标签缓冲文件所在文件夹
    cfg.OUTPUT_DIR = os.path.join(cfg.PASCAL_PATH, 'output')                   #保存生成的网络模型和日志文件所在的文件夹
    cfg.WEIGHTS_DIR = os.path.join(cfg.PASCAL_PATH, 'weights')                 #检查点文件所在的目录

    cfg.WEIGHTS_FILE = os.path.join(cfg.WEIGHTS_DIR, weights_file)

我们主要来说一下另一个函数main()函数，先解析命令行参数，然后创建YOLONet、pascal_voc、Solver对象，最后开始训练。

def main():
    #创建一个解析器对象，并告诉它将会有些什么参数。当程序运行时，该解析器就可以用于处理命令行参数。
    #https://www.cnblogs.com/lovemyspring/p/3214598.html
    parser = argparse.ArgumentParser()
    #定义参数
    parser.add_argument('--weights', default="YOLO_small.ckpt", type=str)    #权重文件名
    parser.add_argument('--data_dir', default="data", type=str)              #数据集路径
    parser.add_argument('--threshold', default=0.2, type=float)
    parser.add_argument('--iou_threshold', default=0.5, type=float)
    parser.add_argument('--gpu', default='', type=str)
    #定义了所有参数之后，你就可以给 parse_args() 传递一组参数字符串来解析命令行。默认情况下，参数是从 sys.argv[1:] 中获取
    #parse_args() 的返回值是一个命名空间，包含传递给命令的参数。该对象将参数保存其属性
    args = parser.parse_args()

    #判断是否是使用gpu
    if args.gpu is not None:
        cfg.GPU = args.gpu

    #设定数据集路径，以及检查点文件路径
    if args.data_dir != cfg.DATA_PATH  and args.data_dir is not None:
        update_config_paths(args.data_dir, args.weights)

    #设置环境变量
    os.environ['CUDA_VISIBLE_DEVICES'] = cfg.GPU

    #创建YOLO网络对象
    yolo = YOLONet()
    #数据集对象
    pascal = pascal_voc('train')
    #求解器对象
    solver = Solver(yolo, pascal)

    print('Start training ...')
    #开始训练
    solver.train()
    print('Done training.')

我们执行如下代码，开始训练网络：

if __name__ == '__main__':
    tf.reset_default_graph()
    # python train.py --weights YOLO_small.ckpt --gpu 0
    main()