import torch
from pathlib import Path
from PIL import Image
import streamlit as st
from torchvision import transforms

# Load YOLOv5 model
model = torch.hub.load('ultralytics/yolov5:master', 'yolov5s', pretrained=True)

# Define image transformation
transform = transforms.Compose([
    transforms.ToTensor(),
])

def perform_inference(image):
    # Apply transformation
    input_image = transform(image).unsqueeze(0)

    # Perform inference
    with torch.no_grad():
        results = model(input_image)

    return results

def display_results(image, results):
    # Display the image with bounding boxes
    st.image(image, caption="Input Image", use_column_width=True)

    # Access the bounding box coordinates and class labels
    boxes = results.xyxy[0].cpu().numpy()[:, :-1]
    class_labels = results.xyxy[0].cpu().numpy()[:, -1]

    # Display bounding boxes on the image
    for box, label in zip(boxes, class_labels):
        st.rectangle(
            xy=(box[0], box[1]),
            width=box[2] - box[0],
            height=box[3] - box[1],
            color='red',
            label=f'Class {int(label)}'
        )

# Streamlit app
def main():
    st.title("YOLOv5 Object Detection with Streamlit")

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

    if uploaded_file is not None:
        # Load the image
        image = Image.open(uploaded_file).convert('RGB')

        # Perform inference
        results = perform_inference(image)

        # Display results
        display_results(image, results)

        # Save the result image
        results.save(Path('output'))

if __name__ == "__main__":
    main()