TensorFlow目标检测api训练mobilenet ssd教程

TensorFlow目标检测api训练mobilenet ssd教程

install anaconda

安装依赖

sudo apt-get install libgl1-mesa-glx libegl1-mesa libxrandr2 libxrandr2 libxss1 libxcursor1 libxcomposite1 libasound2 libxi6 libxtst6

下载安装文件

下载地址 python3.7版本地址

安装

chmod +x Anaconda3-2019.10-Linux-x86_64.sh
./Anaconda3-2019.10-Linux-x86_64.sh

开始是一堆霸王条款,一直回车,然后输入yes

Please answer 'yes' or 'no':'
>>> yes

然后选择安装路径,默认就好,直接回车,等待安装即可。

新建环境

conda create --name obj_detection python=3.6

激活环境

conda activate obj_detection

download tensorflow models

git clone https://github.com/tensorflow/models.git

install tf requirements

python

pip install Cython contextlib2 pillow lxml jupyter matplotlib absl-py tensorflow==1.14

注:由于默认下载最新的tensorflow2,由于keras为tf2的御用接口框架,于是slim就找不到了,而且包的结构发生了翻天覆地的变化,找不到tf.contrib了,说白了tf2和tf1是两个不兼容的框架,而该项目是用tf1写的,所以我们指定一个tf1的版本。

当然如果你的机器有gpu想用gpu训练的话,请安装gpu版的tensorflow

pip install tensorflow-gpu==1.14

protoc

下载地址

解压后改名字为protoc,然后复制到刚刚下载的models项目的research目录下,然后在research目录下运行:

./protoc/bin/protoc object_detection/protos/*.proto --python_out=.
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim
python setup.py build
python setup.py install

运行如下Python文件检验环境是否正确

python object_detection/builders/model_builder_test.py

成功的话是这样的:

[此处省略部分打印信息...]
Ran 17 tests in 0.228s

OK (skipped=1)

prepare dataset

  1. 数据集分为训练、验证、测试

    import os
    import random
    import time
    import shutil
    
    # 所有标注的xml文件目录
    xmlfilepath = r'/home/zheshi/h4tv/datasets/VOCdevkit/VOCBall2k/Annotations'
    # 分割数据集后存储的xml文件路径,会存放在这个路径下分成train validation test三个目录
    saveBasePath = r"/home/zheshi/h4tv/datasets/VOCdevkit/VOCBall2k/VienDataset/Annotations"
    
    # 分割比例
    trainval_percent = 0.9
    train_percent = 0.85
    
    total_xml = os.listdir(xmlfilepath)
    num = len(total_xml)
    list = range(num)
    tv = int(num * trainval_percent)
    tr = int(tv * train_percent)
    trainval = random.sample(list, tv)
    train = random.sample(trainval, tr)
    print("train and val size", tv)
    print("train size", tr)
    start = time.time()
    test_num = 0
    val_num = 0
    train_num = 0
    
    for i in list:
       name = total_xml[i]
       if i in trainval:  # train and val set
           if i in train:
               directory = "train"
               train_num += 1
               xml_path = os.path.join(saveBasePath, directory)
               print(xml_path)
               if (not os.path.exists(xml_path)):
                   os.mkdir(xml_path)
               filePath = os.path.join(xmlfilepath, name)
               newfile = os.path.join(saveBasePath, os.path.join(directory, name))
               shutil.copyfile(filePath, newfile)
           else:
               directory = "validation"
               xml_path = os.path.join(saveBasePath, directory)
               print(xml_path)
               if (not os.path.exists(xml_path)):
                   os.mkdir(xml_path)
               val_num += 1
               filePath = os.path.join(xmlfilepath, name)
               newfile = os.path.join(saveBasePath, os.path.join(directory, name))
               shutil.copyfile(filePath, newfile)
       else:  # test set
           directory = "test"
           xml_path = os.path.join(saveBasePath, directory)
           print(xml_path)
           if (not os.path.exists(xml_path)):
               os.mkdir(xml_path)
           test_num += 1
           filePath = os.path.join(xmlfilepath, name)
           newfile = os.path.join(saveBasePath, os.path.join(directory, name))
           shutil.copyfile(filePath, newfile)
    
    # End time
    end = time.time()
    seconds = end - start
    print("train total : " + str(train_num))
    print("validation total : " + str(val_num))
    print("test total : " + str(test_num))
    total_num = train_num + val_num + test_num
    print("total number : " + str(total_num))
    print("Time taken : {0} seconds".format(seconds))
    
  2. xml转csv

    import os
    import glob
    import pandas as pd
    import xml.etree.ElementTree as ET
    
    def xml_to_csv(path):
       xml_list = []
       for xml_file in glob.glob(path + '/*.xml'):
           print(xml_file)
           tree = ET.parse(xml_file)
           # print(root.find('filename').text)
    
           for member in root.findall('object'):
               value = (root.find('filename').text,
                        int(root.find('size')[0].text),  # width
                        int(root.find('size')[1].text),  # height
                        member[0].text,
                        int(member[4][0].text),
                        int(float(member[4][1].text)),
                        int(member[4][2].text),
                        int(member[4][3].text)
                        )
               xml_list.append(value)
       column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']
       xml_df = pd.DataFrame(xml_list, columns=column_name)
       return xml_df
    
    def main():
       # 这里是存放转换后的三个csv的位置
       csv_root = r"/home/zheshi/h4tv/datasets/VOCdevkit/VOCBall2k/VienDataset"
       # 这个是上一步存放分割后的三个xml文件夹的路径
       annotation_root = r"/home/zheshi/h4tv/datasets/VOCdevkit/VOCBall2k/VienDataset/Annotations"
       for directory in ['train', 'test', 'validation']:
           xml_path = os.path.join(annotation_root, directory)
           xml_df = xml_to_csv(xml_path)
           xml_df.to_csv(csv_root + '/ball_{}_labels.csv'.format(directory), index=None)
           print('Successfully converted xml to csv.')
    
    main()
  3. image和label数据转为tfrecord格式: generate_tfrecord.py

    from __future__ import division
    from __future__ import print_function
    from __future__ import absolute_import
    
    import os
    import io
    import pandas as pd
    import tensorflow as tf
    
    from PIL import Image
    from utils import dataset_util
    from collections import namedtuple, OrderedDict
    
    flags = tf.app.flags
    flags.DEFINE_string('csv_input', '', 'Path to the CSV input')
    flags.DEFINE_string('output_path', '', 'Path to output TFRecord')
    FLAGS = flags.FLAGS
    
    # 设置要检测的类型,如果再有就加上elif按照if的格式以此累加,
    # 例如elif row_label == 'shit': return 2
    def class_text_to_int(row_label, filename):
       if row_label == 'ball':
           return 1
       else:
           print("------------------nonetype:", filename)
           return None
    
    def split(df, group):
       data = namedtuple('data', ['filename', 'object'])
       gb = df.groupby(group)
       return [data(filename, gb.get_group(x)) for filename, x in zip(gb.groups.keys(), gb.groups)]
    
    def create_tf_example(group, path):
       with tf.io.gfile.GFile(os.path.join(path, '{}'.format(group.filename)), 'rb') as fid:
           encoded_jpg = fid.read()
       encoded_jpg_io = io.BytesIO(encoded_jpg)
       image = Image.open(encoded_jpg_io)
       width, height = image.size
    
       filename = group.filename.encode('utf8')
       image_format = b'png'
       xmins = []
       xmaxs = []
       ymins = []
       ymaxs = []
       classes_text = []
       classes = []
    
       for index, row in group.object.iterrows():
           xmins.append(row['xmin'] / width)
           xmaxs.append(row['xmax'] / width)
           ymins.append(row['ymin'] / height)
           ymaxs.append(row['ymax'] / height)
           classes_text.append(row['class'].encode('utf8'))
           classes.append(class_text_to_int(row['class'], group.filename))
    
       tf_example = tf.train.Example(features=tf.train.Features(feature={
           'image/height': dataset_util.int64_feature(height),
           'image/width': dataset_util.int64_feature(width),
           'image/filename': dataset_util.bytes_feature(filename),
           'image/source_id': dataset_util.bytes_feature(filename),
           'image/encoded': dataset_util.bytes_feature(encoded_jpg),
           'image/format': dataset_util.bytes_feature(image_format),
           'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),
           'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),
           'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),
           'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),
           'image/object/class/text': dataset_util.bytes_list_feature(classes_text),
           'image/object/class/label': dataset_util.int64_list_feature(classes),
       }))
       return tf_example
    
    def main(_):
       writer = tf.io.TFRecordWriter(FLAGS.output_path)
       # 训练用的图片的路径
       path = '/home/zheshi/h4tv/datasets/VOCdevkit/VOCBall2k/JPEGImages'
       examples = pd.read_csv(FLAGS.csv_input)
       grouped = split(examples, 'filename')
       num = 0
       for group in grouped:
           num += 1
           tf_example = create_tf_example(group, path)
           writer.write(tf_example.SerializeToString())
           if (num % 100 == 0):  # 每完成100个转换,打印一次
               print(num)
    
       writer.close()
       output_path = os.path.join(os.getcwd(), FLAGS.output_path)
       print('Successfully created the TFRecords: {}'.format(output_path))
    
    if __name__ == '__main__':
       tf.compat.v1.app.run()

执行代码:

python generate_tfrecord.py --csv_input=data/ball_train_labels.csv --output_path=data/ball_train.tfrecord

其中csv_input是之前转换的三个csv的路径,output_path是输出的tfrecord的路径,train、test、validation需要分别运行一次。

Config

在项目中创建一个存放配置文件的目录,比如命名为vien_data,然后在其目录下创建标签分类的配置文件label_map.pbtxt,如果需要检测多个,依次往下排,id依次+1

item {
  id: 1
  name: 'ball'
}

从项目的models\research\object_detection\samples\configs\ssd_mobilenet_v1_pets.config复制一份配置的模板文件到vien_data中,我们就命名为ssd_mobilenet_v1_ball.config好了,然后修改配置文件。

如果没有预训练的model文件,配置文件中fine_tune_checkpoint要设置为空。

然后还需要修改训练集和验证集的路径,在文件末尾,修改input_path为你的训练和测试集的tfrecord路径,label_map_path为上面创建的label_map.pbtxt的路径

train_input_reader: {
  tf_record_input_reader {
    input_path: "/home/zheshi/h4tv/datasets/VOCdevkit/VOCBall2k/VienDataset/ball_train.tfrecord"
  }
  label_map_path: "/home/zheshi/tensorflow/models/research/object_detection/vien_data/ball_label_map.pbtxt"
}

eval_config: {
  metrics_set: "coco_detection_metrics"
  num_examples: 1100
}

eval_input_reader: {
  tf_record_input_reader {
    input_path: "/home/zheshi/h4tv/datasets/VOCdevkit/VOCBall2k/VienDataset/ball_validation.tfrecord"
  }
  label_map_path: "/home/zheshi/tensorflow/models/research/object_detection/vien_data/ball_label_map.pbtxt"
  shuffle: false
  num_readers: 1
}

Train

在项目的research目录下执行(其中train_dir是训练出来的结果存放的路径,pipeline_config_path是上面复制修改的配置文件ssd_mobilenet_v1_ball.config路径):

python legacy/train.py --logtostderr \
                     --train_dir=/home/zheshi/tensorflow/models/research/object_detection/vien_train_gpu_models \
                     --pipeline_config_path=/home/zheshi/tensorflow/models/research/object_detection/vien_data/ssd_mobilenet_v1_ball.config

如果遇到ModuleNotFoundError: No module named 'object_detection' ,在research目录下执行

export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim
python setup.py build
python setup.py install

如果没有问题,加载配置后会开始训练,中间过程生成的文件和model都存在刚刚运行训练脚本时设置的vien_train_gpu_models目录中

可以查看图形化训练状态数据(修改logdir==training:后面的路径为执行训练脚本设置的vien_train_gpu_models目录):

 tensorboard --logdir==training:/home/zheshi/tensorflow/models/research/object_detection/vien_train_gpu_models --host=127.0.0.1

然后浏览器访问http://127.0.0.1:6006即可