TF_model_zoo_runner.py

import os
import tarfile
from distutils.version import StrictVersion

import numpy as np
import six.moves.urllib as urllib
import tensorflow as tf
from IPython.display import HTML, clear_output, display
from matplotlib import pyplot as plt
from PIL import Image

from object_detection.utils import label_map_util
from object_detection.utils import ops as utils_ops
from object_detection.utils import visualization_utils as vis_util

os.chdir("/content/models/research")


# This is needed since the notebook is stored in the object_detection folder.
# sys.path.append("..")


if StrictVersion(tf.__version__) < StrictVersion("1.12.0"):
    raise ImportError(
        "Please upgrade your TensorFlow installation to v1.12.*.")


class TF_model_zoo_runner:
    def __init__(self):
        self._cwd_ = "/content/models/research"
        os.chdir(self._cwd_)
        # What model to download.
        self.MODEL_NAME = "ssd_mobilenet_v1_coco_2017_11_17"
        self.MODEL_FILE = self.MODEL_NAME + ".tar.gz"
        self.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.
        self.PATH_TO_FROZEN_GRAPH = self.MODEL_NAME + "/frozen_inference_graph.pb"
        self.PATH_TO_CHECK_POINT = self.MODEL_NAME + "/checkpoint"
        self.PATH_TO_SAVED_MODEL = self.MODEL_NAME + "/saved_model/"
        # self.PATH_TO_SAVED_MODEL = self.MODEL_NAME + '/saved_model/saved_model.pb'

        # List of the strings that is used to add correct label for each box.
        self.PATH_TO_LABELS = os.path.join(
            "object_detection/data", "mscoco_label_map.pbtxt"
        )

    def run_prepare(self):
        self.download_graph()
        self.load_graph()
        self.load_category_mapping()
        self.setup_test_params()

    def run_logic(self):
        self.run_prepare()
        self.run_test()

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

    def download_saved_model(self):
        opener = urllib.request.URLopener()
        opener.retrieve(self.DOWNLOAD_BASE + self.MODEL_FILE, self.MODEL_FILE)
        tar_file = tarfile.open(self.MODEL_FILE)
        for file in tar_file.getmembers():
            file_name = os.path.basename(file.name)
            print(file_name)
            if "saved_model" in file_name:
                # print('cwd: ' + os.getcwd())
                tar_file.extract(file, os.getcwd())

    def load_graph(self):
        self.detection_graph = tf.Graph()
        with self.detection_graph.as_default():
            od_graph_def = tf.GraphDef()
            with tf.gfile.GFile(self.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="")

    def load_category_mapping(self):
        self.category_index = label_map_util.create_category_index_from_labelmap(
            self.PATH_TO_LABELS, use_display_name=True
        )

    @staticmethod
    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)
        )

    def setup_test_params(self):
        # For the sake of simplicity we will use only 2 images:
        # image1.jpg
        # image2.jpg
        # If you want to test the code with your images, just add path to the images to the TEST_IMAGE_PATHS.
        self.PATH_TO_TEST_IMAGES_DIR = "object_detection/test_images"
        self.TEST_IMAGE_PATHS = [
            os.path.join(self.PATH_TO_TEST_IMAGES_DIR, "image{}.jpg".format(i))
            for i in range(1, 3)
        ][::-1]

        # Size, in inches, of the output images.
        self.IMAGE_SIZE = (12, 8)

    @staticmethod
    def run_inference_for_single_image(image, graph):
        with graph.as_default():
            with tf.Session() as sess:
                # Get handles to input and output tensors
                ops = tf.get_default_graph().get_operations()
                all_tensor_names = {
                    output.name for op in ops for output in op.outputs}
                tensor_dict = {}
                for key in [
                    "num_detections",
                    "detection_boxes",
                    "detection_scores",
                    "detection_classes",
                    "detection_masks",
                ]:
                    tensor_name = key + ":0"
                    if tensor_name in all_tensor_names:
                        tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(
                            tensor_name
                        )
                if "detection_masks" in tensor_dict:
                    # The following processing is only for single image
                    detection_boxes = tf.squeeze(
                        tensor_dict["detection_boxes"], [0])
                    detection_masks = tf.squeeze(
                        tensor_dict["detection_masks"], [0])
                    # Reframe is required to translate mask from box coordinates to image coordinates and fit the image size.
                    real_num_detection = tf.cast(
                        tensor_dict["num_detections"][0], tf.int32
                    )
                    detection_boxes = tf.slice(
                        detection_boxes, [0, 0], [real_num_detection, -1]
                    )
                    detection_masks = tf.slice(
                        detection_masks, [0, 0, 0], [
                            real_num_detection, -1, -1]
                    )
                    detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(
                        detection_masks, detection_boxes, image.shape[1], image.shape[2]
                    )
                    detection_masks_reframed = tf.cast(
                        tf.greater(detection_masks_reframed, 0.5), tf.uint8
                    )
                    # Follow the convention by adding back the batch dimension
                    tensor_dict["detection_masks"] = tf.expand_dims(
                        detection_masks_reframed, 0
                    )
                image_tensor = tf.get_default_graph().get_tensor_by_name(
                    "image_tensor:0"
                )

                # Run inference
                output_dict = sess.run(tensor_dict, feed_dict={
                                       image_tensor: image})

                # all outputs are float32 numpy arrays, so convert types as appropriate
                output_dict["num_detections"] = int(
                    output_dict["num_detections"][0])
                output_dict["detection_classes"] = output_dict["detection_classes"][
                    0
                ].astype(np.int64)
                output_dict["detection_boxes"] = output_dict["detection_boxes"][0]
                output_dict["detection_scores"] = output_dict["detection_scores"][0]
                if "detection_masks" in output_dict:
                    output_dict["detection_masks"] = output_dict["detection_masks"][0]
        return output_dict

    def run_test(self):
        for image_path in self.TEST_IMAGE_PATHS:
            image = Image.open(image_path)
            # the array based representation of the image will be used later in order to prepare the
            # result image with boxes and labels on it.
            image_np = self.load_image_into_numpy_array(image)
            # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
            image_np_expanded = np.expand_dims(image_np, axis=0)
            # Actual detection.
            output_dict = self.run_inference_for_single_image(
                image_np_expanded, self.detection_graph
            )
            # Visualization of the results of a detection.
            vis_util.visualize_boxes_and_labels_on_image_array(
                image_np,
                output_dict["detection_boxes"],
                output_dict["detection_classes"],
                output_dict["detection_scores"],
                self.category_index,
                instance_masks=output_dict.get("detection_masks"),
                use_normalized_coordinates=True,
                line_thickness=8,
            )
            plt.figure(figsize=self.IMAGE_SIZE)
            plt.imshow(image_np)

    def check_graph(self):
        # tensorf_names = [n.name for n in self.detection_graph.as_graph_def().node]
        # Visualizing the network graph. Be sure expand the "mixed" nodes to see their
        # internal structure. We are going to visualize "Conv2D" nodes.
        graph_def = self.detection_graph.as_graph_def()
        tmp_def = rename_nodes(graph_def, lambda s: "/".join(s.split("_", 1)))
        show_graph(tmp_def)

    def set_model(self, name):
        self.MODEL_NAME = name
        self.MODEL_FILE = self.MODEL_NAME + ".tar.gz"
        self.PATH_TO_FROZEN_GRAPH = self.MODEL_NAME + "/frozen_inference_graph.pb"
        self.PATH_TO_CHECK_POINT = self.MODEL_NAME + "/checkpoint"
        self.PATH_TO_SAVED_MODEL = self.MODEL_NAME + "/saved_model/"

    def load_model(self):
        self.download_saved_model()

        loaded = tf.saved_model.load(self.PATH_TO_SAVED_MODEL)



def strip_consts(graph_def, max_const_size=32):
    """Strip large constant values from graph_def."""
    strip_def = tf.GraphDef()
    for n0 in graph_def.node:
        n = strip_def.node.add()
        n.MergeFrom(n0)
        if n.op == "Const":
            tensor = n.attr["value"].tensor
            size = len(tensor.tensor_content)
            if size > max_const_size:
                tensor.tensor_content = tf.compat.as_bytes(
                    "<stripped %d bytes>" % size)
    return strip_def


def rename_nodes(graph_def, rename_func):
    res_def = tf.GraphDef()
    for n0 in graph_def.node:
        n = res_def.node.add()
        n.MergeFrom(n0)
        n.name = rename_func(n.name)
        for i, s in enumerate(n.input):
            n.input[i] = rename_func(
                s) if s[0] != "^" else "^" + rename_func(s[1:])
    return res_def


def show_graph(graph_def, max_const_size=32):
    """Visualize TensorFlow graph."""
    if hasattr(graph_def, "as_graph_def"):
        graph_def = graph_def.as_graph_def()
    strip_def = strip_consts(graph_def, max_const_size=max_const_size)
    code = """
        <script>
          function load() {{
            document.getElementById("{id}").pbtxt = {data};
          }}
        </script>
        <link rel="import" href="https://tensorboard.appspot.com/tf-graph-basic.build.html" onload=load()>
        <div style="height:600px">
          <tf-graph-basic id="{id}"></tf-graph-basic>
        </div>
    """.format(
        data=repr(str(strip_def)), id="graph" + str(np.random.rand())
    )

    iframe = """
        <iframe seamless style="width:800px;height:620px;border:0" srcdoc="{}"></iframe>
    """.format(
        code.replace('"', "&quot;")
    )
    with open("graph_detail.htm", "w") as f:
        f.write(iframe)

    display(HTML(iframe))