import streamlit as st
from transformers import AutoImageProcessor, DetaForObjectDetection
from PIL import Image
import requests

st.title("Object Detection")

# Sidebar instructions
st.sidebar.header("Instructions")
st.sidebar.write("1. Enter an image URL in the text input below.")
st.sidebar.write("2. Click the 'Detect Objects' button to process the image.")

# Image URL input
image_url = st.text_input("Enter image URL:", "http://images.cocodataset.org/val2017/000000039769.jpg")

if st.button("Detect Objects"):
    try:
        # Load the image
        image = Image.open(requests.get(image_url, stream=True).raw)

        # Initialize the image processor and model
        image_processor = AutoImageProcessor.from_pretrained("jozhang97/deta-swin-large")
        model = DetaForObjectDetection.from_pretrained("jozhang97/deta-swin-large")

        # Process the image
        inputs = image_processor(images=image, return_tensors="pt")
        outputs = model(**inputs)

        # Convert outputs to Pascal VOC format
        target_sizes = torch.tensor([image.size[::-1]])
        results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[0]

        # Display the image and detected objects
        st.image(image, use_column_width=True)
        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
            box = [round(i, 2) for i in box.tolist()]
            st.write(f"Detected {model.config.id2label[label.item()]} with confidence {round(score.item(), 3)} at location {box}")
    except:
        st.write("Error: Unable to process the image. Please check the URL and try again.")