Image Classification with Tensorflow 2.0 | by Rich Folsom | Towards Data Science

In this article, we will address one of the most common AI problems. Here is the scenario: you’re surfing the web at a bar enjoying a beer and some chicken wings when you start to wonder “could I write an image classifier using Tensorflow 2.0 and Transfer Learning?” This article will show an end to end process for accomplishing this. Note, I will not be going into how the model works, that would be a whole other article. This article will show you how to get up and running with Tensorflow 2.0. Once we have that we can go back and play with the params that might make the code more efficient.
The only prerequisites are:
A computer (obviously). I ran mine on a small laptop with no GPU’s.A cell phone with a camera. We’re going to be generating our own images.(Optional) A cloud image storage. I used Amazon photos.Docker, which we will use to host our Tensorflow 2.0 environment. If you don’t have it, installation instructions can be found here
A computer (obviously). I ran mine on a small laptop with no GPU’s.
A cell phone with a camera. We’re going to be generating our own images.
(Optional) A cloud image storage. I used Amazon photos.
Docker, which we will use to host our Tensorflow 2.0 environment. If you don’t have it, installation instructions can be found here
The steps we will follow are:
Install Tensorflow 2.0 Docker image.Acquire a set of images to train/validate/test our model.Organize our images into a directory structure suitable for our model.Download a pre-trained deep learning model.Customize our model for our specific use case (beer or wing).Train our customized model.Visualize model predictions.
Install Tensorflow 2.0 Docker image.
Acquire a set of images to train/validate/test our model.
Organize our images into a directory structure suitable for our model.
Download a pre-trained deep learning model.
Customize our model for our specific use case (beer or wing).
Train our customized model.
Visualize model predictions.
I followed the instructions on this page. Here are the specific commands I used. Note that for this tutorial the image must contain Jupyter.
docker pull tensorflow/tensorflow:2.0.0a0-py3-jupyter
First off, cd to the directory where you will store your source code. From there, we will start our image I created a simple shell script (note that I’m running on Linux, if you’re on a Mac or Windows you probably don’t need the sudo command:
sudo docker run -it -p 8888:8888 -v $PWD:/tf/data -w /tf/data tensorflow/tensorflow:2.0.0a0-py3-jupyter
The most important things to know here are:
-p is a port mapping, our Jupyter notebook runs on port 8888 in Docker so we will map the port 8888 on our machine to match. You should only change this if port 8888 is already in use on your machine.
-v is how to mount a volume in docker. In this case, we’re mounting our current directory ($PWD) as /tp/data. The Jupyter notebook in this Docker image runs in the /tp directory, so when you open jupyter you should see this:
Where the data directory maps to your current directory on your machine.
After you run the above command, you’ll see something like this in your command prompt.
The important line here is the one after “The Jupyter Notebook is running at”. you’ll want to copy that line starting with the “:8888”. Then go to your browser and type http://localhost and paste in the line you copied. In this case:
http://localhost:8888/?token=cd37ab44fab55bce5e44ac6f4bb187a4b34b713c5fbeac9e
At this point, you have Tensorflow 2.0 up and running in a Docker container with access to your local filesystem.
This step is very easy. I used my phone and took about 30 pictures of a beer glass from different angles, and 30 pictures of a chicken wing from different angles. I am a big Amazon Prime user, so I have my phone set up to back up my images to Amazon Photos. You can use whatever cloud environment you prefer (ICloud, Google photos, etc...) Or even email yourself the photos. The point here is to copy the photos to your PC.
The first step is to label our images. There are several ways to do this, but in the end, you’ll want to copy all of your “beer” pictures into a directory called “beer” and your “wing” pictures into a directory called “wings”. From there you are going to want to create a directory structure that looks like this:
Here are the commands I used in Linux:
mkdir trainmkdir testmkdir valmkdir train/beermkdir train/wingsmkdir test/beermkdir test/wingsmkdir val/beermkdir val/wings
At this point, you’ll want to move subsets of your data to the val and test directories. After some google searches, I found this command:
shuf -n 6 -e * | xargs -i mv {} target-directory
I implemented this using the following commands:
cd beershuf -n 6 -e * | xargs -i mv {} ../test/beershuf -n 6 -e * | xargs -i mv {} ../val/beermv * ../train/beercd ..cd wingsshuf -n 6 -e * | xargs -i mv {} ../test/wingsshuf -n 6 -e * | xargs -i mv {} ../val/wingsmv * ../train/wings
This code moves 6 images each to our val and test folders and the rest to our train folders. After these steps, your directory structure should look something like this:
At this point, go back to your Jupyter notebook in your browser and create a new Notebook. First, we need to import the Python libraries we will be using:
import numpy as npimport tensorflow.kerasfrom tensorflow.keras.models import Sequential, Modelfrom tensorflow.keras.layers import  Dropout, Inputfrom tensorflow.keras.layers import Dense, Flattenfrom tensorflow.keras.optimizers import Adamfrom tensorflow.keras.metrics import categorical_crossentropyfrom tensorflow.keras.preprocessing.image import ImageDataGeneratorimport itertoolsimport matplotlib.pyplot as plt%matplotlib inline
Important troubleshooting step: You may get errors with missing libraries. Depending on the version of your Docker image, you may have to run this step:
!pip install --upgrade pip!pip install pillow!pip install scipy!pip install pandas
After you run it, you’ll need to click on the restart kernel button and rerun the import statements.
Now that we have our Python imports complete, we need to generate ImageGenerator objects for each of our image folders. ImageGenerators take an input image and slightly modify it to provide uniformity and shape to train a Neural Network. Note our pictures will be 224x224.
train_path = '/tf/data/beer_wings/train'valid_path = '/tf/data/beer_wings/val'test_path = '/tf/data/beer_wings/test'train_batches = ImageDataGenerator().flow_from_directory(train_path, target_size=(224,224), classes=['beer', 'wings'], batch_size=32)valid_batches = ImageDataGenerator().flow_from_directory(valid_path, target_size=(224,224), classes=['beer', 'wings'], batch_size=32)test_batches = ImageDataGenerator().flow_from_directory(test_path, target_size=(224,224), classes=['beer', 'wings'], batch_size=32)
Here’s a helpful function to see what our ImageGenerator is doing I found this from a very helpful Youtube series:
# plots images with labels within jupyter notebookdef plots(ims, figsize=(24,12), rows=4, interp=False, titles=None):    if type(ims[0]) is np.ndarray:        ims = np.array(ims).astype(np.uint8)        if (ims.shape[-1] != 3):            ims = ims.transpose((0,2,3,1))    f = plt.figure(figsize=figsize)    cols = len(ims)//rows if len(ims) % 2 == 0 else len(ims)//rows + 1    for i in range(len(ims)):        sp = f.add_subplot(rows, cols, i+1)        sp.axis('Off')        if titles is not None:            sp.set_title(titles[i], fontsize=32)        plt.imshow(ims[i], interpolation=None if interp else 'none')imgs, labels = next(train_batches)plots(imgs, titles=labels)
And the output:
Note [0,1] = wings, [1,0] = beer.
Finally, we are ready to download our pretrained model. In this case, we will be using the VGG16 model. Tensorflow 2.0 has numerous models built in. They are defined here.
This is the code to import the pretrained VGG16 model:
vgg16_model = tensorflow.keras.applications.vgg16.VGG16(weights='imagenet', include_top=False, input_tensor=Input(shape=(224,224,3)))
Simple as that! Well, sorta, the important thing is that we set include_top = False, since we’re going to create our own final layers, also note that our Input shape is (224,224,3). The (224,224) matches the ImageGenerators above. The extra 3 is the color channels (red, blue, green).
Now that we have downloaded a pretrained model which can generally predict image classifications, let’s customize it for our needs. Theoretically speaking, the first few layers of models like these simplify parts of the image and identify shapes within them. Those early labels are pretty generic(lines, cirles, squares, etc...), so we don’t want to retrain them. We want to only train the last few layers of our network along with the new layers we add.
First, let’s disable training on all but the last 4 layers of the pretrained model.
for layer in vgg16_model.layers[:-4]:    layer.trainable = False
Now, let’s add our own final layers to the network:
# Create the modelmodel = Sequential() # Add the vgg convolutional base modelmodel.add(vgg16_model) # Add new layersmodel.add(Flatten())model.add(Dense(1024, activation='relu'))model.add(Dropout(0.5))model.add(Dense(2, activation='softmax')) # Show a summary of the model. Check the number of trainable parametersmodel.summary()
There you have it, our own customized model based on VGG16!
Now that we have our model defined, let’s compile it and train it.
model.compile(loss='categorical_crossentropy',              optimizer=tensorflow.keras.optimizers.RMSprop(lr=1e-4),              metrics=['acc'])history = model.fit_generator(      train_batches,      steps_per_epoch=train_batches.samples/train_batches.batch_size ,      epochs=5,      validation_data=valid_batches,      validation_steps=valid_batches.samples/valid_batches.batch_size,      verbose=1)
Now, let’s feed our trained model a set of images that it has never seen. The most important part of the code is these 2 lines:
test_imgs, test_labels = next(test_batches)predictions = model.predict(test_imgs)
The first one generates a new batch of previously unseen images. Let’s see what our model predicted for these images:
import pandas as pddef to_label(value):    if value==0:        return 'beer'    else:        return 'wings'test_imgs, test_labels = next(test_batches)predictions = model.predict(test_imgs)df = pd.DataFrame()df['actual'] = test_labels[:,1]df['predicted'] = np.round(predictions[:,1])df['predicted_labels']=df['predicted'].map(lambda x: to_label(x))plots(test_imgs, titles=df['predicted_labels'])
Obviously, this is a nonsensical example of implementing image classification, but it did provide some valuable information, which can be applied to future projects. Namely, data acquisition, transfer learning and model evaluation. Note that the code can be easily modified to allow for multiple classifications(our example only had 2).