Spaces:
Sleeping
Sleeping
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)}") | |