import streamlit as st
from PIL import Image
import torch
import torchvision.transforms as T
import torchvision
import matplotlib.pyplot as plt
import matplotlib.patches as patches

# Load pre-trained object detection model
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()

# COCO class labels
COCO_INSTANCE_CATEGORY_NAMES = [
    '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign',
    'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
    'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag',
    'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite',
    'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
    'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana',
    'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
    'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table',
    'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
    'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
    'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]

# Streamlit UI
st.set_page_config(page_title="AI Object Detector", layout="centered")
st.title("🎯 Object Detection with AI")
st.markdown("Upload an image and let the AI detect objects with names!")

uploaded_file = st.file_uploader("📷 Upload an image", type=["jpg", "png", "jpeg"])

if uploaded_file:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Transform image for model input
    transform = T.Compose([T.ToTensor()])
    image_tensor = transform(image).unsqueeze(0)

    with st.spinner("🔍 Detecting objects..."):
        predictions = model(image_tensor)[0]

    # Draw bounding boxes
    fig, ax = plt.subplots(1)
    ax.imshow(image)

    threshold = 0.7  # confidence threshold
    for idx in range(len(predictions["boxes"])):
        score = predictions["scores"][idx].item()
        if score > threshold:
            box = predictions["boxes"][idx].detach().numpy()
            label = COCO_INSTANCE_CATEGORY_NAMES[predictions["labels"][idx]]
            x1, y1, x2, y2 = box
            rect = patches.Rectangle((x1, y1), x2 - x1, y2 - y1,
                                     linewidth=2, edgecolor='lime', facecolor='none')
            ax.add_patch(rect)
            ax.text(x1, y1 - 10, f"{label} ({score:.2f})", color='lime',
                    fontsize=10, backgroundcolor='black')

    st.pyplot(fig)