Using Caffe In A Python Flask Application

I'm going to show you how to set up a barebones Flask application that can classify images using Caffe. This requires that you have Caffe installed and a pre-trained model ready to go (we require a .caffemodel file and the corresponding deploy.prototxt; if you're using digits, you can find these files in digits/digits/jobs/<job_id>). We're going to use some Python code from a DIGITS example to get this to work. If you haven't got DIGITS already, use git clone http://github.com/NVIDIA/digits to retrieve it. The files we need are digits/examples/classification/example.py and requirements.txt in the same directory. Copy these files into a new project folder in your home directory. I'm calling my project example. If you're working off my last guide on How to install Caffe and DIGITS on an Amazon EC2 Instance like I am, use the following commands:

$ cp -r ~/digits/examples/classification/ ~/example;
$ cd ~/example;
$ rm README.md; rm use_archive.py;
$ mv example.py classifier.py;

We also need to set up a virtual environment for our project. If you haven't already got virtualenv, use sudo pip install virtualenv to get it. Now, I've had a lot of trouble satisfying the required dependencies but ultimately this works:

$ virtualenv venv;
$ source venv/bin/activate;
$ pip install -r requirements.txt;
$ pip install -r ~/digits/requirements.txt;
$ pip install scikit-image;

Setting Up The Model

Having not trained any models with DIGITS yet, I'll be using the Caffe model generated by training LeNet on the MNIST dataset with these example scripts:

$ ~/caffe/data/mnist/get_mnist.sh;
$ ~/caffe/examples/mnist/create_mnist.sh;
$ ~/caffe/examples/mnist/train_lenet.sh;

The files are located in caffe/examples/mnist, copy them into our project directory with these commands:

$ cp ~/caffe/examples/mnsit/lenet_iter_10000.caffemodel ~/example/lenet.caffemodel;
$ cp ~/caffe/examples/mnsit/lenet.prototxt ~/example/deploy.prototxt;

If you are using a model trained with digits, similarly copy them from digits/digits/jobs/<job_id>. You may also want to copy mean.binaryproto and labels.txt located in the job folder that corresponds to the dataset you trained your model with. We need to slightly modify the classify function in classifier.py so that it returns our classifications instead of just printing them. Add return classifications to the end of the classify function.

Writing The Application

Now we're going to build the actual web application. I've included a copy of the completed project here. The requirements.txt included there is complete, also. You need to create the following files/folders:

example
|- app.py
|- forms.py
|- templates/
|     |- home.html
|     |- show.html
|- uploads/

Here is the code for them:

# app.py

import os, classifier, datetime
from flask import Flask, render_template, request
from forms import ImageForm
from PIL import Image

BASE_DIR = os.path.abspath(os.path.dirname(__file__))
CAFFE_MODEL = BASE_DIR + "/lenet.caffemodel"
DEPLOY_FILE = BASE_DIR + "/deploy.prototxt"
#MEAN = BASE_DIR + "/mean.binaryproto"
#LABELS_FILE = BASE_DIR + "/labels.txt"
MEAN_FILE = None
LABELS_FILE = None
UPLOAD_FOLDER = BASE_DIR + "/uploads/"

# Edit this to fit the specifications of your model.
def pre_process(filepath) :
  size=(28, 28)
  im = Image.open(filepath)
  im = im.convert('L') # Grayscale
  return im.resize(size)

app = Flask(__name__)
app.debug = True
app.secret_key = 'A0Zr98j/3yX R~XHH!jmN]LWX/,?RT'

@app.route('/', methods=['GET', 'POST'])
def home():
  form = ImageForm()
  if request.method == 'POST':
    image_file = form.image.data
    extension = os.path.splitext(image_file.filename)[1]
    filepath = os.path.join(UPLOAD_FOLDER, \
      datetime.datetime.utcnow().strftime('%Y%m%d%H%M%S%f')) + extension
    image_file.save(filepath)
    pre_process(filepath).save(filepath)

    image_files = [filepath]
    classifications = classifier.classify(
      caffemodel=CAFFE_MODEL,
      deploy_file=DEPLOY_FILE,
      image_files=image_files,
      labels_file=LABELS_FILE,
      mean_file=MEAN_FILE,
      use_gpu=True
    )

    return render_template('show.html', classifications=classifications)
  else:
    return render_template('home.html')

if __name__== "__main__":
        app.run(host="0.0.0.0")

# forms.py

from flask_wtf import Form
from flask_wtf.file import FileField, FileRequired

class ImageForm(Form):
  image = FileField('image',
    validators=[
      FileRequired(message="Please include 'image' field.")
    ])
home.html

<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <title>Classify Image</title>
</head>
<body>
  <form action="/" method="POST" enctype="multipart/form-data">
        <input type="file" name="image" accept="image/*">
        <button type="submit">Submit</button>
  </form>
</body>
</html>
show.html

<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <title>Classify Image</title>
</head>
<body>
  {{classifications}}
</body>
</html>

With that completed, you should be able to run your server with python app.py and access it at http://localhost:5000 or http://<ec2-public-ip>:5000. Ensure the security group that your EC2 instance is a member of is open to inbound requests on port 5000.