import os 

try:
    import detectron2
except:
    os.system('pip install lib/detectron2')

import numpy as np
import os, json, cv2, random

from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog


MODEL_YAML='COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml'

cfg = get_cfg()

cfg.merge_from_file(model_zoo.get_config_file(MODEL_YAML))
#cfg.DEVICE = 'cpu'
cfg.MODEL.DEVICE = 'cpu'
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.MODEL.WEIGHTS = "weights/model_final_2d9806.pkl"

predictor = DefaultPredictor(cfg)


import gradio as gr
from PIL import Image

def infer(input_filename):
    # Predictor takes BGR.
    cv2_image = cv2.imread(input_filename)
    v = Visualizer(cv2_image[:, :, ::-1], # Suppose RGB
                   MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
                   scale=1.2)

    results = predictor(cv2_image)

    output_image = v.draw_instance_predictions(results["instances"].to("cpu")).get_image()
    return Image.fromarray(np.uint8(output_image)).convert('RGB')


with gr.Blocks(title="Detectron2 Object Detection - ClassCat",
                    css=".gradio-container {background:lightyellow;}"
               ) as demo:
    #sample_index = gr.State([])

    gr.HTML("""<div style="font-family:'Times New Roman', 'Serif'; font-size:16pt; font-weight:bold; text-align:center; color:royalblue;">Detectron2 Object Detection</div>""")

    gr.HTML("""<h4 style="color:navy;">1-a. Select an example by clicking a thumbnail below.</h4>""")

    gr.HTML("""<h4 style="color:navy;">1-b. Or upload an image by clicking on the canvas.</h4>""")

    with gr.Row():
        input_image = gr.Image(label="Input image", type="filepath")
        output_image = gr.Image(label="Output image with predicted instances", type="numpy")
    
    gr.Examples(['samples/detectron2.png', 'samples/cat.jpg', 'samples/hotdog.jpg'], inputs=input_image)

    gr.HTML("""<br/>""")
    gr.HTML("""<h4 style="color:navy;">2. Then, click "Infer" button to predict object instances. It will take about 15-20 seconds (on cpu)</h4>""")

    send_btn = gr.Button("Infer")
    send_btn.click(fn=infer, inputs=[input_image], outputs=[output_image])

    gr.HTML("""<br/>""")
    gr.HTML("""<h4 style="color:navy;">Reference</h4>""")
    gr.HTML("""<ul>""")
    gr.HTML("""<li><a href="https://colab.research.google.com/drive/16jcaJoc6bCFAQ96jDe2HwtXj7BMD_-m5" target="_blank">Detectron2 Tutorial</a>""")
    gr.HTML("""</ul>""")


#demo.queue()
demo.launch() # debug=True)


### EOF ###