TensorFlow 使用经过训练的模型进行检测

Question

我正在为个人项目开发废物/垃圾检测器。我依靠 Tensorflow（在 Python 3 中）来训练我自己的数据集。

我有一个从头开始创建和训练模型的脚本。然后，我冻结检查点以获取 PB 文件进行检测。

我的检测代码 (found here) 需要两个文件才能工作：之前的 PB 文件和 labelmap.txt。

# Path to frozen detection graph. This is the actual model that is used for the object detection.
PATH_TO_CKPT = 'frozen_inference_graph.pb'

# List of the strings that is used to add correct label for each box.
PATH_TO_LABELS = 'label_map.pbtxt'

我知道 labelmap.txt 的样子，实际上自己编写很简单，但我不知道如何生成它，因为它将每个类都链接到一个 ID，而我不知道 ID。

我尝试在网上搜索，当人们提到labelmap.txt时，它涉及到Tfrecords。但是，我的项目没有使用 Tf 记录，我提取每个感兴趣的区域并将它们保存在子文件夹中，一个类的子文件夹（罐头，瓶子......）。

由于我是 Tensorflow 的新手，我可能在训练过程中误解了一些东西。你有任何线索，所以我可以通过测试来看看我的模型是否准确？如果您需要，我可以提供一些代码。

提前谢谢你，

Answer 1

labelmap.pbtxt 文件将网络内部使用的 ID 映射到标签名称。你不能简单地在训练后生成一个。您需要确保在训练期间使用相同的 ID-标签映射，否则您可能会得到不正确的结果。

如果您使用 tensorflow object_detection 模型的训练说明，那么您将在某个时候生成此标签映射文件，您可以重复使用它。

查看您用于训练网络的步骤或在此处发布。

Answer 2

在训练之前，我收集并标记了数千张图像，提取了每个标记区域，调整了每个区域的大小，并根据它们的类别将它们拆分到不同的文件夹中。

训练步骤涉及多个文件。我最初是从这个repository 中检索代码并添加了恢复训练的可能性。

trainer.py

import os
import tensorflow as tf
import model_architecture

from utils import utils
from build_model import model_tools

# Images directory.
data_path = os.path.join('dataset' + os.sep)# contains subfolders, one per item
all_classes = os.listdir(data_path)
number_of_classes = len(all_classes)

# Images dimensions.
height = 64
width = 64

# Checkpoints directory.
output_dir = os.path.join(os.pardir + os.sep, 'checkpoints' + os.sep)
model_pattern = 'model.ckpt'
model_base_path = os.path.join(output_dir, model_pattern)
meta_file_path = model_base_path + '.meta'

# Training params.
color_channels = 3
start = 0
epochs = 5
batch_size = 10
batch_counter = 0

# Create Placeholders for images and labels.
images_ph = tf.placeholder(tf.float32, shape=[None, height, width, color_channels])
labels_ph = tf.placeholder(tf.float32, shape=[None, number_of_classes])


def trainer(network, number_of_images):
    cross_entropy = tf.nn.softmax_cross_entropy_with_logits_v2(logits=network, labels=labels_ph)
    cost = tf.reduce_mean(cross_entropy)
    optimizer = tf.train.AdamOptimizer().minimize(cost)
    tf.summary.scalar('cost', cost)
    tf.add_to_collection('optimizer', optimizer)

    global_step = tf.Variable(0, name='global_step', trainable=False)
    saver = tf.train.Saver()

    # Launch the graph in a session
    with tf.Session() as sess:
        # Initialize all variables.
        tf.global_variables_initializer().run()

        # Read checkpoints directory.
        ckpt = tf.train.get_checkpoint_state(output_dir)

        if ckpt and ckpt.model_checkpoint_path:
            saver.restore(sess, ckpt.model_checkpoint_path)
            print('Reloading existing model.')
        else:
            init = tf.global_variables_initializer()
            sess.run(init)
            print('Creating a new model.')

        # Get last epoch index.
        start = global_step.eval()

        writer = tf.summary.FileWriter(output_dir, graph=tf.get_default_graph())
        merged = tf.summary.merge_all()
        saver = tf.train.Saver(write_version=tf.train.SaverDef.V2, max_to_keep=5)
        counter = 0

        # Training.
        for epoch in range(start, epochs):
            tools = utils()
            for batch in range(int(number_of_images / batch_size)):
                counter += 1
                images, labels = tools.batch_dispatch()
                if images is None:
                    break
                loss, summary = sess.run([cost, merged], feed_dict={images_ph: images, labels_ph: labels})
                sess.run(optimizer, feed_dict={images_ph: images, labels_ph: labels})
                print('Epoch number {epoch} batch {batch} complete - loss {loss}'.format(
                    epoch=epoch, batch=batch, loss=loss))
                writer.add_summary(summary, counter)
            global_step.assign(epoch).eval()

            # Save progression.
            saver.save(sess, model_base_path, global_step=epoch)


# Main program.
if __name__ == '__main__':
    tools = utils()
    model = model_tools()
    network = model_architecture.generate_model(images_ph, number_of_classes)
    number_of_images = sum([len(files) for r, d, files in os.walk('dataset')])
    trainer(network, number_of_images)

model_tools.py

class model_tools:

    def add_weights(self, shape):
        return tf.Variable(tf.truncated_normal(shape=shape, stddev=0.05))

    def add_biases(self, shape):
        return tf.Variable(tf.constant(0.05, shape=shape))

    def conv_layer(self, layer, kernel, input_shape, output_shape, stride_size):
        weights = self.add_weights([kernel, kernel, input_shape, output_shape])
        biases = self.add_biases([output_shape])
        stride = [1, stride_size, stride_size, 1]
        layer = tf.nn.conv2d(layer, weights, strides=stride, padding='SAME') + biases
        return layer

    def pooling_layer(self, layer, kernel_size, stride_size):
        kernel = [1, kernel_size, kernel_size, 1]
        stride = [1, stride_size, stride_size, 1]
        return tf.nn.max_pool(layer, ksize=kernel, strides=stride, padding='SAME')

    def flattening_layer(self, layer):
        input_size = layer.get_shape().as_list()
        new_size = input_size[-1] * input_size[-2] * input_size[-3]
        return tf.reshape(layer, [-1, new_size]), new_size

    def fully_connected_layer(self, layer, input_shape, output_shape):
        weights = self.add_weights([input_shape, output_shape])
        biases = self.add_biases([output_shape])
        layer = tf.matmul(layer, weights) + biases
        return layer

    def activation_layer(self, layer):
        return tf.nn.relu(layer)

utils.py

import cv2
import random    

class utils:
    image_count = []
    count_buffer = []
    class_buffer = all_classes[:]

    def __init__(self):
        self.image_count = []
        self.count_buffer = []
        for i in os.walk(data_path):
            if len(i[2]):
                self.image_count.append(len(i[2]))
        self.count_buffer = self.image_count[:]

    def batch_dispatch(self, batch_size=batch_size):
        global batch_counter
        if sum(self.count_buffer):

            class_name = random.choice(self.class_buffer)
            choice_index = all_classes.index(class_name)
            choice_count = self.count_buffer[choice_index]
            if choice_count == 0:
                class_name = all_classes[self.count_buffer.index(max(self.count_buffer))]
                choice_index = all_classes.index(class_name)
                choice_count = self.count_buffer[choice_index]

            slicer = batch_size if batch_size < choice_count else choice_count
            img_ind = self.image_count[choice_index] - choice_count
            indices = [img_ind, img_ind + slicer]
            images = self.generate_images(class_name, indices)
            labels = self.generate_labels(class_name, slicer)

            self.count_buffer[choice_index] = self.count_buffer[choice_index] - slicer
        else:
            images, labels = (None,) * 2
        return images, labels

    def generate_labels(self, class_name, number_of_samples):
        one_hot_labels = [0] * number_of_classes
        one_hot_labels[all_classes.index(class_name)] = 1
        one_hot_labels = [one_hot_labels] * number_of_samples
        return one_hot_labels

    def generate_images(self, class_name, indices):
        batch_images = []
        choice_folder = os.path.join(data_path, class_name)
        selected_images = os.listdir(choice_folder)[indices[0]:indices[1]]
        for image in selected_images:
            img = cv2.imread(os.path.join(choice_folder, image))
            batch_images.append(img)
        return batch_images

model_architecture.py 包含 3 层图像分类器的结构。

当我运行 trainer.py 时，我得到一个包含元和索引文件的检查点文件夹。看来是对的。

关于导出模型，我很尴尬，因为我不知道该为管道配置路径提供什么参数。

python3 export_inference_graph.py \ --input_type image_tensor \ --trained_checkpoint_prefix "/home/user/model/model.ckpt-4" \ --pipeline_config_path ???? \ --output_directory /home/user/exports/

为了获取 PB 文件，我使用了这个：

checkpoint_location = 'checkpoints/model.ckpt-0' 
export_dir = 'frozen/'
loaded_graph = tf.Graph()
with tf.Session(graph=loaded_graph) as sess:
    loader = tf.train.import_meta_graph(checkpoint_location+ '.meta')
    loader.restore(sess, checkpoint_location)

    builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
    builder.add_meta_graph_and_variables(sess,
                                         [tf.saved_model.tag_constants.TRAINING],
                                         strip_default_attrs=True)
builder.add_meta_graph([tf.saved_model.tag_constants.SERVING], strip_default_attrs=True)
builder.save()

它会创建一个 save_model.pb 文件，但不会创建 labelmap.pbtxt。

我应该完全改变我训练模型的方式吗？