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('/kaggle/working/classification_project/yolo_classification/weights/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)
    
    # Extract predictions
    predictions = []
    if results:
        for result in results:
            for box in result.boxes:
                class_name = result.names[box.label]
                confidence = box.conf.item()  # Convert tensor to a Python float
                predictions.append({
                    "Class": class_name,
                    "Confidence": confidence
                })
    return predictions

# 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()

            # Predict using YOLO
            predictions = predict_with_yolo(image)

            # Display predictions
            if predictions:
                st.write("Predictions:")
                for pred in predictions:
                    st.write(f"Class: {pred['Class']}, Confidence: {pred['Confidence']:.2f}")
            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)}")