top_wear / app.py
vishalkatheriya18's picture
Update app.py
3500d9f verified
raw
history blame
1.99 kB
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)}")