# Step 1: Import libraries
from transformers import DetrImageProcessor, DetrForObjectDetection
from PIL import Image, ImageDraw
import torch
import gradio as gr

# Step 2: Load model and processor
model = DetrForObjectDetection.from_pretrained("hilmantm/detr-traffic-accident-detection")
processor = DetrImageProcessor.from_pretrained("hilmantm/detr-traffic-accident-detection")

# Step 3: Define the inference logic
def detect_traffic_accidents(image):
    # Preprocess image
    inputs = processor(images=image, return_tensors="pt")
    
    # Run model inference
    outputs = model(**inputs)

    # Extract detections
    target_sizes = torch.tensor([image.size[::-1]])  # Width, height
    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
    
    # Annotate image with bounding boxes
    draw = ImageDraw.Draw(image)
    for box, score, label in zip(results["boxes"], results["scores"], results["labels"]):
        if score > 0.9:  # Confidence threshold
            x_min, y_min, x_max, y_max = box
            draw.rectangle([x_min, y_min, x_max, y_max], outline="red", width=3)
            draw.text((x_min, y_min), f"{model.config.id2label[label]}: {score:.2f}", fill="red")

    return image

# Step 4: Define the Gradio Interface
def api_endpoint(input_image):
    output_image = detect_traffic_accidents(input_image)
    return output_image

# Step 5: Launch Gradio app
iface = gr.Interface(
    fn=api_endpoint,
    inputs=gr.Image(type="pil"),  # Accepts images
    outputs=gr.Image(type="pil"),  # Returns annotated image
)

iface.launch(server_name="0.0.0.0", server_port=7860)  # Accessible via LAN/WAN