# -*- coding: utf-8 -*-
"""Copy of inspect_balloon_model.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1MLAg2B1HwiLTDBG8xdXnPir66v8KTjft

# Mask R-CNN - Inspect Ballon Trained Model

Code and visualizations to test, debug, and evaluate the Mask R-CNN model.

## Setup
"""

import os
import sys
import skimage.io
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import cv2


# Root directory of the project
ROOT_DIR = os.path.abspath("../../")
ROOT_DIR = os.path.join(ROOT_DIR,"Mask_RCNN")
#how many frames you want to take(intervall of 1 is every frame)
INTERVALL = 10

# Import Mask RCNN
sys.path.append(ROOT_DIR)  # To find local version of the library
from mrcnn import visualize
from mrcnn.visualize import display_images
import mrcnn.model as modellib
from mrcnn.model import log
import mrcnn.utils as utils

from samples.balloon import balloon

# %matplotlib inline 

# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")

# Path to Ballon trained weights
# You can download this file from the Releases page
# https://github.com/matterport/Mask_RCNN/releases
# This path will differ depending on where you saved the Weights and how they are named
BALLON_WEIGHTS_PATH = os.path.join(ROOT_DIR, "data","model", "mask_rcnn_tool_disassembly_0030.h5")

"""## Configurations"""

config = balloon.BalloonConfig()
BALLOON_DIR = os.path.join(ROOT_DIR, "data","gear")


# Override the training configurations with a few
# changes for inferencing.
class InferenceConfig(config.__class__):
    # Run detection on one image at a time
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1


config = InferenceConfig()
config.display()

"""## Notebook Preferences"""

# Device to load the neural network on.
# Useful if you're training a model on the same 
# machine, in which case use CPU and leave the
# GPU for training.
DEVICE = "/cpu:0"  # /cpu:0 or /gpu:0

# Inspect the model in training or inference modes
# values: 'inference' or 'training'
# TODO: code for 'training' test mode not ready yet
TEST_MODE = "inference"


def get_ax(rows=1, cols=1, size=16):
    """Return a Matplotlib Axes array to be used in
    all visualizations in the notebook. Provide a
    central point to control graph sizes.
    
    Adjust the size attribute to control how big to render images
    """
    _, ax = plt.subplots(rows, cols, figsize=(size * cols, size * rows))
    return ax


"""## Load Validation Dataset"""

# Load validation dataset
dataset = balloon.BalloonDataset()
dataset.load_balloon(BALLOON_DIR, "val")

# Must call before using the dataset
dataset.prepare()



"""## Load Model"""

# Create model in inference mode
with tf.device(DEVICE):
    model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR,
                              config=config)


# Run this cell to mount your Google Drive.
# click on the link and copy the password


# copy the weights from your drive to the logs directory


# Set path to balloon weights file

# Download file from the Releases page and set its path
# https://github.com/matterport/Mask_RCNN/releases
# weights_path = "/path/to/mask_rcnn_balloon.h5"

# Or, load the last model you trained
weights_path = BALLON_WEIGHTS_PATH
images_path = os.path.join(ROOT_DIR,"data","video")

# Load weights
print("Loading weights ", weights_path)
model.load_weights(weights_path, by_name=True)

"""## Run Detection"""
cam = cv2.VideoCapture(os.path.join(images_path,"video1.mp4"))
result = []
counter = 0
while(True):
    ret, image = cam.read()

    if not ret:
        break
    if (not counter%INTERVALL==0):
        counter+=1
        continue
    # Run object detection
    results = model.detect([image], verbose=1)
    # Display results
    ax = get_ax(1)
    r = results[0]
    visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'],
                               dataset.class_names, r['scores'], ax=ax,
                                title="Predictions")

    for index,roi in enumerate(r['rois']):
        result.append([counter//INTERVALL,np.array(roi),r['class_ids'][index]])

    counter+=1

#comment this line if you do not want to inspect the results
plt.show()

file = open("allAOIs","w")

for counter, aoi in enumerate(result):

    text = str(aoi[0]) + " " + str(aoi[1][0]) + " " + str(aoi[1][1]) + " " + str(aoi[1][2]) + " " + str(aoi[1][3]) + " " + dataset.class_names[aoi[2]] + "\n"
    file.write(text)
file.close()
