import os
os.system("pip install git+https://github.com/facebookresearch/detectron2.git")
os.system("git clone https://github.com/Visual-AI/Mr.DETR.git MrDETR && cd MrDETR && rm -f requirements.txt && cd ..")
os.system("cp multi_scale_deform_attn.py MrDETR/detrex/layers/ && cd MrDETR && pip install . & cd ..")
import sys 
sys.path.append("MrDETR/")

import gradio as gr
from demo.predictors import VisualizationDemo
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import LazyConfig, instantiate
import numpy as np

if __name__ == "__main__":
    gr.close_all()
    cfg = LazyConfig.load("MrDETR/projects/mr_detr_align/configs/deformable_detr_swinl_two_stage_12ep_plusplus.py")
    cfg["model"].device = "cpu"
    cfg["train"].device = "cpu"
    model = instantiate(cfg.model)
    checkpointer = DetectionCheckpointer(model)
    checkpointer.load("https://github.com/Visual-AI/Mr.DETR/releases/download/weights/MrDETR_align_swinL_12ep_900q_safe.pth")

    model.eval()
    vis_demo = VisualizationDemo(
        model=model,
        min_size_test=800,
        max_size_test=1333,
        img_format="RGB",
        metadata_dataset="coco_2017_val",
    )
    
    def inference(img, confidence):
        img = np.array(img)
        _, results = vis_demo.run_on_image(img, confidence)
        return results 
    
    demo = gr.Interface(
        fn=inference,
        inputs=[
            gr.Image(type="pil", image_mode="RGB"),
            gr.Number(precision=2, minimum=0.0, maximum=1.0, value=0.5)
        ],
        outputs="image",
        examples=[
            ["MrDETR/assets/000000014226.jpg", 0.5],
            ["MrDETR/assets/000000028449.jpg", 0.5]
        ]
    )
    demo.launch()