Tensorflow use trained model for detection - python-3.x

Question

I'm working on a waste/garbage detector for a personal project. I rely on Tensorflow (in Python 3) to train my own dataset.
I have a script that creates and trains a model from scratch. Then, I freeze the checkpoints to get a PB file for detection.
The code I have for the detection (found here) requires two files to work: the previous PB file and a 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'
I know how a labelmap.txt looks like and it is quite simple to write it myself actually, but I don't know how to generate it because it links each class to an ID and the ID is unknown to me.
I tried to search on the Internet, when people mention the labelmap.txt, it involves Tfrecords. However, I don't use Tf records for my project, I extract each region of interest and save them in subfolders, one subfolder for a class (can, bottle...).
As I am new to Tensorflow, I may have misunderstood something in the training process. Do you have any lead so I can see if my model is accurate by testing it ? I can provide some codes if you need it.
Thanking you in advance,

Answer 1

The labelmap.pbtxt file maps the IDs used internally in the network to the label names. You cannot simply generate one after training. You need to make sure to use same ID-label mapping was used during training or you might get incorrect results.
If you use the training instructions for the tensorflow object_detection model then you will have generated this labelmap-file at some point and you can just re-use it.
Check out the steps you used to train the network or post them here.

Answer 2

Before training, I gathered and labelled thousands of images, extracted each labelled area, resized each of them and, according to their classes, I splitted them in different folders.
There are several files involved in the training step. I originally retrieved the code from this repository and added the possibility to resume training.
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 contains the structure of the 3 layered Image classifier.
When I run trainer.py, I get a checkpoints folder filled with meta and index files. It seems correct.
About exporting the model, I'm embarrassed as I don't know what to give as parameter for the pipeline config path.
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/
To get the PB file, I used this:
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()
It creates a save_model.pb file but not a labelmap.pbtxt.
Should I completely change the way I train my model ?