tensorflow-hub.py

import numpy as np
import tensorflow as tf
from os import path
from glob import glob
from tensorflow.python.platform import gfile

def load_graph(model_file):
    graph = tf.Graph()

    graph_def = tf.GraphDef()
    graph_def.ParseFromString(open(model_file, 'rb').read())
    
    with graph.as_default():
        tf.import_graph_def(graph_def)
    return graph

def read_tensor_from_image_file(file_name, input_height=299, input_width=299, input_mean=0, input_std=255):
    file_reader = tf.read_file(file_name, "file_reader")
    
    if file_name.endswith(".png"):
        image_reader = tf.image.decode_png(file_reader, channels=3, name="png_reader")
    elif file_name.endswith(".gif"):
        image_reader = tf.squeeze(tf.image.decode_gif(file_reader, name="gif_reader"))
    elif file_name.endswith(".bmp"):
        image_reader = tf.image.decode_bmp(file_reader, name="bmp_reader")
    else:
        image_reader = tf.image.decode_jpeg(file_reader, channels=3, name="jpeg_reader")
    
    float_caster = tf.cast(image_reader, tf.float32)
    dims_expander = tf.expand_dims(float_caster, 0)
    resized = tf.image.resize_bilinear(dims_expander, [input_height, input_width])
    normalized = tf.divide(tf.subtract(resized, [input_mean]), [input_std])
    sess = tf.Session()
    result = sess.run(normalized)

    return result

def load_image_batch(folder_path):
    list_files = glob(path.join(folder_path, '*'))
    n_files = len(list_files)
    
    image_tensor = np.empty(shape=(n_files, 224, 224, 3), dtype=np.float32)
    for i, file in enumerate(list_files):
        print('{} of {}'.format(i+1, n_files), end='\r')
        tensor_4d = read_tensor_from_image_file(file, input_height=224, input_width=224)
        image_tensor[i] = tensor_4d[0]
    print()
    
    return image_tensor

def predict(graph, image):
    input_name = "import/Placeholder"
    output_name = "import/final_result"
    input_operation = graph.get_operation_by_name(input_name)
    output_operation = graph.get_operation_by_name(output_name)
    
    with tf.Session(graph=graph) as sess:
        feed_dict = {input_operation.outputs[0]: image}
        scores = sess.run(output_operation.outputs[0], feed_dict)

    scores = np.squeeze(scores)
    return bool(scores.argmax()), scores

graph = load_graph('PATH/TO/GRAPH_FILE')
image_tensor = load_image_batch('PATH/TO/FOLDER_WITH_IMAGES/')
for img in image_tensor:
    y_pred, score = predict(graph, image=np.expand_dims(img, 0))