Lola Detector

lola_detector.py

import tensorflow as tf, sys
import simplejson as json
from flask import Flask, jsonify

app = Flask(__name__)


def start_tf_session():

  # Loads label file, strips off carriage return
    label_lines = [line.rstrip() for line 
                       in tf.gfile.GFile("tfmodel/retrained_labels_cats.txt")]
    
    # Unpersists graph from file
    with tf.gfile.FastGFile("tfmodel/retrained_graph_cats.pb", 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        _ = tf.import_graph_def(graph_def, name='')
    
    print('Tensorflow session (for cats) created.')
    return tf.Session()


# Create the "Tensor"
def get_tensor(sess):
    print('softmax_tensor (for cats) got gotten.')
    return sess.graph.get_tensor_by_name('final_result:0')


# This function gets called by the web service. It runs the Tensor model.
def image_labels(image_name):
    
    # Load the image from disk
    print('loading image')
    image_data = tf.gfile.FastGFile(image_name, 'rb').read()
    
    # Use TensorFlow!
    print('running prediction')
    predictions = sess.run(softmax_tensor, \
             {'DecodeJpeg/contents:0': image_data})
    
    
    print('sorting and organizing results')
    # Sort to show labels of first prediction in order of confidence
    top_k = predictions[0].argsort()[-len(predictions[0]):][::-1]
    
    results = dict()
        
    # Save the results from the TensorFlow function as a Dictionary  
    for node_id in top_k:
        human_string = label_lines[node_id]
        score = predictions[0][node_id]
        print('%s (score = %.5f)' % (human_string, score))
        results[human_string] = str(round(score,5))

    return results    

# Start the app!
print('Starting Tensorflow (for cats)')

# Start a TensorFlow session, and save the session as a variable
sess = start_tf_session()

# Get a reference to the softmax tensor with the newly created Session
softmax_tensor = get_tensor(sess)

# Load the labels file from the custom trained model
label_lines = [line.rstrip() for line in tf.gfile.GFile("tfmodel/retrained_labels_cats.txt")]


# Create a web service to analyze the raspberry pi image
@app.route("/lola")
def lola():
    print('Analyzing image for Lola.')
    
    # Get the labels associated with the input image
    labels = image_labels("rpicam.jpg")
    
    # Return a JSON object of labels and probabilities
    print('returning json')
    return jsonify(**labels)

  
# Run Flask web service
if __name__ == "__main__":
    app.run()