Spaces:
Running
Running
| import streamlit as st | |
| from PIL import Image | |
| import os | |
| import shutil | |
| from ultralytics import YOLO | |
| # 1. Welcome message and brief description | |
| st.title("Car Brand Detector") | |
| st.subheader("Upload the image of a car (or multiple cars) and we will detect the brand!") | |
| # 2. Initialize YOLO with our custom trained model name | |
| model = YOLO("best.pt") | |
| # 3. Add a file uploader widget to the sidebar | |
| uploaded_file = st.sidebar.file_uploader("Upload an image", type=["jpg", "jpeg", "png"]) | |
| uploaded_image_path = '' | |
| # Check if a file was uploaded | |
| if uploaded_file is not None: | |
| col1, col2 = st.columns((1, 2)) | |
| with col1: | |
| # Display the uploaded image | |
| image = Image.open(uploaded_file) | |
| # Calculate the new width based on the aspect ratio | |
| aspect_ratio = image.width / image.height | |
| new_width = int(400 * aspect_ratio) | |
| # Resize the image while maintaining aspect ratio | |
| image_resized = image.resize((new_width, 400)) | |
| st.image(image_resized, caption='Uploaded Image') | |
| # Save the uploaded file to a directory named 'uploaded_files' | |
| save_dir = 'uploaded_files' | |
| os.makedirs(save_dir, exist_ok=True) # Create the directory if it doesn't exist | |
| file_path = os.path.join(save_dir, uploaded_file.name) | |
| with open(file_path, 'wb') as f: | |
| f.write(uploaded_file.getvalue()) | |
| uploaded_image_path = file_path | |
| with col2: | |
| # Create a button for brand detection | |
| if st.button('Detect Brand'): | |
| # Perform brand detection when the button is clicked | |
| # Use the pre-trained YOLO model to detect the car brand from the uploaded image | |
| results = model.predict(source=uploaded_image_path, save=True, save_txt=True) | |
| # Display the uploaded image | |
| image = Image.open('runs/detect/predict/'+uploaded_file.name) | |
| # Calculate the new width based on the aspect ratio | |
| aspect_ratio = image.width / image.height | |
| new_width = int(400 * aspect_ratio) | |
| # Resize the image while maintaining aspect ratio | |
| image_resized = image.resize((new_width, 400)) | |
| classes = [] | |
| for detection in results: | |
| # Assuming detection.boxes.cls and detection.boxes.conf are tensors with multiple elements | |
| for class_id_tensor, confidence_tensor in zip(detection.boxes.cls, detection.boxes.conf): | |
| class_id = int(class_id_tensor.item()) | |
| confidence = round(float(confidence_tensor.item()), 2) | |
| class_name = results[0].names[class_id].capitalize() | |
| classes.append((class_name, confidence)) | |
| caption = 'Brand(s) detected: '+str(classes) | |
| st.image(image_resized, caption=caption) | |
| for c in classes: | |
| st.write("Brand:", c[0],"--- Confidence:", c[1]) | |
| # Delete the 'uploaded_files' directory after brand detection | |
| if os.path.exists('uploaded_files'): | |
| shutil.rmtree('uploaded_files') | |
| # Delete the 'runs' directory after brand detection | |
| if os.path.exists('runs'): | |
| shutil.rmtree('runs') |