import gradio as gr
import torch
from transformers import DetrImageProcessor, DetrForObjectDetection
from PIL import Image, ImageDraw

# Load the pre-trained DETR model and processor
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")

def detect_objects(image: Image.Image) -> Image.Image:
    try:
        # Preprocess the image
        inputs = processor(images=image, return_tensors="pt")
        outputs = model(**inputs)

        # Convert outputs to bounding boxes and labels
        target_sizes = torch.tensor([image.size[::-1]])
        results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]

        # Draw bounding boxes on the image
        draw = ImageDraw.Draw(image)
        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
            box = [round(i, 2) for i in box.tolist()]
            label_text = f"{model.config.id2label[label.item()]}: {round(score.item(), 3)}"
            draw.rectangle(box, outline="red", width=3)
            draw.text((box[0], box[1]), label_text, fill="red")
        return image
    except Exception as e:
        print("Error during detection:", e)
        return image  # In a robust production system, consider returning a message or a default image

# Create a Gradio interface
iface = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=gr.Image(label="Detection Result"),
    title="Robust Object Detection with DETR",
    description="Upload an image to detect objects using a pre-trained DETR model from Hugging Face Hub."
)

if __name__ == "__main__":
    iface.launch()