try:
    import detectron2
except:
    import os 
    os.system('pip install git+https://github.com/facebookresearch/detectron2.git')

import cv2

from matplotlib.pyplot import axis
import gradio as gr
import requests
import numpy as np
from torch import nn
import requests

import torch

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


models = [
{
    "name": "Version 1 (2-class)",
    "model_path": "https://huggingface.co/dbmdz/detectron2-model/resolve/main/model_final.pth",
    "classes": ["Illumination", "Illustration"],
    "cfg": None,
    "metadata": None
},
{
    "name": "Version 2 (4-class)",
    "model_path": "https://huggingface.co/dbmdz/detectron2-v2-model/resolve/main/model_final.pth",
    "classes": ["ILLUSTRATION", "OTHER", "STAMP", "INITIAL"],
    "cfg": None,
    "metadata": None
},
]

model_name_to_id = {model["name"] : id_ for id_, model in enumerate(models)}

for model in models:

    model["cfg"] = get_cfg()
    model["cfg"].merge_from_file("./configs/detectron2/faster_rcnn_R_50_FPN_3x.yaml")
    model["cfg"].MODEL.ROI_HEADS.NUM_CLASSES = len(model["classes"])
    model["cfg"].MODEL.WEIGHTS = model["model_path"]
    
    model["metadata"] = MetadataCatalog.get(model["name"])
    model["metadata"].thing_classes = model["classes"]

    if not torch.cuda.is_available():
      model["cfg"].MODEL.DEVICE = "cpu"


def inference(image_url, image, min_score, model_name):
    if image_url:
        r = requests.get(image_url)
        if r:
            im = np.frombuffer(r.content, dtype="uint8")
            im = cv2.imdecode(im, cv2.IMREAD_COLOR)
    else:
        # Model expect BGR!
        im = image[:,:,::-1]

    model_id = model_name_to_id[model_name]

    models[model_id]["cfg"].MODEL.ROI_HEADS.SCORE_THRESH_TEST = min_score
    predictor = DefaultPredictor(models[model_id]["cfg"])

    outputs = predictor(im)

    v = Visualizer(im, models[model_id]["metadata"], scale=1.2)
    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))

    return out.get_image()


title = "# DBMDZ Detectron2 Model Demo"
description = """
This demo introduces an interactive playground for our trained Detectron2 model.

Currently, two models are supported that were trained on manually annotated segments from digitized books:

* [Version 1 (2-class)](https://huggingface.co/dbmdz/detectron2-model): This model can detect *Illustration* or *Illumination* segments on a given page.
* [Version 2 (4-class)](https://huggingface.co/dbmdz/detectron2-v2-model): This model is more powerful and can detect *Illustration*, *Stamp*, *Initial* or *Other* segments on a given page.
"""
footer = "Made in Munich with ❤️ and 🥨."

with gr.Blocks() as demo:
    gr.Markdown(title)
    gr.Markdown(description)
    
    with gr.Tab("From URL"):
        url_input = gr.Textbox(label="Image URL", placeholder="https://api.digitale-sammlungen.de/iiif/image/v2/bsb10483966_00008/full/500,/0/default.jpg")
    
    with gr.Tab("From Image"):
        image_input = gr.Image(type="numpy", label="Input Image")

    min_score = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Minimum score")

    model_name = gr.Radio(choices=[model["name"] for model in models], value=models[0]["name"], label="Select Detectron2 model")

    output_image = gr.Image(type="pil", label="Output")

    inference_button = gr.Button("Submit")
    
    inference_button.click(fn=inference, inputs=[url_input, image_input, min_score, model_name], outputs=output_image)

    gr.Markdown(footer)

demo.launch()

#gr.Interface(
#    inference,
#    [gr.inputs.Textbox(label="Image URL", placeholder="https://api.digitale-sammlungen.de/iiif/image/v2/bsb10483966_00008/full/500,/0/default.jpg"),
#     gr.inputs.Image(type="numpy", label="Input Image"),
#     gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Minimum score"),
#     gr.Radio(choices=[model["name"] for model in models], value=models[0]["name"], label="Select Detectron2 model"),
#    ], 
#    gr.outputs.Image(type="pil", label="Output"),
#    title=title,
#    description=description,
#    article=article,
#    examples=[]).launch()