# Path: utils/kitti.py
# KITTI_TRAIN_DIR_PREFIX = '/usr/local/google/home/limeng/Downloads/kitti/data_road/training/image_2/'
# KITTI_GT_DIR_PREFIX = '/usr/local/google/home/limeng/Downloads/kitti/data_road/training/gt_image_2/'
# UM_TRAIN_TEMPLATE = "um_0000%02d.png"
# UU_TRAIN_TEMPLATE = "uu_0000%02d.png"
# UMM_TRAIN_TEMPLATE = "umm_0000%02d.png"
# UU_GT_ROAD_TEMPLATE = "uu_road_0000%02d.png"
# UM_GT_LANE_TEMPLATE = "um_lane_0000%02d.png"
# UM_GT_ROAD_TEMPLATE = "um_road_0000%02d.png"
# UMM_GT_ROAD_TEMPLATE = "umm_road_0000%02d.png"
# class Kitti(object):
#     def __init__(self):
#     def _read_datasets(self,
#                        train_data_dir=KITTI_TRAIN_DIR_PREFIX,
#                        gt_data_dir=KITTI_GT_DIR_PREFIX,
#                        cat='uu'):
#     def next_batch(self, batch_id):
# def get_training_data(file_name):
# def get_ground_truth(fileNameGT):
# def main(_):


import copy
import tensorflow as tf
import scipy as scp
import scipy.misc
import matplotlib as mpl
import matplotlib.cm
from utils import kitti
from self_driving.road_seg import fcn8_vgg

def color_image(image, num_classes=20):
    norm = mpl.colors.Normalize(vmin=0., vmax=num_classes)
    mycm = mpl.cm.get_cmap('Set1')
    return mycm(norm(image))


def save_output(index, training_image, prediction, label):
    prediction_label = 1 - prediction[0]
    output_image = copy.copy(training_image)
    # Save prediction
    up_color = color_image(prediction[0], 2)
    scp.misc.imsave('output/decision_%d.png' % index, up_color)
    # Merge true positive with training images' green channel
    true_positive = prediction_label * label[..., 0][0]
    merge_green = (1 - true_positive) * training_image[..., 1] + true_positive * 255
    output_image[..., 1] = merge_green
    # Merge false positive with training images' red channel
    false_positive = prediction_label * label[..., 1][0]
    merge_red = (1 - false_positive) * training_image[..., 0] + false_positive * 255
    output_image[..., 0] = merge_red
    # Merge false negative with training images' blue channel
    false_negative = (1 - prediction_label) * label[..., 0][0]
    merge_blue = (1 - false_negative) * training_image[..., 2] + false_negative * 255
    output_image[..., 2] = merge_blue
    # Save images
    scp.misc.imsave('merge/decision_%d.png' % index, output_image)


def main(_):

# Path: utils/kitti.py
# KITTI_TRAIN_DIR_PREFIX = '/usr/local/google/home/limeng/Downloads/kitti/data_road/training/image_2/'
# KITTI_GT_DIR_PREFIX = '/usr/local/google/home/limeng/Downloads/kitti/data_road/training/gt_image_2/'
# UM_TRAIN_TEMPLATE = "um_0000%02d.png"
# UU_TRAIN_TEMPLATE = "uu_0000%02d.png"
# UMM_TRAIN_TEMPLATE = "umm_0000%02d.png"
# UU_GT_ROAD_TEMPLATE = "uu_road_0000%02d.png"
# UM_GT_LANE_TEMPLATE = "um_lane_0000%02d.png"
# UM_GT_ROAD_TEMPLATE = "um_road_0000%02d.png"
# UMM_GT_ROAD_TEMPLATE = "umm_road_0000%02d.png"
# class Kitti(object):
#     def __init__(self):
#     def _read_datasets(self,
#                        train_data_dir=KITTI_TRAIN_DIR_PREFIX,
#                        gt_data_dir=KITTI_GT_DIR_PREFIX,
#                        cat='uu'):
#     def next_batch(self, batch_id):
# def get_training_data(file_name):
# def get_ground_truth(fileNameGT):
# def main(_):


# Path: self_driving/road_seg/convnet.py
import copy
import tensorflow as tf
import scipy as scp
import scipy.misc
import matplotlib as mpl
import matplotlib.cm
from utils import kitti
from self_driving.road_seg import fcn8_vgg


def color_image(image, num_classes=20):
    norm = mpl.colors.Normalize(vmin=0., vmax=num_classes)
    mycm = mpl.cm.get_cmap('Set1')
    return mycm(norm(image))


def save_output(index, training_image, prediction, label):
    prediction_label = 1 - prediction[0]
    output_image = copy.copy(training_image)
    # Save prediction
    up_color = color_image(prediction[0], 2)
    scp.misc.imsave('output/decision_%d.png' % index, up_color)
    # Merge true positive with training images' green channel
    true_positive = prediction_label * label[..., 0][0]
    merge_green = (1 - true_positive) * training_image[..., 1] + true_positive * 255
    output_image[..., 1] = merge_green
    # Merge false positive with training images' red channel
    false_positive = prediction_label * label[..., 1][0]
    merge_red = (1 - false_positive) * training_image[..., 0] + false_positive * 255
    output_image[..., 0] = merge_red
    # Merge false negative with training images' blue channel
    false_negative = (1 - prediction_label) * label[..., 0][0]
    merge_blue = (1 - false_negative) * training_image[..., 2] + false_negative * 255
    output_image[..., 2] = merge_blue
    # Save images
    scp.misc.imsave('merge/decision_%d.png' % index, output_image)


def main(_):

    kitti_data = kitti.Kitti()