app.py

import spaces
import os
import gradio as gr
import shutil

@spaces.GPU
def find_cuda():
    # Check if CUDA_HOME or CUDA_PATH environment variables are set
    cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')

    if cuda_home and os.path.exists(cuda_home):
        return cuda_home

    # Search for the nvcc executable in the system's PATH
    nvcc_path = shutil.which('nvcc')

    if nvcc_path:
        # Remove the 'bin/nvcc' part to get the CUDA installation path
        cuda_path = os.path.dirname(os.path.dirname(nvcc_path))
        return cuda_path

    return None

os.environ["CUDA_HOME"] = find_cuda()

if cuda_path:
    print(f"CUDA installation found at: {cuda_path}")
else:
    print("CUDA installation not found")



@spaces.GPU(duration=180)
def install_setup():
    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 ..")
    # os.system("cd MrDETR && export CUDA_HOME=$(dirname $(dirname $(which nvcc))) && pip install . & cd ..")
    os.system("cd MrDETR && pip install . & cd ..")
    import sys 
    sys.path.append("MrDETR/")
install_setup()


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