import spaces
import os
import gradio as gr
import shutil
import sys 
import subprocess




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 ..")
subprocess.run(
    shlex.split(
        "pip install detrex-0.3.0-cp310-cp310-linux_x86_64.whl"
    )
)
    
sys.path.append("MrDETR/")
install_setup()

detrex-0.3.0-cp310-cp310-linux_x86_64.whl


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


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 = "cuda"
    cfg["train"].device = "cuda"

    # @spaces.GPU(duration=40, progress=gr.Progress(track_tqdm=True))
    # def 
    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()
    model.cuda()
    vis_demo = VisualizationDemo(
        model=model,
        min_size_test=800,
        max_size_test=1333,
        img_format="RGB",
        metadata_dataset="coco_2017_val",
    )
    
    @spaces.GPU
    def inference(img, confidence):
        img = np.array(img)
        _, results = vis_demo.run_on_image(img, confidence)
        results = Image.fromarray(results.get_image()[:, :, ::-1])
        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)
            gr.Slider(minimum=0.0, maximum=1.0, value=0.5, step=0.05)
        ],
        outputs="image",
        examples=[
            ["MrDETR/assets/000000014226.jpg", 0.5],
            ["MrDETR/assets/000000028449.jpg", 0.3],
            ["MrDETR/assets/000000070048.jpg", 0.5],
            ["MrDETR/assets/000000218997.jpg", 0.5],
            ["MrDETR/assets/000000279774.jpg", 0.5],
            ["MrDETR/assets/000000434459.jpg", 0.5],
            ["MrDETR/assets/000000448448.jpg", 0.5],
            ["MrDETR/assets/000000560474.jpg", 0.5],
        ],
        title="[CVPR 2025] Mr. DETR: Instructive Multi-Route Training for Detection Transformers",
        description='''
            [![Paper](https://img.shields.io/badge/arXiv-2412.10028-red)](https://arxiv.org/abs/2412.10028)
            [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/mr-detr-instructive-multi-route-training-for/object-detection-on-coco-2017-val)](https://paperswithcode.com/sota/object-detection-on-coco-2017-val?p=mr-detr-instructive-multi-route-training-for)
        '''
    )
    demo.launch()