app.py

"""
Copyright (c) 2024 The D-FINE Authors. All Rights Reserved.
"""
import gradio as gr
import os
import sys
import torch
import torch.nn as nn
import torchvision.transforms as T
import supervision as sv
from PIL import Image
import requests
import yaml
import numpy as np

from src.core import YAMLConfig


model_configs = {
    "dfine_n_coco":
        {"cfgfile": "configs/dfine/dfine_hgnetv2_n_coco.yml",
         "classinfofile": "configs/coco.yml",
         "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_n_coco.pth"},
    "dfine_s_coco":
        {"cfgfile": "configs/dfine/dfine_hgnetv2_s_coco.yml",
         "classinfofile": "configs/coco.yml",
         "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_s_coco.pth"},
    "dfine_m_coco":
        {"cfgfile": "configs/dfine/dfine_hgnetv2_m_coco.yml",
         "classinfofile": "configs/coco.yml",
         "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_m_coco.pth"},
    "dfine_l_coco":
        {"cfgfile": "configs/dfine/dfine_hgnetv2_l_coco.yml",
         "classinfofile": "configs/coco.yml",
         "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_coco.pth"},
    "dfine_x_coco":
        {"cfgfile": "configs/dfine/dfine_hgnetv2_x_coco.yml",
         "classinfofile": "configs/coco.yml",
         "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_x_coco.pth"},
    "dfine_s_obj365":
        {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_s_obj365.yml",
         "classinfofile": "configs/obj365.yml",
         "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_s_obj365.pth"},
    "dfine_m_obj365":
        {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_m_obj365.yml",
         "classinfofile": "configs/obj365.yml",
         "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_m_obj365.pth"},
    "dfine_l_obj365":
        {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_l_obj365.yml",
         "classinfofile": "configs/obj365.yml",
         "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_obj365.pth"},
    "dfine_l_obj365_e25":
        {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_l_obj365.yml",
         "classinfofile": "configs/obj365.yml",
         "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_obj365_e25.pth"},
    "dfine_x_obj365":
        {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_x_obj365.yml",
         "classinfofile": "configs/obj365.yml",
         "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_x_obj365.pth"},
    "dfine_s_obj2coco":
        {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_s_obj2coco.yml",
         "classinfofile": "configs/coco.yml",
         "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_s_obj2coco.pth"},
    "dfine_m_obj2coco":
        {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_m_obj2coco.yml",
         "classinfofile": "configs/coco.yml",
         "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_m_obj2coco.pth"},
    "dfine_l_obj2coco_e25":
        {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_l_obj2coco.yml",
         "classinfofile": "configs/coco.yml",
         "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_l_obj2coco_e25.pth"},
    "dfine_x_obj2coco":
        {"cfgfile": "configs/dfine/objects365/dfine_hgnetv2_x_obj2coco.yml",
         "classinfofile": "configs/coco.yml",
         "weights": "https://github.com/Peterande/storage/releases/download/dfinev1.0/dfine_x_obj2coco.pth"},
}


def download_weights(model_name):
    """Download model weights if not already present"""
    weights_url = model_configs[model_name]["weights"]
    # Directory path to save weight files
    weights_dir = os.path.join(os.path.dirname(__file__), "weights")
    # Weight file path
    weights_path = os.path.join(weights_dir, model_name + ".pth")
    
    # Create weights directory if it doesn't exist
    if not os.path.exists(weights_dir):
        os.makedirs(weights_dir)
        print(f"Created directory: {weights_dir}")
    
    # Check if file already exists
    if os.path.exists(weights_path):
        print(f"Weights file already exists at: {weights_path}")
        return weights_path
    
    # Download file
    print(f"Downloading weights from {weights_url} to {weights_path}...")
    
    response = requests.get(weights_url, stream=True)
    response.raise_for_status()  # Check for download errors
    
    with open(weights_path, 'wb') as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)
    
    print(f"Downloaded weights to: {weights_path}")
    return weights_path


def process_image_for_gradio(model, device, image, model_name, threshold=0.4):
    """Process image function for Gradio interface"""
    if isinstance(image, np.ndarray):
        # Convert NumPy array to PIL image
        im_pil = Image.fromarray(image)
    else:
        im_pil = image
    
    # Load class information
    classinfofile = model_configs[model_name]["classinfofile"]
    classinfo = yaml.load(open(classinfofile, "r"), Loader=yaml.FullLoader)["names"]
    indexing_method = "0-based" if "coco" in classinfofile else "1-based"
    
    w, h = im_pil.size
    orig_size = torch.tensor([[w, h]]).to(device)

    transforms = T.Compose(
        [
            T.Resize((640, 640)),
            T.ToTensor(),
        ]
    )
    im_data = transforms(im_pil).unsqueeze(0).to(device)

    output = model(im_data, orig_size)
    labels, boxes, scores = output

    # Visualize results
    detections = sv.Detections(
        xyxy=boxes[0].detach().cpu().numpy(),
        confidence=scores[0].detach().cpu().numpy(),
        class_id=labels[0].detach().cpu().numpy().astype(int),
    )
    detections = detections[detections.confidence > threshold]

    text_scale = sv.calculate_optimal_text_scale(resolution_wh=im_pil.size)
    line_thickness = sv.calculate_optimal_line_thickness(resolution_wh=im_pil.size)

    box_annotator = sv.BoxAnnotator(thickness=line_thickness)
    label_annotator = sv.LabelAnnotator(text_scale=text_scale, smart_position=True)

    label_texts = [
        f"{classinfo[class_id if indexing_method == '0-based' else class_id - 1]} {confidence:.2f}"
        for class_id, confidence
        in zip(detections.class_id, detections.confidence)
    ]

    result_image = im_pil.copy()
    result_image = box_annotator.annotate(scene=result_image, detections=detections)
    result_image = label_annotator.annotate(
        scene=result_image,
        detections=detections,
        labels=label_texts
    )

    detection_info = [
        f"{classinfo[class_id if indexing_method == '0-based' else class_id - 1]}: {confidence:.2f}, bbox: [{xyxy[0]:.1f}, {xyxy[1]:.1f}, {xyxy[2]:.1f}, {xyxy[3]:.1f}]"
        for class_id, confidence, xyxy
        in zip(detections.class_id, detections.confidence, detections.xyxy)
    ]
    
    return result_image, "\n".join(detection_info)


class ModelWrapper(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.model = cfg.model.deploy()
        self.postprocessor = cfg.postprocessor.deploy()

    def forward(self, images, orig_target_sizes):
        outputs = self.model(images)
        outputs = self.postprocessor(outputs, orig_target_sizes)
        return outputs


def load_model(model_name):
    cfgfile = model_configs[model_name]["cfgfile"]
    weights_path = download_weights(model_name)
    
    cfg = YAMLConfig(cfgfile, resume=weights_path)

    if "HGNetv2" in cfg.yaml_cfg:
        cfg.yaml_cfg["HGNetv2"]["pretrained"] = False

    checkpoint = torch.load(weights_path, map_location="cpu")
    state = checkpoint["ema"]["module"] if "ema" in checkpoint else checkpoint["model"]
    
    cfg.model.load_state_dict(state)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = ModelWrapper(cfg).to(device)
    model.eval()

    return model, device


# Dictionary to store loaded models
loaded_models = {}

def process_image(image, model_name, confidence_threshold):
    """Main processing function for Gradio interface"""
    global loaded_models
    
    # Load model if not already loaded
    if model_name not in loaded_models:
        print(f"Loading model: {model_name}")
        model, device = load_model(model_name)
        loaded_models[model_name] = (model, device)
    else:
        print(f"Using cached model: {model_name}")
        model, device = loaded_models[model_name]
    
    # Process the image
    return process_image_for_gradio(model, device, image, model_name, confidence_threshold)


# Create Gradio interface
demo = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Dropdown(
            choices=list(model_configs.keys()), 
            value="dfine_n_coco", 
            label="Model Selection"
        ),
        gr.Slider(
            minimum=0.1, 
            maximum=0.9, 
            value=0.4, 
            step=0.05, 
            label="Confidence Threshold"
        )
    ],
    outputs=[
        gr.Image(type="pil", label="Detection Result"),
        gr.Textbox(label="Detected Objects")
    ],
    title="D-FINE Object Detection Demo",
    description="Upload an image to see object detection results using the D-FINE model. You can select different models and adjust the confidence threshold.",
    examples=[
        ["examples/image1.jpg", "dfine_n_coco", 0.4],
    ]
)

demo.launch(debug=True)