Spaces:
Sleeping
Sleeping
File size: 3,280 Bytes
a4dc79a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
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') |