#!/usr/bin/env python
# coding: utf-8

# # Object Detection Demo
# Welcome to the object detection inference walkthrough!  This notebook will walk you step by step through the process of using a pre-trained model to detect objects in an image. Make sure to follow the [installation instructions](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/installation.md) before you start.

# # Imports

# In[ ]:


import numpy as np
import os
import sys
import tensorflow as tf
from tensorflow.keras import backend as K
import pickle

# for download url and extract zip
# import six.moves.urllib as urllib
# import tarfile
# import zipfile

# legacy utils
# from collections import defaultdict
# from io import StringIO

from matplotlib import pyplot as plt
from PIL import Image

# zhulei custom config patch
config = tf.ConfigProto()
config.gpu_options.allow_growth = True  # dynamically grow the memory used on the GPU
sess = tf.Session(config=config)
K.set_session(sess)  # set this TensorFlow session as the default session for Keras

# This is needed since the notebook is stored in the object_detection folder.
sys.path.append("..")
from object_detection.utils import ops as utils_ops

print(tf.__version__)


# ## Env setup

# In[ ]:


myhost = os.uname()[1]
print(">>>>    Hostname: ", myhost)
print("\nCWD: ", os.getcwd())


# ## Object detection imports
# Here are the imports from the object detection module.

# In[ ]:


from utils import label_map_util

from utils import visualization_utils as vis_util


# # Model preparation 

# ## Variables
# 
# Any model exported using the `export_inference_graph.py` tool can be loaded here simply by changing `PATH_TO_CKPT` to point to a new .pb file.  
# 
# By default we use an "SSD with Mobilenet" model here. See the [detection model zoo](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/detection_model_zoo.md) for a list of other models that can be run out-of-the-box with varying speeds and accuracies.

# In[ ]:


# Folder name containing the Trained Obj-det model/graph
MODEL_NAME = 'Axial_1-491_Resnet_Jun142020_graph'

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

# List of the strings that is used to add correct label for each box.
PATH_TO_LABELS = os.path.join('training', 'ak_detection.pbtxt')


# 3 for axial, 1 for sagittal
NUM_CLASSES = 3

# save previews of the overlay images
save_previews = False

if save_previews:
    script_dir = os.getcwd()
    results_dir = os.path.join(script_dir, 'Sag_Resnet_1-491-preview_obj-det-results-Jun222020/')

    if not os.path.isdir(results_dir):
        os.makedirs(results_dir)

# ## Load a (frozen) Tensorflow model into memory.

# In[ ]:


# zhulei updated with compat.v1 and tf.io
detection_graph = tf.Graph()
with detection_graph.as_default():
    od_graph_def = tf.compat.v1.GraphDef()
    with tf.io.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def, name='')



label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)
# zhulei custom modify
#category_index = {1: {'id': 1, 'name': 'left'}, 2: {'id': 2, 'name': 'center'}, 3: {'id': 3, 'name': 'right'}}
print("\ncategory_index: ", category_index)


# ## Helper code

# In[ ]:


# zhulei modify to detect image.mode
def load_image_into_numpy_array(image):
    (im_width, im_height) = image.size
    # 'L' for Grayscale, 'RGB' : for 3 channel images
    channel_dict = {'L':1, 'RGB':3}
    return np.array(image.getdata()).reshape(
        (im_height, im_width, channel_dict[image.mode])).astype(np.uint8)


# # Detection

# In[ ]:


print("current dir: ", os.getcwd())

# If you want to test the code with your images, just add path to the images to the TEST_IMAGE_PATHS.
PATH_TO_TEST_IMAGES_DIR = 'images/test'


generated_pickle = './{}/obj-det.pickle'.format(
    PATH_TO_TEST_IMAGES_DIR)

# TEST_IMAGE_PATHS = [ os.path.join(PATH_TO_TEST_IMAGES_DIR, 'image{}.jpg'.format(i)) for i in range(1, 3) ]
TEST_IMAGE_PATHS = []
for file in os.listdir(PATH_TO_TEST_IMAGES_DIR):
    if '.jpg' in file or '.png' in file or '.JPG' in file:
        TEST_IMAGE_PATHS.append(os.path.join(PATH_TO_TEST_IMAGES_DIR, file))

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

print("\nTest images: ", TEST_IMAGE_PATHS[:5])

print("\nNum test images: ", len(TEST_IMAGE_PATHS))


# In[ ]:


def run_inference_for_single_image(image, graph):
  with graph.as_default():
    with tf.compat.v1.Session() as sess:
      # Get handles to input and output tensors
      ops = tf.compat.v1.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.compat.v1.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[0], image.shape[1])
        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.compat.v1.get_default_graph().get_tensor_by_name('image_tensor:0')

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

      # 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.uint8)
      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


# In[ ]:

detection_result = {}

count_img = 0

for image_path in TEST_IMAGE_PATHS:
    # zhulei add counting img
    print('process: ', str(count_img), image_path)
    count_img += 1

    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 = load_image_into_numpy_array(image)

    # zhulei check for image_np
    if image_np.shape[2] != 3:
        # Duplicating the Content
        image_np = np.broadcast_to(image_np, (image_np.shape[0], image_np.shape[1], 3)).copy()
        ## adding Zeros to other Channels
        ## This adds Red Color stuff in background -- not recommended
        # z = np.zeros(image_np.shape[:-1] + (2,), dtype=image_np.dtype)
        # image_np = np.concatenate((image_np, z), axis=-1)
    
    # 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 = run_inference_for_single_image(image_np, 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'],
      category_index,
      instance_masks=output_dict.get('detection_masks'),
      use_normalized_coordinates=True,
      line_thickness=2)  # default 8

    # zhulei custom
    boxes = output_dict['detection_boxes']
    classes = output_dict['detection_classes']
    scores = output_dict['detection_scores']
    max_boxes_to_draw = 20
    min_score_thresh = 0.5
    box_to_class = {}

    # zhulei custom noting down the box_to_classes
    for i in range(min(max_boxes_to_draw, boxes.shape[0])):
        if scores[i] > min_score_thresh:
            box = tuple(boxes[i].tolist())
            box_to_class[box] = classes[i]

    # for export to pickle later
    detection_result[image_path] = box_to_class

    if save_previews:
        plt.figure(figsize=IMAGE_SIZE)
        # need imshow to save the imgs
        plt.imshow(image_np)
        plt.savefig('{}/{}'.format(
            results_dir, os.path.basename(image_path))
        )

# In[ ]:

# export detection_result dict with pickle
with open(generated_pickle, 'wb') as f:
    pickle.dump(detection_result, f, protocol=pickle.HIGHEST_PROTOCOL)


# verify the generated pickle is saved
# generated_pickle = './{}/{}_detection.pickle'.format(PATH_TO_TEST_IMAGES_DIR, labeler)
print(os.path.exists(generated_pickle))
print(generated_pickle, "exists")
print(os.path.getsize(generated_pickle), "byte")