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即可