import streamlit as st
from ultralytics import YOLO
from PIL import Image
import requests
from io import BytesIO
import time

# Load the YOLO model only once using Streamlit session state
if 'models_loaded' not in st.session_state:
    st.session_state.yolo_model = YOLO('best.pt')  # Update with your model path
    st.session_state.models_loaded = True

# Define function for inference using YOLO
def predict_with_yolo(image):
    # Run inference on the image using the YOLO model
    results = st.session_state.yolo_model(image)
    
    # Print the classification results
    if results:
        # YOLOv8 classification results contain label and confidence for the whole image
        for result in results:
            return {
                "Class": result.names[result.probs.top1],
                "Confidence": float(result.probs.top1conf)  # Convert confidence to float
            }
    else:
        st.warning("No classification results found.")
        return None
    

# Streamlit app UI
st.title("Clothing Detection with YOLO")

url = st.text_input("Paste image URL here...")
if url:
    try:
        response = requests.get(url)
        if response.status_code == 200:
            image = Image.open(BytesIO(response.content)).convert('RGB')
            st.image(image.resize((200, 200)), caption="Uploaded Image", use_column_width=False)
            
            start_time = time.time()
            image_resized = image.resize((224, 224))
            # Predict using YOLO
            predictions = predict_with_yolo(image_resized)

            # Display predictions
            if predictions:
                st.write("Predictions:", predictions)
            else:
                st.write("No objects detected.")

            st.write(f"Time taken: {time.time() - start_time:.2f} seconds")
        else:
            st.error("Failed to load image from URL. Please check the URL.")
    except Exception as e:
        st.error(f"Error processing the image: {str(e)}")