helpers.py

"""
 Code has been taken and modified from Object Detection Demo jupyter notebook present
 in object_detection folder and originally in the following url:
 https://github.com/tensorflow/models/blob/master/research/object_detection/object_detection_tutorial.ipynb

"""

import six.moves.urllib as urllib
import tarfile
import os

import tensorflow as tf
import numpy as np

# What model to download.
MODEL_NAME = 'ssd_mobilenet_v1_coco_2018_01_28' # UPDATED VERSION TO 2018
MODEL_FILE = MODEL_NAME + '.tar.gz'
DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/'

# Path to frozen detection graph. This is the actual model that is used for the object detection.
PATH_TO_FROZEN_GRAPH = MODEL_NAME + '/frozen_inference_graph.pb'

def download_model():
  opener = urllib.request.URLopener()
  opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE)
  tar_file = tarfile.open(MODEL_FILE)
  for file in tar_file.getmembers():
    file_name = os.path.basename(file.name)
    if 'frozen_inference_graph.pb' in file_name:
      tar_file.extract(file, os.getcwd())

# Perhaps I shouldn't load model into memory for each call?
def load_model_into_memory():
  detection_graph = tf.Graph()
  with detection_graph.as_default():
    od_graph_def = tf.GraphDef()
    with tf.gfile.GFile(PATH_TO_FROZEN_GRAPH, 'rb') as fid:
      serialized_graph = fid.read()
      od_graph_def.ParseFromString(serialized_graph)
      tf.import_graph_def(od_graph_def, name='')
  return detection_graph

def load_image_into_numpy_array(image):
  (im_width, im_height) = image.size
  return np.array(image.getdata()).reshape(
      (im_height, im_width, 3)).astype(np.uint8)