9. Beyond Classification

Written by Matthijs Hollemans

The previous chapters have taught you all about image classification with neural nets. But neural networks can be used for many other computer vision tasks. In this chapter and the next, you’ll look at two advanced examples:

• Object detection: find multiple objects in an image.
• Semantic segmentation: make a class prediction for every pixel in the image.

Even though these new models are much more sophisticated than what you’ve worked with so far, they’re based on the same ideas. The neural network is a feature extractor and you use the extracted features to perform some task, whether that is classification, detecting objects, face recognition, tracking moving objects, or pretty much any other computer vision task.

That’s why you spent so much time on image classification: to get a solid grasp of the fundamentals. But now it’s time to take things a few steps further…

Where is it?

Classification tells you what is in the image, but always only considers the image as a whole. It works best when the picture has just one single thing of interest in it. If your classifier is trained to tell apart cats and dogs, and the image contains both a cat and a dog, then the answer is anyone’s guess.

An object detection model has no problem dealing with such images. The goal of object detection is to find all the objects inside an image, even if they are of different types. You can think of it as a classifier for specific image regions.

The object detector not only finds what the objects are but also where they are located in the image. It does this by predicting one or more bounding boxes, which are simply rectangular regions in the image.

A bounding box is described by four numbers, representing either the corner points of the rectangle or the center point plus a width and height:

Both types are used in practice, but this chapter uses the one with the corner points.

Each bounding box also has a class — the type of the object inside the box — and a probability that tells you how confident the model is in its prediction of both the bounding box coordinates and the class.

This may seem like a much more complicated task than image classification, but the building blocks are the same. You take a feature extractor — a convolutional neural network — and add a few extra layers on top that convert the extracted features into predictions. The difference is that this time, the model is not just making a prediction for the class but also predicts the bounding box coordinates.

Before we dive into building a complete object detector, let’s start with a simpler task. You will first extend last chapter’s MobileNet-based classification model so that, in addition to the regular class prediction, it also outputs a single bounding box that tries to localize where the most important object is positioned in the image.

Just predict one bounding box, how hard could it be? (Answer: It’s actually easier than you might think.)

The ground-truth will set you free

First, we should revisit the dataset.

Even though this new neural network will now make a different kind of prediction, the training procedure is still the same: You provide a dataset that consists of the images and the targets. You also provide a suitable loss function that calculates how wrong the model’s predictions are by comparing them to the targets. Then you use a Stochastic Gradient Descent optimizer, such as Adam, to find the values for the model’s learnable parameters that make the loss value as small as possible. Been there, done that.

No matter what task your neural network performs, whether it’s predicting classes or bounding boxes — or the weather or stock prices or anything else — the training process is always the same. However, each task needs its own kind of training data. And for object detection tasks, the training data must contain bounding box information.

Previously, the targets were just the class names for the images, but now they must also include the so-called ground-truth bounding boxes that tell you where the objects are located inside the training images. Without these bounding box annotations, the loss function wouldn’t be able to calculate how wrong the model is, and training the model to predict bounding boxes would be impossible.

We have provided the bounding box annotations for the snacks dataset as a set of CSV files. To get a feel for how they work, you’ll now take a closer look at those annotations. Create a new Jupyter notebook or follow along with final/Localization.ipynb from this chapter’s resources.

Note: As before, you’ll be working with the kerasenv Python environment. Set up this environment with Anaconda Navigator or with conda create if necessary. If you don’t already have it from previous chapters, download the snacks dataset by double-clicking starter/snacks-download-link.webloc and unzip this file. It contains the images on which you’ll train the model, including the ground-truth annotations.

The easiest way to deal with CSV files in Python is by using the Pandas library. As usual, first import the needed packages — NumPy, Matplotlib and Pandas — and define the paths to where you downloaded the dataset:

import os, sys
import numpy as np
import pandas as pd

%matplotlib inline
import matplotlib.pyplot as plt

data_dir = "snacks"
train_dir = os.path.join(data_dir, "train")
val_dir = os.path.join(data_dir, "val")
test_dir = os.path.join(data_dir, "test")

Then load the annotations-train.csv file into a new Pandas DataFrame object:

path = os.path.join(data_dir, "annotations-train.csv")
train_annotations = pd.read_csv(path)
train_annotations.head()

The new dataframe train_annotations literally contains the exact same data as the CSV file. Pandas offers a lot of useful functions to manipulate this data. It’s like using the functionality of SQL and Excel but inside Python, which is awesome if you’re into that sort of thing.

The train_annotations.head() command gives as output the “head” of this dataframe, which is the first five rows:

The dataframe is actually much bigger: When you do len(train_annotations) it should print 7040. The dataframe has one row for each annotation. There are only about 4,800 images in the training set but some pictures have multiple objects in them — that’s why there are more annotations than training images.

Some training images even contain objects of different types. The first three rows in the dataframe all belong to the same image, a picture of a cake, but apparently, there’s also some ice cream in that image (see rows 1 and 2).

There are also a number of images in the training set that do not have bounding box annotations at all.

The coordinates of the bounding box are given by four numbers: x_min, x_max, y_min and y_max. The top-left corner of the box is (x_min, y_min), the bottom-right corner is (x_max, y_max). These are floating-point values — or “real-valued” numbers in math speak — between 0 and 1, also known as normalized coordinates.

It’s convenient to use normalized coordinates because it makes them independent of the actual size of the image. This is important: remember that we scale down images to 224×224 pixels during training. If the bounding box coordinates were given in pixels as well, you’d have to remember to scale these down by the same amount… it gets messy really quick. With normalized coordinates, you don’t have to worry about this.

The dataframe has two columns that contain class names: class_name, which is the class of the object inside this bounding box, and folder, which is where the image is stored in the dataset. folder is also the name of the class you used for training the classifier in the previous chapters. From now on, you’ll only use the class_name for training, but you still need folder to know whence to load the image file.

While you’re at it, you might as well load the annotations for the validation and test sets into their own dataframes. These each have about 1,400 rows:

val_annotations = pd.read_csv(os.path.join(data_dir,
                                   "annotations-val.csv"))
test_annotations = pd.read_csv(os.path.join(data_dir,
                                   "annotations-test.csv"))

Show me the data!

Now, let’s have a proper look at these bounding boxes. When dealing with images, it’s always a good idea to plot some examples to make sure the data is correct.

Remember the old adage, “Garbage in equals garbage out.” If you’re training your model on data that doesn’t make sense, then neither will the model’s predictions and you just wasted a lot of time and electricity. Don’t be that person!

The code for plotting the images isn’t terribly exciting, and so we’ve hidden this away in a file helpers.py that you can find in this chapter’s downloads. It’s a good idea to keep your notebook clean and put big functions and reusable code in separate Python files.

Copy helpers.py into the same folder that your Jupyter notebook is in, and then write:

image_width = 224
image_height = 224

from helpers import plot_image

This imports the plot_image function from the helpers.py module. plot_image() takes as arguments an image and a list of one or more bounding boxes and then draws the bounding boxes on top of the image.

Feel free to have a look inside helpers.py to see how this function works. You can also run plot_image? in a new cell to see its documentation, or plot_image?? to see the full source code.

To get a single row from the dataframe, you can write the following:

train_annotations.iloc[0]

Here, 0 is the row index so this returns the fields from the first row:

image_id      009218ad38ab2010
x_min                  0.19262
x_max                 0.729831
y_min                 0.127606
y_max                 0.662219
class_name                cake
folder                    cake
Name: 0, dtype: object

This is a so-called Pandas Series object and you can index it by name to get any of these fields, just like you would a dictionary. Now, grab an image from a single row in the dataframe and plot it together with its bounding box:

from keras.preprocessing import image

def plot_image_from_row(row, image_dir):
    # Load the image from "folder/image_id.jpg"
    image_path = os.path.join(image_dir, row["folder"],
                              row["image_id"] + ".jpg")
    img = image.load_img(image_path,
                    target_size=(image_width, image_height))

    # Put the box coordinates and class name into a tuple
    bbox = (row["x_min"], row["x_max"],
            row["y_min"], row["y_max"], row["class_name"])

    # Draw the bounding box on top of the image
    plot_image(img, [bbox])

Now, call this new function to make the plot for a given annotation:

annotation = train_annotations.iloc[0]
plot_image_from_row(annotation, train_dir)

This draws the following image and bounding box (on the left):

You can see that the bounding box for the “cake” annotation neatly fits around the actual slice of cake in the picture. So it looks like the data is loaded correctly!

This training image actually has three annotations. The other two are for the ice cream dessert in the top-right corner of the photo. On the right is shown the annotation from row 2. The bounding box from row 1 is very similar and covers the same object.

In the Google Open Images dataset that these images and annotations come from, often the same object in the image has multiple annotations, created by different people. That doesn’t appear to be a problem, as long as these annotations aren’t too different. After all, more training data is usually better.

However, many images have fewer annotations than there are objects, which is not ideal. For example, the image at index 3,500 in the train_annotations dataframe, with image_id 0c429e9be7f72342, has four strawberries but only three annotations, two of which are for the same strawberry. Ideally, this image would have a unique annotation for each individual object.

To get a feel for what the dataset is like, have a look at some of the other images from the training, validation and test annotations.

Because not all objects from all images have annotations, and some have duplicates, this dataset isn’t ideal — but, with over 7,000 annotations, it should still be good enough to train a decent object detection model. When you start building your own models, you’ll find that you’ll be spending a lot of time cleaning up your training data, filling in missing values, and so on. Your model will only ever be as good as the quality of the dataset, so it’s worth putting in the time.

What about images without annotations?

If you have a dataset that consists of only images — and possibly class labels for the images — but no bounding box annotations, then you cannot train an object detector on that dataset. Not gonna happen; ain’t no two ways about it.

First, you’ll have to create the bounding box annotations for each image. This can be a time-consuming process, especially since you need lots of images, but fortunately there are tools that can help. A few suggestions:

• RectLabel, available on the Mac App Store. This is a powerful tool with many options, but it expects the annotations to be provided as a separate XML file for each image. This is not unusual — it’s how the popular Pascal VOC dataset does things — but it won’t be able to handle our CSV files. If you’re getting serious about training your own object detectors, definitely give this tool a try.
• Labelbox at labelbox.io is an online tool for labeling training data for many different tasks, including object detection. This is a paid service but there is a free tier.
• Simple Image Annotator from github.com/sgp715 is a Python program that runs as a local web service. As its name implies, it’s pretty simple to use and offers only basic editing features. The output is a CSV file but it’s not 100% compatible with the CSV format we’re using.
• Sloth, which is available at sloth.readthedocs.io, and is an advanced labeling tool. Requires Linux.
• CVAT, or Computer Vision Annotation Tool, which is available at github.com/opencv/cvat.

This is by no means an exhaustive list, and new annotation tools and services are springing up left and right.

There are about 800 images in the snacks dataset that do not have annotations. For the purposes of this book, you’re just going to ignore those images and only train on the images that already do have annotations. But, if you’re bored at home on a rainy Sunday afternoon and you feel like labeling the remaining images, don’t let us stop you.

Note: We just mentioned that RectLabel uses a different format for storing the annotations (XML) and that Simple Image Annotator does use a CSV file but with different fields. Some of the other tools output JSON files. This sort of thing is common. Every dataset will store its data in a slightly different way, and you’ll often find yourself writing small Python scripts to convert data from one format to the other. A large part of any machine-learning project consists of finding data, cleaning it up and annotating it. Once the data is in the format you want, doing the actual machine learning is usually quite straightforward.

Your own generator

Previously, you used ImageDataGenerator and flow_from_directory() to automatically load the images and put them into batches for training. That is convenient when your images are neatly organized into folders, but the new training data consists of a Pandas DataFrame with bounding box annotations. You’ll need a way to read the rows from this dataframe into a batch. Fortunately, Keras lets you write your own custom generator.

Note: Instead of training on images, you’ll now train on the combination of an image plus a bounding box annotation. For images that have more than one annotation, it’s therefore possible that the same image appears multiple times in the same batch, although each time with a different bounding box.

The code for this generator is again in helpers.py. First, let’s see the generator in action and then we’ll describe how it works:

from helpers import BoundingBoxGenerator

batch_size = 32
train_generator = BoundingBoxGenerator(
  train_annotations,
  train_dir,
  image_height,
  image_width,
  batch_size,
  shuffle=True)

This imports the BoundingBoxGenerator class from the helpers module and creates a new instance. You have to give it the following information:

• The DataFrame that contains the annotations, train_annotations.
• The folder that contains the images for this DataFrame, in this case, snacks/train.
• The image size that the neural network will expect.
• A batch size, i.e., how many training examples the generator should combine into a mini-batch. Here, you’re using a batch size of 32 images.
• Whether you want to randomly shuffle the examples or not. For training, this should be True, for validation and testing this is usually False.

Now, run the following cell to grab a batch of training data:

train_iter = iter(train_generator)
X, (y_class, y_bbox) = next(train_iter)

The iter() function turns train_generator into a so-called iterator object, and next() asks this iterator to return its next element. The generator, in other words, is simply a collection of training examples that you can iterate over. Keras does exactly the same thing in its training loop: it calls next() over and over until it has seen all 7,040 rows from the dataframe.

The NumPy array X now contains thirty-two training images (because the batch size is 32), while y_class and y_bbox will contain the class labels and ground-truth bounding boxes for these images. You can verify this by printing the shape of these arrays:

X.shape

This prints (32, 224, 224, 3) because it contains thirty-two 224×224 color images. The shape of y_class is (32,) because it has thirty-two class labels. And the shape of y_bbox is (32, 4) because it has thirty-two bounding boxes — one per image — and each box is made up of four coordinates.

If you print y_bbox it will look like this:

array([[ 0.348343,  0.74359 ,  0.55838 ,  0.936911],
       [ 0.102564,  0.746717,  0.062909,  0.93219 ],
       [ 0.      ,  1.      ,  0.135843,  0.98036 ],
       [ 0.448405,  0.978111,  0.288574,  0.880734],
       ...

The numbers you’ll see will be different because the generator randomly shuffles the examples. y_class will be something like this:

array([ 9, 16, 12,  7,  8, 18, 10,  1, 14,  2,  7, 17, ...])

These are the indices of the classes that belong to the bounding boxes. To turn this back into text labels, you can do the following:

from helpers import labels
list(map(lambda x: labels[x], y_class))

The labels variable contains the class names corresponding to these indices and is defined in helpers.py. Using the map() function, which works the same way as Swift’s map, you can convert from y_class’s numeric indices back to text labels. The helpers module also has a label2index dictionary that does the mapping the other way around, from text labels to numeric class indices.

Now, have a look at how exactly this generator works. Open helpers.py to view the complete code, but here are the highlights. BoundingBoxGenerator is a subclass of the Keras Sequence object that overrides a couple of methods:

class BoundingBoxGenerator(keras.utils.Sequence):
    def __len__(self):
        return len(self.df) // self.batch_size

    def __getitem__(self, index):
        # ... code ommitted ...
        return X, [y_class, y_bbox]

    def on_epoch_end(self):
        self.rows = np.arange(len(self.df))
        if self.shuffle:
            np.random.shuffle(self.rows)

The __len()__ method determines how many batches this generator can produce: the number of rows in the dataframe, len(self.df), divided by the size of the batch. The // operator in Python means integer division.

When you write len(train_generator), Python automatically invokes this __len()__ method. It should output 220. The generator produces exactly 220 batches because 7,040 rows / 32 rows per batch = 220 batches. Usually, the size of the training set doesn’t divide so neatly by batch_size, in which case the last, incomplete batch is ignored or is padded with zeros to make it a full batch. (We ignore it.)

The on_epoch_end() method is called by Keras after it completes an epoch of training, i.e., after the generator has run out of batches. Here, on_epoch_end() creates an instance variable self.rows that contains the indices of the rows in the DataFrame. Normally self.rows is [0, 1, 2, ..., len-1] but if shuffle is true, the indices in self.rows get randomly reordered. BoundingBoxGenerator’s constructor, called __init__ in Python, also calls on_epoch_end() to make sure the rows are properly shuffled before the first epoch starts.

The meat of the work happens in __getitem__(). This method is called when you do next() or when you write train_generator[some_index]. This is where the batch gets put together. __getitem__() does the following:

1. Create new NumPy arrays to hold the images X, and the targets y_class and y_bbox for one batch. These arrays are initially empty.

2. Get the indices of the rows to include in this batch. It looks these up in self.rows.

3. For every row index, grab the corresponding row from the DataFrame. Load the image, preprocess it using the standard MobileNet normalization function, and put it into X. Also get the class name, use the label2index dictionary to convert it to a number and put it into y_class. Finally, get the bounding box coordinates and put them into y_bbox.

4. Return X, as well as y_class and y_bbox, to the caller.

Note: __getitem__() returns a tuple of two elements: The first one holds the array with the training images X, the second element holds the targets. But you have two different targets here, one for the classes and one for the bounding boxes. This is why earlier you wrote X, (y_class, y_bbox) = next(train_iter), to unpack this second tuple element into separate y_class and y_bbox variables.

To test that the generator works OK, plot the images and bounding boxes that it returns:

def plot_image_from_batch(X, y_class, y_bbox, img_idx):
    class_name = labels[y_class[img_idx]]
    bbox = y_bbox[img_idx]
    plot_image(X[img_idx], [[*bbox, class_name]])

plot_image_from_batch(X, y_class, y_bbox, 0)

This uses the plot_image() function again but this time the image and bounding box comes from the batch. You need to supply the index of the image in the batch (0 to 31).

You may get a warning message now, “Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).” This is matplotlib telling you that it has trouble interpreting the image data from X. That’s because the pixel values are no longer between 0 and 255 but between -1 and +1 due to the normalization performed by the generator. Matplotlib will still display the images but they are a bit darker than usual.

To grab a new batch of images, simply repeat this statement:

X, (y_class, y_bbox) = next(train_iter)

You can also do the same on validation and test iterators. The only difference is that they don’t shuffle their images, so they’ll always appear in the same order.

There you have it: a dataset with bounding box annotations that’s ready for training. All you need to do now is create a suitable model for it.

Note: In the last chapter, you saw that data augmentation was a neat trick to increase the number of available training examples. The generator is the ideal place to do this sort of thing. To keep the code simple, BoundingBoxGenerator is currently not doing any data augmentation. If you’re up for a challenge, try adding data augmentation code to the generator — but don’t forget that the bounding boxes should be transformed too along with the images!

A simple localization model

You’re now going to extend the existing MobileNet snacks classifier so that it has the ability to predict a bounding box as well as a class label.

To create the classifier, you took the MobileNet feature extractor and added a logistic regression on top, made up of a Dense layer, a softmax activation, as well as a Dropout layer for regularization. Guess what: There’s no reason why you can’t add another bunch of layers that branch off of the feature extractor. These new layers will now predict the bounding box coordinates:

This new model has two outputs: one for the classification results, and one for the bounding box predictions. Both sets of layers are built on the same features from the MobileNet feature extractor, but because you train them on different targets they learn to predict different things. The classification portion of the model is still the same as before and outputs a probability distribution over the 20 possible classes.

The bounding box predictor outputs four real-valued numbers: x_min, x_max, y_min and y_max. If you train the model well, these four numbers will form the corners of a proper bounding box that encloses the object in the image.

Note: Neural networks can have as many outputs as you like, one for every task that you want the model to perform. Best of all, you can train the model to learn all of these tasks at the same time. Models can even have multiple inputs. For example, a second input could be a table with extra information about the image such as its EXIF data, which contains the time of day the image was taken, where it was taken, and other metadata. The only requirement is that you are able to turn this input data into numbers somehow, for example by one-hot encoding it.

To save some training time, you’ll start with the classifier model from the last chapter. After all, this has already learned how to classify snacks and so it already contains a lot of knowledge about the problem domain. What you’re going to do in this section is to add some additional knowledge about bounding boxes to the model.

To load the best model from last time, do the following:

import keras
from keras.models import Sequential
from keras.layers import *
from keras.models import Model, load_model
from keras import optimizers, callbacks
import keras.backend as K

checkpoint = "checkpoints/multisnacks-0.7162-0.8419.hdf5"
classifier_model = load_model(checkpoint)

This simply grabs the best checkpoint and loads it back in. You can find this checkpoint in this chapter’s starter folder.

Tip: Call classifier_model.summary() to check that the model was loaded correctly.

To add the bounding box predictor layers on top of this checkpoint requires a bit of trickery, because you want to keep most of the existing model but also add a new output. It’s easiest to build a new model but reuse some of the layers. Since this new model will involve a branching structure, you can’t use the Sequential model API anymore but you have to use the Keras functional API as you saw in last chapter’s SqueezeNet section.

The code is as follows. First, you reconstruct the classifier model from last time:

num_classes = 20

# The MobileNet feature extractor is the first "layer".
base_model = classifier_model.layers[0]

# Add a global average pooling layer after MobileNet.
pool = GlobalAveragePooling2D()(base_model.outputs[0])

# Reconstruct the classifier layers.
clf = Dropout(0.7)(pool)
clf = Dense(num_classes, kernel_regularizer=regularizers.l2(0.01),
            name="dense_class")(clf)
clf = Activation("softmax", name="class_prediction")(clf)

A quick reminder of how the functional API works: You create a layer object, such as GlobalAveragePooling2D(), and then call this layer object on a tensor, such as base_model.outputs[0], which is the output from the MobileNet feature extractor. This, in turn, gives a new tensor, pool. Then, you create a new layer, Dropout(0.7), apply this to the pool tensor to get the next tensor, and so on. After you run this code, clf is now the tensor that refers to the model’s classification output.

Here is the new bit for the bounding box predictor:

bbox = Conv2D(512, 3, padding="same")(base_model.outputs[0])
bbox = BatchNormalization()(bbox)
bbox = Activation("relu")(bbox)
bbox = GlobalAveragePooling2D()(bbox)
bbox = Dense(4, name="bbox_prediction")(bbox)

This adds a new Conv2D layer that also works directly on the output of the MobileNet feature extractor, given by the tensor base_model.outputs[0]. As is common, the convolution layer is followed by batch normalization and a ReLU. After this comes a GlobalAveragePooling2D layer and the final Dense layer that has four outputs for the bounding box coordinates. bbox is now the tensor for the model’s bounding box output.

Note that the Dense layer for the bounding box prediction does not have an activation function, also sometimes called a linear activation. That means this part of the model performs linear regression, the kind of machine learning that predicts real numbers. Applying a softmax activation here wouldn’t make sense because you’re not trying to predict a probability distribution — you definitely want four independent numbers.

Note: Because the four predicted numbers for the bounding box ought to be normalized coordinates between 0 and 1, in theory it’s possible to apply a sigmoid activation to this Dense layer. The sigmoid function always returns 0, 1, or a value in between. Applying a sigmoid function is a common mathematical trick to restrict numbers to the range [0, 1]. However, the author found that using a linear activation — i.e., having no activation function — worked better.

Finally, you combine everything into a new Model object:

model = Model(inputs=base_model.inputs, outputs=[clf, bbox])

Don’t forget to set the layers of the MobileNet base model to non-trainable, unless you’re interested in fine-tuning the entire model:

for layer in base_model.layers:
    layer.trainable = False

You could also set the classifier layers to be non-trainable since they’ve already been trained before, but it’s probably a good idea to keep training them. The class is now taken from the object in the bounding box, which is not necessarily 100% the same as the class of the entire image.

The model.summary() shows the extra layers, but it can be tricky to understand how they’re connected. To get a good idea of the branching structure, it’s useful to make a plot:

from keras.utils import plot_model
plot_model(model, to_file="bbox_model.png")

The bottom part of this file looks like this:

The conv_pw_13 layers at the top are part of MobileNet. On the right, it shows the classifier branch, and on the left the new bounding box prediction branch. Note that the bounding box branch is slightly larger: it has an extra convolution layer between the MobileNet output and the global average pooling layer.

Note: You may be wondering exactly why you’ve added another Conv2D layer, here. Why not do the same as in the classifier branch and just have a Dense layer that immediately follows the global pooling? Good question. The answer is that the author tried both and adding the convolution layer gave much better results. This is probably because this extra layer helps to convert from image-level features to features that are more useful for predicting bounding boxes. The downside is that having this extra Conv2D layer adds over 4 million additional parameters to the model. Yikes. In the next chapter, you’ll look at a more refined approach to building bounding box predictors that uses way fewer parameters.

Now, at this point, there is an important step you shouldn’t overlook. Because you reconstructed the model’s classification layers, the weights for these layers are still initialized with random numbers. If you’d use this model to make a classification, it would predict a random class. So before you continue, first put the weights back:

layer_dict = {layer.name:i for i, layer in enumerate(model.layers)}

# Get the weights from the checkpoint model.
weights, biases = classifier_model.layers[-2].get_weights()

# Put them into the new model.
model.layers[layer_dict["dense_class"]].set_weights([weights,
                                                     biases])

The layer_dict lets you look up layers in the Keras model by name. That’s why you gave the new layers names when you created them. "dense_class" is the name of the Dense layer in the classification branch. With get_weights() you can grab a layer’s weights, and biases if it has them; with set_weights(), you can change the weights on a layer.

Note: In the original classifier model you didn’t give the layers names. In that case, Keras will automatically choose names and you can’t really depend on them having a certain name. That’s why to load the weights, you use layers[-2]. In Python notation, a negative index means that you’re indexing the array from the back, so layers[-1] would be the last layer, which is the softmax activation, making layers[-2] the classification layer. Using indices is fine but giving the layers clear names is better.

The new loss function

With the definition of the model complete, you now can compile it:

model.compile(loss=["sparse_categorical_crossentropy", "mse"],
              loss_weights=[1.0, 10.0],
              optimizer=optimizers.Adam(lr=1e-3),
              metrics={ "class_prediction": "accuracy" })

There are a few new things going on, here. Previously, you specified a single loss, categorical_crossentropy. Here, you have specified not one but two loss functions: sparse_categorical_crossentropy and mse. The model has two outputs and each predicts a different thing, so you want to use a different loss function for each output.

The cross-entropy loss function is great for classification tasks, but it’s not suitable for the bounding box prediction.

Note: The sparse categorical cross-entropy you’re using here, does the same thing as the regular one you’ve used in the previous chapters. It compares the predicted probability distribution with the true class label. The difference is one of convenience. Recall that the BoundingBoxGenerator returns the target y_class as a list of class indices. In Chapter 6, “Taking Control of Training with Keras,” you saw that such targets need to be one-hot encoded, so y_class really ought to be a tensor of size (batch_size, 20) with the classes as one-hot encoded vectors. But Keras is clever: if you use the sparse_categorical_crossentropy loss function instead of the regular categorical_crossentropy, it will one-hot encode the class labels on-the-fly, saving you the effort of doing it yourself.

The loss function for the bounding box predictions is "mse" or mean squared error. This is a typical loss function for regression tasks, i.e., when the output of the model consists of real-valued numbers, such as bounding box coordinates. The math for this loss function looks like this:

mse_loss = sum( (truth - prediction)**2 ) / (4*batch_size)

Let’s unpack this:

• First, it finds the difference between the ground-truth value and the prediction by subtracting the two numbers: truth - prediction. This is the error in mean squared error.

• Then it takes the square, which in Python is done with **2, so that this difference will always be a positive number. This also makes larger errors count more since the square of a large number is much bigger than the square of a small number. This is a common mathematical trick that you see all the time in machine learning. So now you have the squared error.

• Finally, it sums up all these squared differences and divides by how many there are. The loss is computed over a batch at a time, and there are four predicted numbers for each bounding box. In other words, it takes the average — the mean — of the squared errors for all the predictions in the batch. Put it all together and you get the mean squared error.

You don’t need to remember this math; just realize that it’s a really simple formula and that "mse" is the loss function to use when dealing with predictions that are just numbers, as opposed to probability distributions.

model.compile() now also has a loss_weights argument. Because there are two outputs, the loss computed during training looks like this:

loss = crossentropy_loss + mse_loss + L2_penalties

But not all of these loss terms will have the same scale, so some will count more than others in the final sum. Or perhaps you decide that some of them should count more than others. That’s why each of these terms is weighted. The choices we’ve made with loss_weights=[1.0, 10.0] result in a final loss function that looks like this:

loss = 1.0*crossentropy_loss + 10.0*mse_loss + 0.01*L2_penalties

Because this model has already been trained on the classification task but hasn’t learned anything about the bounding box prediction task yet, we’ve decided that the MSE loss for the bounding boxes should count more heavily. That’s why it has a weight of 10.0 versus a weight of 1.0 for the cross-entropy loss. This will encourage the model to pay more attention to errors from the bounding box output.

Note: Recall the that L2 penalties are extra terms that are added to the loss for regularization purposes. The 0.01 weight for the L2 penalties comes from kernel_regularizer=regularizers.l2(0.01) in the definition of the model.

Sanity checks

At this point, it’s a good idea to see what happens when you load an image and make a prediction. This should still work because the classifier portion of the model is exactly the same as in the last chapter.

from keras.applications.mobilenet import preprocess_input
from keras.preprocessing import image

img = image.load_img(train_dir + "/salad/2ad03070c5900aac.jpg",
                     target_size=(image_width, image_height))

Now, normalize the image and let the model loose on it:

x = image.img_to_array(img)
x = np.expand_dims(x, axis=0)
x = preprocess_input(x)

preds = model.predict(x)

The preds variable is a list containing two NumPy arrays: The first array, preds[0], is the 20-element probability distribution from the classifier output. The second array, preds[1], has the four numbers for the bounding box.

Right now, the bounding box prediction is completely bogus because those layers haven’t been trained yet, but the classification result should be reasonable. If not, something is wrong with the model. An easy way to check is to plot the predicted probabilities as a bar chart:

plt.figure(figsize=(10, 5))
plt.bar(range(num_classes), preds[0].squeeze())
plt.xticks(range(num_classes), labels, rotation=90, fontsize=20)
plt.show()

Which indeed shows this is an image of a salad:

In fact, if you do classifier_model.predict(x), which uses the last chapter’s model without the bounding box layers added, then you should get the exact same probability distribution. (Try it!)

Of course, you can also use the generator to make predictions:

preds = model.predict_generator(train_generator)

This will create predictions for all the rows in the train_annotations dataframe, an array of size (7040, 20) for the classification output, and an array of size (7040, 4) for the bounding box output. But as you’ve seen, the bounding box predictions don’t make much sense yet… at least until you train the model.

Train it!

Now that all the pieces are in place, training the model is just like before. This model is again trained best on a machine with a fast GPU. (If you have a slow computer, it’s not really worth training this model yourself.)

First, create a generator for the validation set, with shuffle set to False:

val_generator = BoundingBoxGenerator(val_annotations, val_dir,
                                     image_height, image_width,
                                     batch_size, shuffle=False)

Some of the helper code now lives in helpers.py, so import those functions:

from helpers import combine_histories, plot_loss, plot_bbox_loss
histories = []

And then train for a number of epochs:

histories.append(model.fit_generator(train_generator,
                           steps_per_epoch=len(train_generator),
                           epochs=5,
                           validation_data=val_generator,
                           validation_steps=len(val_generator),
                           workers=8))

Because there is more going on in the model, Keras also prints out more information during training:

Epoch 1/5
220/220 [==============================] - 14s 64ms/step - loss: 1.8093 - class_prediction_loss: 0.4749 - bbox_prediction_loss: 0.1187 - class_prediction_acc: 0.8709 - val_loss: 1.2640 - val_class_prediction_loss: 0.5931 - val_bbox_prediction_loss: 0.0522 - val_class_prediction_acc: 0.8168

There is class_prediction_loss, which has the cross-entropy loss for the classifier output. There is also bbox_prediction_loss with the Mean Squared Error loss for the bounding box prediction. The names of these metrics are taken from the names of the output layers, which is another reason for giving your layers meaningful identifiers.

Notice how the bounding box loss is much smaller than the class loss, 0.1187 versus 0.4749. You can’t really compare these values because they were computed using completely different formulas. It’s only important that they go down over time.

The total loss value is the sum of these two losses, weighed by the loss_weights you supplied to model.compile(), plus the L2 penalty from the classifier’s Dense layer. This overall loss again is just an indication of what the model is doing — the number itself is meaningless.

Keras also prints out a class_prediction_acc metric that measures the accuracy of the classifications over the training set, but there is no such metric for the bounding box predictions. That’s because you told model.compile() that you only wanted metrics={ "class_prediction": "accuracy" }. After all, what would it mean for a bounding box prediction to be “accurate”? We’ll actually come back to this topic soon because there is a useful metric you can use here, but it’s not accuracy.

It looks like the overall loss is going down during these first five epochs, but it’s hard to say whether this is due to either the classification loss or the bounding box loss. So let’s plot only the bounding box loss and see what that does:

history = combine_histories(histories)
plot_bbox_loss(history)

The training loss certainly went down significantly but the validation loss doesn’t look particularly impressive. So is the model actually learning anything useful? It’s hard to say because the loss itself doesn’t tell you much about how well the model works. The only thing you can say for sure is that the model works better when it has a lower loss than when it has a higher loss — not very enlightening.

At least for the classification output, you can compute the accuracy, which is more interpretable than the loss. If the loss goes down by 10%, what does that mean? Who knows… But the accuracy going up by 10% makes a lot of sense.

Fortunately, for the bounding box predictions, there is also a metric that gives us some intuition about the quality of the model: IOU.

IOU

Sorry, this doesn’t mean I owe you any money. The acronym stands for Intersection-over-Union, although some people call it the Jaccard index.

To measure how well the predicted bounding box matches the ground-truth box from the training data, you can compute how much they overlap, or their intersection. But just the overlap is not enough, what also matters is how much they don’t overlap.

The IOU takes the intersection between the two bounding boxes and divides it by their total area, the union, to get a number between 0 and 1. The more similar the two boxes are, the higher the number. A perfect match is 1, while 0 means the boxes don’t overlap at all.

The helpers.py module has a simple function iou() for computing the Intersection-over-Union between two bounding boxes. You use it like this:

from helpers import iou

bbox1 = [0.2, 0.7, 0.3, 0.6, "bbox1"]
bbox2 = [0.4, 0.6, 0.2, 0.5, "bbox2"]
iou(bbox1, bbox2)

This prints 0.235 (rounded off), meaning that these boxes have only about one-fourth in common. You can see this using plot_image:

plot_image(img, [bbox1, bbox2])

That seems about right:

You can use the average IOU over the validation set as a metric of how good the model’s bounding box predictions are. That’s more enlightening than just the loss value.

To use this metric, you need to compile the model, again:

from helpers import iou, MeanIOU, plot_iou

model.compile(loss=["sparse_categorical_crossentropy", "mse"],
              loss_weights=[1.0, 10.0],
              optimizer=optimizers.Adam(lr=1e-3),
              metrics={ "class_prediction": "accuracy",
                        "bbox_prediction": MeanIOU().mean_iou })

The only difference is the addition of the last line. Now Keras computes the mean IOU for the predictions coming from the model’s "bbox_prediction" output. The MeanIOU object is a simple wrapper class that lets Keras and TensorFlow use the iou() function.

If you train the model again, Keras now also prints out the bbox_prediction_mean_iou metric, which gradually increases from 0.25 to about 0.43 for the training set, but only gets up to approximately 0.34 for the validation set.

You can plot how the IOU developed over time using plot_iou(history). Here is the plot for 15 training epochs, where the learning rate was manually decreased by a factor of 10 after every five epochs.

This shows that the Mean IOU definitely improved over time, at least for the training set. After every five epochs, there’s a nice bump when the learning rate was lowered.

The curve for the validation set isn’t as impressive, though (or as smooth). No doubt there’s some overfitting going on here since that one extra Conv2D layer you added has more parameters than the rest of the model put together…

By the way, this plot is slightly misleading. It may seem as if the validation IOU doesn’t really improve very much, but keep in mind that the validation score is measured after each epoch, so at this point, the model had already seen one epoch of training. On the untrained model, the mean validation IOU is actually close to 0. (Hint: you can see this with model.evaluate_generator(val_generator, steps=len(val_generator)) before you start training.)

So how good is this simple localization model? Well, let’s look at some pictures from the test set and see with your own eyes.

Note: You can also use a loss based on the IOU value, known as the DICE loss. Currently, you’re using the MSE loss, which tries to make each individual corner coordinate of the bounding box as close to the ground-truth as possible. But the model doesn’t really know these four numbers are related. With the DICE loss, you optimize the bounding box as a whole, where the goal is to make the box overlap as large as possible.

Trying out the localization model

Just to get a qualitative idea of how well the model works, a picture says more than a thousand loss curves. So, write a function that makes a prediction on an image and plots both the ground-truth bounding box and the predicted one:

def plot_prediction(row, image_dir):
    # Same as before:
    image_path = os.path.join(image_dir, row["folder"],
                              row["image_id"] + ".jpg")
    img = image.load_img(image_path,
                         target_size=(image_width, image_height))

    # Get the ground-truth bounding box:
    bbox_true = [row["x_min"], row["x_max"],
                 row["y_min"], row["y_max"],
                 row["class_name"].upper()]

    # Make the prediction:
    x = image.img_to_array(img)
    x = np.expand_dims(x, axis=0)
    x = preprocess_input(x)
    pred = model.predict(x)
    bbox_pred = [*pred[1][0], labels[np.argmax(pred[0])]]

    # Plot both bounding boxes and print the IOU:
    plot_image(img, [bbox_true, bbox_pred])   
    print("IOU:", iou(bbox_true, bbox_pred))

This is very similar to the plot_image_from_row() function from earlier, but this time it also makes a prediction on the image and plots the predicted bounding box in addition to the ground-truth box. The function also prints the IOU between the two boxes.

To view the results for a random image from the test set, do the following:

row_index = np.random.randint(len(test_annotations))
row = test_annotations.iloc[row_index]
plot_prediction(row, test_dir)

Here’s an example of a pretty good prediction. The ground-truth box’s label is in uppercase, the predicted box in lowercase.

It’s not an exact match, but the model has definitely located where the hot dog is in the image. The IOU between the boxes is 0.67. IOU values over 0.5 are generally considered to be correct matches.

Unfortunately, there are also many images where the model doesn’t do so well:

An IOU of about 0.03, that’s very bad. But can you really blame this on the model? This image has many apples, and you can argue that the model did indeed find (a portion of) an apple, just not the one in the annotation.

This image really isn’t a fair test of our simple localization model, which was only trained to find a single object at a time. In the next chapter, you’ll train a proper object detection model that can handle images like these and will find all the apples.

Another example:

In this image, the bounding boxes do overlap, but less than the IOU of 50% that you’d like to see. Plus, the model actually found a different class. Again, it’s not a completely wrong answer because this image does have a salad in it.

This is a typical result of a model that can only predict a single bounding box when there are multiple objects in the scene. In such situations, the model tends to predict a bounding box that’s in between the two objects. It tries to hedge its bets and predicts an average box that sits somewhere in the middle. Quite clever, actually.

Conclusion: not bad, could be better

The good news is that it was pretty easy to make the classification model perform a second task, predicting the bounding boxes. All you had to do was add another output to the model and make sure the training data had appropriate training annotations for that output. Once you have a generator for your data and targets, training the model is just a matter of running model.fit_generator().

This is a key benefit of deep learning: You can use the same techniques for building neural network-based models for pretty much any problem domain, whether that’s computer vision, language processing, audio recognition, and many others. As long as you have a dataset with training data and target labels, as well as an appropriate loss function, you’re good to go!

Granted, the simple localization model you built here isn’t super. On the validation set, it had an average IOU of a little over 30%. In general, we only consider a bounding box prediction correct when its IOU is over 0.5 or 50%. The model has definitely learned a few things about bounding boxes but it is still more wrong than it is right.

This is partially the fault of the dataset: If you look through the training images, you’ll see that many images have more than one object — sometimes from different classes — but not annotations for all of these objects. Plus, this simple model can only predict a single bounding box at a time, which obviously doesn’t work so well on images with multiple objects. So there’s still room for improvement.

The solution: create a model that can predict more than one bounding box. That’s what the next chapter is all about. Now stuff is getting serious!

Key points

• Object detection models are more powerful than classifiers: They can find many different objects in an image. It’s easy to make a simple localization model that predicts a single bounding box, but more tricky to make a full object detector.

• To train an object detector, you need a dataset that has bounding box annotations. There are various tools that let you create these annotations. You may need to write your own generator to use the annotations in Keras. Data wrangling is a big part of machine learning.

• A model can perform more than one task. To predict a bounding box in addition to classification probabilities, simply add a second output to the model. This output needs to have its own targets in the training data and its own loss function.

• The loss function to use for linear regression tasks, such as predicting bounding boxes, is MSE or Mean Squared Error. An interpretable metric for the accuracy of the bounding box predictions is IOU or Intersection-over-Union. An IOU of 0.5 or greater is considered a good prediction.

• When working with images, make plenty of plots to see if your data is correct. Don’t just look at the loss and other metrics, also look at the actual predictions to check how well the model is doing.