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('/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)}") | |