Finding Where's Waldo using Mask R-CNN

When I was a kid, I really loved solving Where's Waldo. There were few books (it used to be called Where's Wally) in our school library on which I spent hours finding Waldo. For people who do not know what it is, basically Waldo - a unique character is hidden among hundreds of other characters and you have to find him in all the chaos in the image.

Now that I am too old to be solving it and too busy to spend hours on such things, I decided to build a system that uses deep learning to automatically solve it and spent weeks to build it. 

I started off by treating this like a classification problem with two classes - Waldo and not Waldo, similar to Hot dog - not Hot dog . Once we can get the classification problem successfully solved, we can just apply a classification action mapping (CAM) layer to find Waldo's activations in the image and thus finding Waldo. However I couldn't find enough images of Waldo. I found this repo which has about 20 images. And as there are only 20 Waldo vs thousands of not-Waldo characters, there is very high imbalance in the classes. I still tried though. But the results weren't that great.

When I looked if someone has already worked on it, I found a medium post which used Tensorflow's Faster R-CNN model to do this. But I didn't want to just find bounding boxes, I wanted to actually mask out Waldo in the image. But I got more images of Where's Waldo from it.

Then I came across this paper on Mask R-CNN which sounded promising for this usecase. And it was indeed much better than my earlier approach:

Waldo masked out in the image

Original Image

In this post I would like to share how I was able to get the data, tag it and train a model to be able to solve Where's Waldo. You can checkout my code on github here.

Fork deepwaldo on Github

Mask R-CNN 

The main idea here is to:

1. Take the input image and pass it into a set of convolutional layers that sort of generates a feature map for the given image.
2. Now, you take this feature map and pass it into a region proposal network which generates rectangular regions that say that for the set of final classes, we might have an image in this region. This will have its own classification loss (rpn_class_loss) and bounding box loss (rpn_bbox_loss).
3. Now you take these regions and pass it into a ROI pooling layer using something like non-max-suppression.
4. The regions are then reshaped and passed on to set of convolution layers which predict if there is an object in them or not. This again will have its own classification and bounding box losses. 
5. Now, you have a separate convolution layers which predicts, for every pixel in the bounding box predicted, is it the given class or not. This essentially gives the mask required. Here in addition to the bounding box and classification losses, we also have mask loss.
6. You run all these networks together backpropogating all the losses.

If you want a more clear explanation, checkout the lecture in CS231n. 

Mask R-CNN arch from CS231n (In this case we have a 28x28 mask instead of 14x14)

Data

As I mentioned earlier, I got 20 images from this repo and few more images from the medium post that used Faster R-CNN.  So,  a total of 29 images. I split this into 26 for training and 3 for validation. Then I used the via-via tool (used to tag VGG) to manually draw masks over Waldo in every image.  You can find the images and annotations in my github repo here.

Training

I trained the model for 30 epochs with 100 steps per epoch. The losses on tensorboard:

If you want to train on your own dataset, first set your configurations in the waldo_config.py file or use the default.

This will download the Mask-RCNN model trained on coco dataset to the MODEL_DIR folder and trains a model with the data in the DATA_DIR folder.

python train.py

For prediction, you can do the following which shows a popup with waldo detected in the image.

python predict.py [MODEL PATH] [PATH TO IMAGE]
# for example
python predict.py models/logs/waldo20180612T1628/mask_rcnn_waldo_0030.h5 data/val/5.jpg

In conclusion, the Mask R-CNN algorithm works fairly well to find Waldo for cases where it has already seen similar type of waldo image. Also, it looks like it works much better when the image quality is good and waldo is clearly visible. But I think it is still great since we only had a very tiny training data to train on.