Build an Image Classification App
A guide to building an AI model using TensorFlow, Python, matplotlib, and packaging it for the web using Ruby on Rails
I was recently tasked with building a web API for fruit classification. I realized that such a product would need to have the ability to detect general objects to weed out images of non-fruits. So I decided that I would start the project by getting my hands dirty by exposing a TensorFlow model to the web.
In this article, I will cover the functionality and architecture of this app from a practical, code-first approach. At a high level, the project will be split between its AI and web layers, and I will give you a guided tour of how I glued them together.
In case you want to jump straight to the code, check out the GitHub repository. An installation guide is included in the repo so you can take the app for a spin locally (tested on Mac and Ubuntu). You can also check out the live demo here.
So without further ado, here is a general object classification model I hacked together with TensorFlow, Python, and Ruby on Rails.
- Users should be able to upload an image and see a preview.
- Users should be associated with their uploaded image (but without login/authentication) so we can present the correct image back to them once it is analyzed.
- Image analysis should present the user with labeled bounding boxes drawn on the image they uploaded.
- Ensure that the resolution of the image output is fixed regardless of the original image size (scale all images to 800 x 800).
To associate users with their uploads, we will use Ahoy, which will generate a unique token for each visitor. This token will be used to reference the uploaded image and its analyzed output in the file system.
Once the file is uploaded, saved, and scaled to 800 x 800, we will use some basic Bash shell calls from Ruby to pass the image to the Python script for processing. Once it's processed, we will grab it from the output directory (where it's placed by the Python script after analysis) using Ruby and present it to the user through the Rails stack.
Now that we have a gameplan in place, let’s get to work!
The Rails Side
The Rails application in my repository is stored under the
object-detectr/potashdirectory. Set up Active Storage in Rails as per this tutorial, and set up the following entity relationships:
rails g model Image analyzing:boolean analysis_completed:boolean assertion:stringrails g model ImageAnalysis image:references
After generating the models, install Ahoy. Once the tables are set up, your database schema should look like this:
app/models/image.rb and drop in the following:
This will tell Active Storage that this model has an attachable file named
captured_image and it has one
image_analysisassociation. We will use the
captured_image attribute to store the image given by the user and the
image_analysis association to store the analyzed image which, at the end, will be presented to the user.
app/models/image_analysis.rb looks identical to
image.rb. (This is a code-smell for DRY, and you can unify this into one model easily..)
With this in place, let’s generate a controller that will render a view and handle the file upload.
rails g controller Images new create
For convenience, we will change the route file to present the user with the page for image upload when they visit the root path.
Drop in the plumbing that handles the image upload, scaling, and Python script invocation. While you’re at it, feel free to call the code police because I just broke the single-responsibility principle to smithereens. I put it all in one file in anticipation of making explaining what’s going on a little easier — guilty as charged.
So let’s break it down. The
new action renders the view that allows the user to attach a file for analysis (more on this a little later) and the
before_action Hook ensures that the user must attach a file before moving forward.
The action happens in the
create method, where we grab the unique token for the current user session from
Ahoy with the call to
current_visit.visitor_token. Using that token as an identifier, we create the image, scale it, and grab the path to the scaled image.
Once the paths to the scaled image and the Python script are calculated, we make system calls (lines 15–16) to copy the scaled image to the directory that expects the input (on the Python side) and run the Python script that will call the TensorFlow model.
After the Python script is run, we compute the expected file name of the output and the path to the directory where it will be placed. If the file has been created without error (line 20), we create an
ImageAnalysis instance and attach the analyzed image to it.
Let’s take a look at the views rendered to the user. The
new action will present the user with a form and a submit button:
onChange listener is added to the form, which fires the function
readURL that sets the source attribute of the
img tag to the selected image. With this code in place, we are able to render a preview as soon as the user attaches an image to the form. Let’s take a look at the view that renders the output once they click Analyze!:
We check if the processed image exists (line 8) and if it does, we render it; otherwise, we present an error, asking the user to give us another chance.
All right, that’s all for Rails/Ruby. Let’s move on to the Python side!
The Python Side
The TensorFlow models live under the
object-detectr/models/research directory of the project. These models were cloned from the official TensorFlow repository.
After making sure that your Python version is 3.7.9, run the setup script and build the project:
python setup.py build
python setup.py installpython -m pip install TensorFlow==1.15 lxml pillow matplotlib jupyter contextlib2 cython tf_slim
After building the project and installing dependencies, jump to
I grabbed the original version of this script from this repository and adapted it to my use case. If you want an in-depth view of the script, check out this excellent video by the author of the repository mentioned above.
My changes to the script are as follows:
- Allow for a string to be passed as a command-line argument, and use that to point to the input file name when constructing the
TEST_IMAGE_PATHSarray (line 102).
- Modify the output (line 147) to prefix the generated image name with the string that was passed in from the command line.
These changes allow Rails/Ruby to programmatically place an image in the
models/research/object_detection/test_images/ directory using the unique identifier generated by
Ahoy as the file name. Once that is done, the script can be invoked with
python run.py some_unique_identifier.jpeg. Since
run.py will include the file name that is passed in from before as the prefix to the generated file name, we can expect to have the analyzed file placed in the directory
models/research/object_detection/outputs/with the name
If this file exists at the end of this somewhat unwieldy Ruby-and-Python song and dance, it will be presented to the user.
In case you made it this far, for your efforts, here is an analysis I ran on my housecat, Tiberius:
Thanks for reading and keep on hacking!
- Bengemon825. “Bengemon825/TF_Object_Detection2020.” GitHub, 10 July 2020, github.com/Bengemon825/TF_Object_Detection2020/blob/master/scripts/updated_old_example.py.
- Clark, Mike. The Pragmatic Studio, pragmaticstudio.com/tutorials/using-active-storage-in-rails.
- Tech, Lazy. “How to Install TensorFlow Object Detection in 2020 (Webcam and Images!).” YouTube, YouTube, www.youtube.com/watch?v=usR2LQuxhL4.