# -*- coding: utf-8 -*-
"""InstColorization.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/github/ericsujw/InstColorization/blob/master/InstColorization.ipynb

## Environment Setting
"""

import os
from os.path import join, isfile
from os import listdir

import torch
from detectron2.utils.logger import setup_logger
setup_logger()

import numpy as np
import cv2

# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg

import sys
from options.train_options import TestOptions
from models import create_model

from tqdm import tqdm_notebook, tqdm

from fusion_dataset import Fusion_Testing_Dataset
from util import util
import multiprocessing

# For testing whether running in a Jupyter notebook,
# or in colab, or in a regular Python shell:
def isnotebook():
    '''
    Tests whether current process is running in a:
          o terminal as a regular Python shell
          o jupyter notebook
          o Google colab
    returns one of {'terminal', 'jupyter', 'colab', None}
    None means could not determine.
    '''
    try:
        # NOTE: get_ipython() won't be defined in a terminal,
        # but will be in Jupyter/Colab/Jupyter-Lab:
        shell = get_ipython().__class__.__name__
        if shell == 'ZMQInteractiveShell':
            return 'jupyter'   # Jupyter notebook or qtconsole
        elif shell == 'TerminalInteractiveShell':
            return 'terminal'  # Terminal running IPython
        elif shell == 'Shell' and  get_ipython().__class__.__module__ == 'google.colab._shell':
            return 'colab'
        else:
            return None  # Other type (?)
    except NameError:
        return 'terminal'      # Probably standard Python interpreter


#from google.colab.patches import cv2_imshow

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml")

if not torch.cuda.is_available():
    cfg.MODEL.DEVICE = 'cpu' 
else:
    cfg.MODEL.DEVICE = 'cuda' 

predictor = DefaultPredictor(cfg)

"""Let's create a bounding box folder to save our prediction results."""

input_dir = "example"
image_list = [f for f in listdir(input_dir) if isfile(join(input_dir, f))]
output_npz_dir = "{0}_bbox".format(input_dir)
if os.path.isdir(output_npz_dir) is False:
    print('Create path: {0}'.format(output_npz_dir))
    os.makedirs(output_npz_dir)

"""Here we simply take L channel as our input 
   and make sure that we can get consistent box 
   prediction results even though the original image is color images."""

print("Ensure consistent box prediction results...")
for image_path in image_list:
    img = cv2.imread(join(input_dir, image_path))
    lab_image = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
    l_channel, a_channel, b_channel = cv2.split(lab_image)
    l_stack = np.stack([l_channel, l_channel, l_channel], axis=2)
    outputs = predictor(l_stack)
    save_path = join(output_npz_dir, image_path.split('.')[0])
    pred_bbox = outputs["instances"].pred_boxes.to(torch.device('cpu')).tensor.numpy()
    pred_scores = outputs["instances"].scores.cpu().data.numpy()
    np.savez(save_path, bbox = pred_bbox, scores = pred_scores)
print("Done ensure consistent box prediction results.")

"""Now we have all the images' prediction results."""

"""### Colorize Images

We first set up some libraries and options
"""

multiprocessing.set_start_method('spawn', True)

torch.backends.cudnn.benchmark = True

sys.argv = [sys.argv[0]]
opt = TestOptions().parse()

"""Then we need to create a results folder to save our predicted color images and read the dataset loader."""

save_img_path = opt.results_img_dir
if os.path.isdir(save_img_path) is False:
    print('Create path: {0}'.format(save_img_path))
    os.makedirs(save_img_path)
opt.batch_size = 1
if torch.cuda.is_available():
    opt.gpu_ids = [0]
else:
    opt.gpu_ids = []
dataset = Fusion_Testing_Dataset(opt, -1)
dataset_loader = torch.utils.data.DataLoader(dataset, batch_size=opt.batch_size)

dataset_size = len(dataset)
print('#Testing images = %d' % dataset_size)

"""Load the pre-trained model."""

model = create_model(opt)
model.setup_to_test('coco_finetuned_mask_256_ffs')

"""Start to colorize every images in `dataset_loader`."""

count_empty = 0
if isnotebook() in ['colab', 'jupyter']:
    data_iterator = tqdm_notebook(dataset_loader)
else:
    data_iterator = tqdm(dataset_loader)
    
for data_raw in data_iterator:
    if torch.cuda.is_available():
        data_raw['full_img'][0] = data_raw['full_img'][0].cuda()
    else:
        data_raw['full_img'][0] = data_raw['full_img'][0].cpu()
    if data_raw['empty_box'][0] == 0:
        if torch.cuda.is_available():
            data_raw['cropped_img'][0] = data_raw['cropped_img'][0].cuda()
        else:
            data_raw['cropped_img'][0] = data_raw['cropped_img'][0].cpu()
        box_info = data_raw['box_info'][0]
        box_info_2x = data_raw['box_info_2x'][0]
        box_info_4x = data_raw['box_info_4x'][0]
        box_info_8x = data_raw['box_info_8x'][0]
        cropped_data = util.get_colorization_data(data_raw['cropped_img'], opt, ab_thresh=0, p=opt.sample_p)
        full_img_data = util.get_colorization_data(data_raw['full_img'], opt, ab_thresh=0, p=opt.sample_p)
        model.set_input(cropped_data)
        model.set_fusion_input(full_img_data, [box_info, box_info_2x, box_info_4x, box_info_8x])
        model.forward()
    else:
        count_empty += 1
        full_img_data = util.get_colorization_data(data_raw['full_img'], opt, ab_thresh=0, p=opt.sample_p)
        model.set_forward_without_box(full_img_data)
    model.save_current_imgs(join(save_img_path, data_raw['file_id'][0] + '.png'))
print('{0} images without bounding boxes'.format(count_empty))

img_name_list = ['000000022969', '000000023781', '000000046872', '000000050145']
show_index = 1

orig_img = cv2.imread('example/'+img_name_list[show_index]+'.jpg')
lab_image = cv2.cvtColor(orig_img, cv2.COLOR_BGR2LAB)
l_channel, _, _ = cv2.split(lab_image)

img = cv2.imread('results/'+img_name_list[show_index]+'.png')
lab_image = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
_, a_pred, b_pred = cv2.split(lab_image)
a_pred = cv2.resize(a_pred, (l_channel.shape[1], l_channel.shape[0]))
b_pred = cv2.resize(b_pred, (l_channel.shape[1], l_channel.shape[0]))
gray_color = np.ones_like(a_pred) * 128

gray_image = cv2.cvtColor(np.stack([l_channel, gray_color, gray_color], 2), cv2.COLOR_LAB2BGR)
color_image = cv2.cvtColor(np.stack([l_channel, a_pred, b_pred], 2), cv2.COLOR_LAB2BGR)

# save_img_path = 'results_origin/'
# if os.path.isdir(save_img_path) is False:
#     print('Create path: {0}'.format(save_img_path))
#     os.makedirs(save_img_path)

# cv2.imwrite('results_origin/'+img_name_list[show_index]+'.png', color_image)

#cv2_imshow(np.concatenate([gray_image, color_image], 1))
cv2.imshow(np.concatenate([orig_img, gray_image, color_image], 1))
input("Press ENTER to close image and exit program: ")