######## Image Object Detection Using Tensorflow-trained Classifier #########
#
# Author: Evan Juras
# Date: 1/15/18
# Description: 
# This program uses a TensorFlow-trained classifier to perform object detection.
# It loads the classifier uses it to perform object detection on an image.
# It draws boxes and scores around the objects of interest in the image.

## Some of the code is copied from Google's example at
## https://github.com/tensorflow/models/blob/master/research/object_detection/object_detection_tutorial.ipynb

## and some is copied from Dat Tran's example at
## https://github.com/datitran/object_detector_app/blob/master/object_detection_app.py

## but I changed it to make it more understandable to me.

# Import packages
import os
import cv2
import numpy as np
import tensorflow as tf
import sys
import seaborn as sb
import matplotlib.pyplot as plt
# This is needed since the notebook is stored in the object_detection folder.
sys.path.append("..")

# Import utilites
from utils import label_map_util
from utils import visualization_utils as vis_util

# Name of the directory containing the object detection module we're using
MODEL_NAME = 'inference_graph'
IMAGE_NAME = 'input.JPG'

# Grab path to current working directory
CWD_PATH = os.getcwd()

# Path to frozen detection graph .pb file, which contains the model that is used
# for object detection.
PATH_TO_CKPT = os.path.join(CWD_PATH,'frozen_inference_graph.pb')

# Path to label map file
PATH_TO_LABELS = os.path.join(CWD_PATH,'labelmap.pbtxt')

# Path to image
PATH_TO_IMAGE = os.path.join(CWD_PATH,IMAGE_NAME)

# Number of classes the object detector can identify
NUM_CLASSES = 1

# Load the label map.
# Label maps map indices to category names, so that when our convolution
# network predicts `5`, we know that this corresponds to `king`.
# Here we use internal utility functions, but anything that returns a
# dictionary mapping integers to appropriate string labels would be fine
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)

# Load the Tensorflow model into memory.
detection_graph = tf.Graph()
with detection_graph.as_default():
    od_graph_def = tf.GraphDef()
    with tf.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='')

    sess = tf.Session(graph=detection_graph)

# Define input and output tensors (i.e. data) for the object detection classifier

# Input tensor is the image
image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')

# Output tensors are the detection boxes, scores, and classes
# Each box represents a part of the image where a particular object was detected
detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')

# Each score represents level of confidence for each of the objects.
# The score is shown on the result image, together with the class label.
detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')

# Number of objects detected
num_detections = detection_graph.get_tensor_by_name('num_detections:0')

# Load image using OpenCV and
# expand image dimensions to have shape: [1, None, None, 3]
# i.e. a single-column array, where each item in the column has the pixel RGB value
data=np.ones((10,10),dtype=int)
image = cv2.imread(PATH_TO_IMAGE)
image=  cv2.resize(image,(3000,3000))
image2=image
image3=np.ones((3000,3000,3),np.uint8)
x1,y1,x2,y2=0,0,300,300
for y in range(10):
    x1,x2=0,300
    for x in range(10):
        img1=image[y1:y2,x1:x2]
        image_expanded = np.expand_dims(img1, axis=0)

    # Perform the actual detection by running the model with the image as input
        (boxes, scores, classes, num) = sess.run(
        [detection_boxes, detection_scores, detection_classes, num_detections],
        feed_dict={image_tensor: image_expanded})

    # Draw the results of the detection (aka 'visulaize the results')

        vis_util.visualize_boxes_and_labels_on_image_array(
        img1,
        np.squeeze(boxes),
        np.squeeze(classes).astype(np.int32),
        np.squeeze(scores),
        category_index,
        use_normalized_coordinates=True,
        line_thickness=8,
        min_score_thresh=0.60)

        final_score = np.squeeze(scores)    
        count = 0
        for i in range(100):
            if scores is None or final_score[i] > 0.6:
                count = count + 1
        im_height, im_width = img1.shape[:2]
        img2=np.ones((300,300,3),np.uint8)
        '''
        print(boxes.shape)
        for i in range(0,1,1):
            position = boxes[0][0]
            print(position)
            (xmin, xmax, ymin, ymax) = (position[1]*im_width, position[3]*im_width, position[0]*im_height, position[2]*im_height)
            xmin=xmin.astype(int)
            xmax=xmax.astype(int)
            ymin=ymin.astype(int)
            ymax=ymax.astype(int)
            a=(xmin+xmax)/2
            a=a.astype(int)
            img2=img2.astype(int)
            img2=cv2.line(img2,(a,ymin),(a,ymax),(0,255,0),5)
        #print(xmin, xmax, ymin, ymax)
        '''

        #print(count)
        image2[y1:y2,x1:x2]=img1
        #image3[y1:y2,x1:x2]=img2
        # All the results have been drawn on image. Now display the image.
        #cv2.imshow(str(x)+str(y)+'Object detector', img1)
        #cv2.waitKey(0)
        data[y][x]=count
        # Clean up
        #cv2.destroyAllWindows()

        x1+=300
        x2+=300
    y1+=300
    y2+=300
cv2.imwrite('output.jpg',image2)
cv2.imshow('res',image2)
cv2.waitKey(0)
cv2.destroyAllWindows()
#heat_map=sb.heatmap(data,yticklabels=False,xticklabels=False,cmap="YlGnBu")
#plt.title("Heatmap for the density of corn plants")

plt.show()
