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/stalyn314/10xmineralmodel/resolve/main/xplx10x_d2.pth",
    "classes": ["minerales", "Afs", "Amp", "Bt", "Ms", "Ol", "Pl", "Px", "Qz"],
    "cfg": None,
    "metadata": None
},
{
    "name": "Version 2 (4-class)",
    "model_path": "https://huggingface.co/stalyn314/10xmineralmodel/resolve/main/10xmodel_d2.pth",
    "classes": ["minerales", "Afs", "Amp", "Bt", "Ms", "Ol", "Pl", "Px", "Qz"],
    "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/mask_rcnn_X_101_32x8d_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, min_score, model_name):
    results = []
    
    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"])

    for img in image:
        # Convert image from BGR to RGB if necessary
        im = img[:,:,::-1]

        # Make the prediction
        outputs = predictor(im)

        v = Visualizer(im, models[model_id]["metadata"], scale=1.2)
        out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
        
        # Convert the result from BGR to RGB
        result_image = out.get_image()
        result_image_rgb = result_image[:, :, ::-1]  # Convert BGR to RGB

        results.append(result_image_rgb)  # Add the processed image to the list
    
    return results  # Return all the results

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 Image"):
        image_input = gr.Image(type="numpy", label="Input Images", elem_id="input_image", multiple=True)

    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_gallery = gr.Gallery(label="Output Images", elem_id="output_images").style(grid=2)

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

    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()