TensorRT成功测试自己的数据集SSD模型二

经过前面的指导我们安装好环境后，现在就准备训练我自己的数据集的SSD模型，这里我跟TensorRT的例子保持一致，都用TensorFlow Object Detection API训练获取模型。那么开始详细介绍TensorFlow Object Detection API的安装以及SSD inception V2训练我自己的数据集。

TensorRT成功测试自己的数据集SSD模型一

TensorRT成功测试自己的数据集SSD模型二

TensorRT成功测试自己的数据集SSD模型三

TensorRT成功测试自己的数据集SSD模型四

首先新建一个名为TensorFlow文件夹

然后定位到该文件夹内

cd TensorFlow

下载TensorFlow Models

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

这里切记下载成功后要切换分支，为什么要切换，我自己实践的经验是用直接下载的model训练和测试是没有问题的，但是在用pb转uff的时候会出现_Cast node 错误，根据网上查到的解决办法将create_node_map中的Tofloat改成Cast,但是有会在create_engine时又会出现index error,就是用自己的五类目标训练得到的模型竟然出现预测的目标class_index为coco label中的index,真的很诡异。一直没能解决是哪里的问题，之后切换分支后就一切正常，所以一定要切换分支。

切换分支方法 

git checkout ae0a9409212d0072938fa60c9f85740bb89ced7e

到现在TensorFlow形式如下：

TensorFlow
└─ models
    ├── official
    ├── research
    ├── samples
    └── tutorials

安装Protobuf

pip install protobuf
cd models/research/
protoc object_detection/protos/*.proto --python_out=.

安装依赖包

# From within TensorFlow/models/research/
pip install .

添加research/slim 到PYTHONPATH

gedit ~/.bashrc
# From within tensorflow/models/research/
export PYTHONPATH=$PYTHONPATH:<PATH_TO_TF>/TensorFlow/models/research/slim
source ~/.bashrc

至此环境变量起作用了

安装COCO API

cd Tensorflow/models/research
git clone https://github.com/cocodataset/cocoapi.git
cd cocoapi/PythonAPI
make
cp -r pycocotools <PATH_TO_TF>/TensorFlow/models/research/

验证安装是否成功

cd TensorFlow\models\research\object_detection
jupyter notebook

运行object_detection_tutorial.ipynb

运行无误说明安装成功。

这里说明一点，如果你要运行model_builder_test，可能会出现有些例子运行失败，这里不用care这些失败，无视就好了，不影响接下来的训练。

python object_detection/builders/model_builder_test.py

---

接着正式开始我们的训练

首先，在TensorFlow下新建workspace/training_demo文件夹

TensorFlow
├─ models
│   ├─ official
│   ├─ research
│   ├─ samples
│   └─ tutorials
└─ workspace
    └─ training_demo
    	├─ annotations
		├─ images
		│   ├─ test
		│   └─ train
		├─ pre-trained-model
		├─ training

至于training_demo下的内容不需要新建，根据下面的步骤需要什么补充什么就行。

images是用来保存你的原始数据图像样本与注释文件的，现在要将其分类train和test，可以通过以下代码实现，9：1的比例也可以根据自身需要进行更改。

partition_dataset.py

""" usage: partition_dataset.py [-h] [-i IMAGEDIR] [-o OUTPUTDIR] [-r RATIO] [-x]

Partition dataset of images into training and testing sets

optional arguments:
  -h, --help            show this help message and exit
  -i IMAGEDIR, --imageDir IMAGEDIR
                        Path to the folder where the image dataset is stored. If not specified, the CWD will be used.
  -o OUTPUTDIR, --outputDir OUTPUTDIR
                        Path to the output folder where the train and test dirs should be created. Defaults to the same directory as IMAGEDIR.
  -r RATIO, --ratio RATIO
                        The ratio of the number of test images over the total number of images. The default is 0.1.
  -x, --xml             Set this flag if you want the xml annotation files to be processed and copied over.
"""
import os
import re
import shutil
from PIL import Image
from shutil import copyfile
import argparse
import glob
import math
import random
import xml.etree.ElementTree as ET


def iterate_dir(source, dest, ratio, copy_xml):
    source = source.replace('\\', '/')
    dest = dest.replace('\\', '/')
    train_dir = os.path.join(dest, 'train')
    test_dir = os.path.join(dest, 'test')

    if not os.path.exists(train_dir):
        os.makedirs(train_dir)
    if not os.path.exists(test_dir):
        os.makedirs(test_dir)

    images = [f for f in os.listdir(source)
              if re.search(r'([a-zA-Z0-9\s_\\.\-\(\):])+(.jpg|.jpeg|.png)$', f)]

    num_images = len(images)
    num_test_images = math.ceil(ratio*num_images)

    for i in range(num_test_images):
        idx = random.randint(0, len(images)-1)
        filename = images[idx]
        copyfile(os.path.join(source, filename),
                 os.path.join(test_dir, filename))
        if copy_xml:
            xml_filename = os.path.splitext(filename)[0]+'.xml'
            copyfile(os.path.join(source, xml_filename),
                     os.path.join(test_dir,xml_filename))
        images.remove(images[idx])

    for filename in images:
        copyfile(os.path.join(source, filename),
                 os.path.join(train_dir, filename))
        if copy_xml:
            xml_filename = os.path.splitext(filename)[0]+'.xml'
            copyfile(os.path.join(source, xml_filename),
                     os.path.join(train_dir, xml_filename))


def main():

    # Initiate argument parser
    parser = argparse.ArgumentParser(description="Partition dataset of images into training and testing sets",
                                     formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument(
        '-i', '--imageDir',
        help='Path to the folder where the image dataset is stored. If not specified, the CWD will be used.',
        type=str,
        default=os.getcwd()
    )
    parser.add_argument(
        '-o', '--outputDir',
        help='Path to the output folder where the train and test dirs should be created. '
             'Defaults to the same directory as IMAGEDIR.',
        type=str,
        default=None
    )
    parser.add_argument(
        '-r', '--ratio',
        help='The ratio of the number of test images over the total number of images. The default is 0.1.',
        default=0.1,
        type=float)
    parser.add_argument(
        '-x', '--xml',
        help='Set this flag if you want the xml annotation files to be processed and copied over.',
        action='store_true'
    )
    args = parser.parse_args()

    if args.outputDir is None:
        args.outputDir = args.imageDir

    # Now we are ready to start the iteration
    iterate_dir(args.imageDir, args.outputDir, args.ratio, args.xml)


if __name__ == '__main__':
    main()

python partition_dataser.py -x -i training_demo\images -r 0.1

创建标签文件

training_demo\annotations\label_map.pbtxt

item {
    id: 1
    name: 'dirty'
}

item {
    id: 2
    name: 'oil'
}

item {
    id: 3
    name: 'pit'
}

item {
    id: 4
    name: 'scratch'
}

item {
    id: 5
    name: 'wire_drawing'
}

生成TensorFlow Records

在TensorFlow下新建scripts\preprocessing文件夹

xml 转 csv

scripts\preprocessing\xml_to_csv.py

"""
Usage:
# Create train data:
python xml_to_csv.py -i [PATH_TO_IMAGES_FOLDER]/train -o [PATH_TO_ANNOTATIONS_FOLDER]/train_labels.csv

# Create test data:
python xml_to_csv.py -i [PATH_TO_IMAGES_FOLDER]/test -o [PATH_TO_ANNOTATIONS_FOLDER]/test_labels.csv
"""

import os
import glob
import pandas as pd
import argparse
import xml.etree.ElementTree as ET


def xml_to_csv(path):
    """Iterates through all .xml files (generated by labelImg) in a given directory and combines them in a single Pandas datagrame.

    Parameters:
    ----------
    path : {str}
        The path containing the .xml files
    Returns
    -------
    Pandas DataFrame
        The produced dataframe
    """

    xml_list = []
    for xml_file in glob.glob(path + '/*.xml'):
        tree = ET.parse(xml_file)
        root = tree.getroot()
        for member in root.findall('object'):
            value = (root.find('filename').text,
                    int(root.find('size')[0].text),
                    int(root.find('size')[1].text),
                    member[0].text,
                    int(member[4][0].text),
                    int(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():
    # Initiate argument parser
    parser = argparse.ArgumentParser(
        description="Sample TensorFlow XML-to-CSV converter")
    parser.add_argument("-i",
                        "--inputDir",
                        help="Path to the folder where the input .xml files are stored",
                        type=str)
    parser.add_argument("-o",
                        "--outputFile",
                        help="Name of output .csv file (including path)", type=str)
    args = parser.parse_args()

    if(args.inputDir is None):
        args.inputDir = os.getcwd()
    if(args.outputFile is None):
        args.outputFile = args.inputDir + "/labels.csv"

    assert(os.path.isdir(args.inputDir))

    xml_df = xml_to_csv(args.inputDir)
    xml_df.to_csv(
        args.outputFile, index=None)
    print('Successfully converted xml to csv.')


if __name__ == '__main__':
    main()

cd TensorFlow/scripts/preprocessing

# Create train data:
python xml_to_csv.py -i [PATH_TO_IMAGES_FOLDER]/train -o [PATH_TO_ANNOTATIONS_FOLDER]/train_labels.csv

# Create test data:
python xml_to_csv.py -i [PATH_TO_IMAGES_FOLDER]/test -o [PATH_TO_ANNOTATIONS_FOLDER]/test_labels.csv

# For example
# python xml_to_csv.py -i C:\Users\sglvladi\Documents\TensorFlow\workspace\training_demo\images\train -o C:\Users\sglvladi\Documents\TensorFlow\workspace\training_demo\annotations\train_labels.csv
# python xml_to_csv.py -i C:\Users\sglvladi\Documents\TensorFlow\workspace\training_demo\images\test -o C:\Users\sglvladi\Documents\TensorFlow\workspace\training_demo\annotations\test_labels.csv

csv 转 record

TensorFlow\scripts\preprocessing\generate_tfrecord.py

"""
Usage:

# Create train data:
python generate_tfrecord.py --label=<LABEL> --csv_input=<PATH_TO_ANNOTATIONS_FOLDER>/train_labels.csv  --output_path=<PATH_TO_ANNOTATIONS_FOLDER>/train.record

# Create test data:
python generate_tfrecord.py --label=<LABEL> --csv_input=<PATH_TO_ANNOTATIONS_FOLDER>/test_labels.csv  --output_path=<PATH_TO_ANNOTATIONS_FOLDER>/test.record
"""

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
import sys
sys.path.append("../../models/research")

from PIL import Image
from object_detection.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.DEFINE_string('label0', '', 'Name of class label')
# if your image has more labels input them as
flags.DEFINE_string('label0', '', 'Name of class[0] label')
flags.DEFINE_string('label1', '', 'Name of class[1] label')
flags.DEFINE_string('label2', '', 'Name of class[2] label')
flags.DEFINE_string('label3', '', 'Name of class[3] label')
flags.DEFINE_string('label4', '', 'Name of class[4] label')
# and so on.
flags.DEFINE_string('img_path', '', 'Path to images')
FLAGS = flags.FLAGS


# TO-DO replace this with label map
# for multiple labels add more else if statements
def class_text_to_int(row_label):
    # if row_label == FLAGS.label:  # 'ship':
    #     return 1
    # comment upper if statement and uncomment these statements for multiple labelling
    if row_label == FLAGS.label0:
      return 1
    elif row_label == FLAGS.label1:
      return 2
    elif row_label == FLAGS.label2:
      return 3      
    elif row_label == FLAGS.label3:
      return 4
    elif row_label == FLAGS.label4:
      return 5
    else:
        return 0


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.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'jpg'
    # check if the image format is matching with your images.
    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']))

    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.python_io.TFRecordWriter(FLAGS.output_path)
    path = os.path.join(os.getcwd(), FLAGS.img_path)
    examples = pd.read_csv(FLAGS.csv_input)
    grouped = split(examples, 'filename')
    for group in grouped:
        tf_example = create_tf_example(group, path)
        writer.write(tf_example.SerializeToString())

    writer.close()
    output_path = os.path.join(os.getcwd(), FLAGS.output_path)
    print('Successfully created the TFRecords: {}'.format(output_path))


if __name__ == '__main__':
    tf.app.run()

cd TensorFlow\scripts\preprocessing

# Create train data:
python generate_tfrecord.py --label0=dirty --label1=oil --label2=pit --label3=scratch --label4=wire_drawing --csv_input=/home/xxx/Tensorflow_new/workspace/training_demo/annotations/train_labels.csv --output_path=/home/xxx/Tensorflow_new/workspace/training_demo/annotations/train.record --img_path=/home/xxx/Tensorflow_new/workspace/training_demo/images/train

# Create test data:
python generate_tfrecord.py --label0=dirty --label1=oil --label2=pit --label3=scratch --label4=wire_drawing --csv_input=/home/xxx/Tensorflow_new/workspace/training_demo/annotations/test_labels.csv --output_path=/home/xxx/Tensorflow_new/workspace/training_demo/annotations/test.record --img_path=/home/xxx/Tensorflow_new/workspace/training_demo/images/test

之后会在training_demo/annotations文件夹下生成test.record 和 train.record 即为转换成功。

创建训练Pipeline

首先下载预训练模型选择ssd_inception_v2_coco下载

将下载好的模型解压到 pre-trained-models

training_demo/
├─ ...
├─ pre-trained-models/
│  └─ ssd_inception_v2_coco_2018_01_28/
│     ├─ saved_model/
│       ├─ pipeline.config
│     └─ ...

创建pipeline.config

1.复制training_demo/pre-trained-models/ssd_inception_v2_coco_2018_01_28/pipeline.config到training_demo/training

2.打开training_demo/training下的 pipeline.config ；

training_demo/training/pipeline.config

# SSD with Inception v2 configuration for MSCOCO Dataset.
# Users should configure the fine_tune_checkpoint field in the train config as
# well as the label_map_path and input_path fields in the train_input_reader and
# eval_input_reader. Search for "PATH_TO_BE_CONFIGURED" to find the fields that
# should be configured.

model {
    ssd {
        num_classes: 5 # Set this to the number of different label classes
        box_coder {
            faster_rcnn_box_coder {
                y_scale: 10.0
                x_scale: 10.0
                height_scale: 5.0
                width_scale: 5.0
            }
        }
        matcher {
            argmax_matcher {
                matched_threshold: 0.5
                unmatched_threshold: 0.5
                ignore_thresholds: false
                negatives_lower_than_unmatched: true
                force_match_for_each_row: true
            }
        }
        similarity_calculator {
            iou_similarity {
            }
        }
        anchor_generator {
            ssd_anchor_generator {
                num_layers: 6
                min_scale: 0.2
                max_scale: 0.95
                aspect_ratios: 1.0
                aspect_ratios: 2.0
                aspect_ratios: 0.5
                aspect_ratios: 3.0
                aspect_ratios: 0.3333
                reduce_boxes_in_lowest_layer: true
            }
        }
        image_resizer {
            fixed_shape_resizer {
                height: 300
                width: 300
            }
        }
        box_predictor {
            convolutional_box_predictor {
                min_depth: 0
                max_depth: 0
                num_layers_before_predictor: 0
                use_dropout: false
                dropout_keep_probability: 0.8
                kernel_size: 3
                box_code_size: 4
                apply_sigmoid_to_scores: false
                conv_hyperparams {
                activation: RELU_6,
                regularizer {
                    l2_regularizer {
                        weight: 0.00004
                    }
                }
                initializer {
                        truncated_normal_initializer {
                            stddev: 0.03
                            mean: 0.0
                        }
                    }
                }
            }
        }
        feature_extractor {
            type: 'ssd_inception_v2' # Set to the name of your chosen pre-trained model
            min_depth: 16
            depth_multiplier: 1.0
            conv_hyperparams {
                activation: RELU_6,
                regularizer {
                    l2_regularizer {
                        weight: 0.00004
                    }
                }
                initializer {
                    truncated_normal_initializer {
                        stddev: 0.03
                        mean: 0.0
                    }
                }
                batch_norm {
                    train: true,
                    scale: true,
                    center: true,
                    decay: 0.9997,
                    epsilon: 0.001,
                }
            }
            override_base_feature_extractor_hyperparams: true
        }
        loss {
            classification_loss {
                weighted_sigmoid {
                }
            }
            localization_loss {
                weighted_smooth_l1 {
                }
            }
            hard_example_miner {
                num_hard_examples: 3000
                iou_threshold: 0.99
                loss_type: CLASSIFICATION
                max_negatives_per_positive: 3
                min_negatives_per_image: 0
            }
            classification_weight: 1.0
            localization_weight: 1.0
        }
        normalize_loss_by_num_matches: true
        post_processing {
            batch_non_max_suppression {
                score_threshold: 1e-8
                iou_threshold: 0.6
                max_detections_per_class: 100
                max_total_detections: 100
            }
            score_converter: SIGMOID
        }
    }
}

train_config: {
    batch_size: 24 # Increase/Decrease this value depending on the available memory (Higher values require more memory and vice-versa)
    optimizer {
        rms_prop_optimizer: {
            learning_rate: {
                exponential_decay_learning_rate {
                    initial_learning_rate: 0.004
                    decay_steps: 800720
                    decay_factor: 0.95
                }
            }
            momentum_optimizer_value: 0.9
            decay: 0.9
            epsilon: 1.0
        }
    }
    fine_tune_checkpoint: "pre-trained-model/model.ckpt" # Path to extracted files of pre-trained model
    from_detection_checkpoint: true
    # Note: The below line limits the training process to 200K steps, which we
    # empirically found to be sufficient enough to train the pets dataset. This
    # effectively bypasses the learning rate schedule (the learning rate will
    # never decay). Remove the below line to train indefinitely.
    num_steps: 200000
    data_augmentation_options {
        random_horizontal_flip {
        }
    }
    data_augmentation_options {
        ssd_random_crop {
        }
    }
}

train_input_reader: {
    tf_record_input_reader {
        input_path: "annotations/train.record" # Path to training TFRecord file
    }
    label_map_path: "annotations/label_map.pbtxt" # Path to label map file
}

eval_config: {
    # (Optional): Uncomment the line below if you installed the Coco evaluation tools
    # and you want to also run evaluation
    # metrics_set: "coco_detection_metrics"
    # (Optional): Set this to the number of images in your <PATH_TO_IMAGES_FOLDER>/train
    # if you want to also run evaluation
    num_examples: 8000
    # Note: The below line limits the evaluation process to 10 evaluations.
    # Remove the below line to evaluate indefinitely.
    max_evals: 10
}

eval_input_reader: {
    tf_record_input_reader {
        input_path: "annotations/test.record" # Path to testing TFRecord
    }
    label_map_path: "annotations/label_map.pbtxt" # Path to label map file
    shuffle: false
    num_readers: 1
}

开始训练

1.复制TensorFlow/models/research/object_detection/train.py到training_demo文件夹

2.cd training_demo, 运行

python train.py --logtostderr --train_dir=training/ --pipeline_config_path=training/pipeline.config

最终在training文件夹保存到训练所得到的模型。