Image Classification with TFLearn and Python

In today’s post we are going to walk through how to build a flexible image classifier using TFLearn and Python. For those not familiar with TFLearn, it is a wrapper around the very popular Tensorflow library from Google. Building a image classifier with TFLearn is relatively simple and today we are going to walk through how to build your own image classifier.


We are going to need to import a number of different libraries in order to build our classifier. For users on Windows the easiest way to install the Scipy library is to use the pre-compiled wheel which can be found here. Once you have installed all the required imports, we can start building our ImageClassify class.

Initializing Our Class

When initializing our class, we are going to need to know a few pieces of information. We are going to need a list of class names. These are the names of the different objects that our classifier is going to classify. We also need to pass in an image size, the classifier will automatically resize our images into a square image of the specified size. So, if we pass in a value of 100, our classifier will end up resizing our images to be 100×100 pixels in size.

Generally, the larger the image size the better the classification we will end up with. This is provided that your images are larger than the specified value. It should be warned that using larger images will increase the time taken to train the algorithm. We store this value in self.image_size.

We also pass in default values for our learning rate and test split. The learning rate dictates how quickly the machine learning algorithm will discover new features. As a default value 0.001 tends to work well. The test split, defines what percentage of samples we will use to validate our model against. Again, using around ten percent of samples for your test size works pretty well.

We also create empty lists which will end up holding our image data and their respective labels

Labeling An Image

Our extract label function takes an image and extracts the label from the image. We begin by creating an array of zeroes, with a zero for each class to be trained. We then split the file name of the image. Our extract label functions expects images in the following format “class.number.png”. Using this format allows us to extract the class name directly from the image. We then look up the index of the class label and set that value in our array of zeros to a 1. Before returning the array itself.

Processing Images

Our process image function first calls our label function. We then read the image using skimage’s io.imread function. We then resize this image to the size specified when we initialized the class. We then append the image data and the labels to self.image_data and self.labels respectively.

Processing images is simply involves us using our process image function on every single image we provide to our image classification class.

Building Our Model

Our build model function simply builds us a convolutional net model, using the parameters we defined when initializing our class. Explaining the working’s of the net are probably beyond the scope of this post. But I will just note that creating our model like this allows our classifier to be used with images of any size and datasets with any number of classes. Creating a build model function also makes it easier to load and predict using pre-trained models.

Training Our Model

Our train_model function, takes a model name, epochs and a batch size parameter. The epochs parameter determines the number of times the model will be run over the entirety of the dataset. The batch size determines the number of samples to be run through the model at once. Generally, the more epochs the more accurate the model will be. Though too many epochs may mean that your model over fits the dataset and you end up with rather inaccurate predictions when you use the model to make predictions. If accuracy hits 100% and loss goes to 0, this is a very strong indication that you have over fit.

We first begin by creating X, y variables using the self.image_data and self.labels variables. We then use our self.test_split value to split the dataset up into training and test sets. We then call the build model function. We then call the fit method on model using both the train and test sets for validation purposes.

Once we have finished training the model. We save the model using the passed in model name and set self.model to be equivalent to our newly trained model.

Loading A Pre-Trained Model & Predicting Images

We can define a very simple function to load a model. This will be useful when we need to predict images sometime after we have trained a model. We can load a model by simply passing in the model’s name.

We then need another function to take an image and transform it to something we can use in our prediction function. This is much like our process image function, with the exception that we have no need to label the image.

Our predict image function takes a path to an image file. We call our _image_to_array function this data can then be fed straight into the model. Our model will then output an array of probabilities. We can line these up with the classes which we provided to the Image Classify class. We then pull out the most probable label, before returning this and the list of probabilities.

Example Usage: Training A Model

Example Usage: Making  A Prediction With Already Trained Model

Full Code & Example Dataset

The full code and an example data set can be found on my Github here. The Github also contains another image classification model which makes use of Google’s Googlenet model. This model is very highly accurate but takes a considerable amount of time to train and is likely to need to be run for a greater number of epics.

Leave a Reply

Your email address will not be published. Required fields are marked *