from transformers import pipeline
from PIL import Image, ImageDraw
import gradio as gr
import torch
import timm  # Required for DETR model

# Load the model
device = 0 if torch.cuda.is_available() else -1
print(f"Using device: {'cuda' if device == 0 else 'cpu'}")

model_pipeline = pipeline(
    "object-detection",
    model="opria123/detr-resnet-50-dc5-hardhat-finetuned",
    device=device,
    threshold=0.2  # Lower threshold to catch more results
)

# Function to classify and draw results
def classify_image(image):
    print("Received image for classification.")

    results = model_pipeline(image)
    print(f"Model results:\n{results}")

    image = image.convert("RGB")
    draw = ImageDraw.Draw(image)

    for item in results:
        box = item["box"]
        label = item["label"]
        score = item["score"]

        print(f"Drawing box for: {label} ({score:.2f}) at {box}")

        draw.rectangle(
            [(box["xmin"], box["ymin"]), (box["xmax"], box["ymax"])],
            outline="red",
            width=3
        )
        draw.text((box["xmin"] + 5, box["ymin"] - 10), f"{label} ({score:.2f})", fill="red")

    return image

# Gradio interface
interface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Hard Hat Detection",
    description="Upload an image to detect hard hats using a fine-tuned DETR model."
)

interface.launch()