Spaces:
Runtime error
Runtime error
Upload 3 files
Browse files- app.py +83 -0
- best.pt +3 -0
- requirements.txt +4 -0
app.py
ADDED
@@ -0,0 +1,83 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import streamlit as st
|
2 |
+
from PIL import Image
|
3 |
+
import os
|
4 |
+
import shutil
|
5 |
+
from ultralytics import YOLO
|
6 |
+
|
7 |
+
# 1. Welcome message and brief description
|
8 |
+
st.title("Car Brand Detector")
|
9 |
+
st.subheader("Upload the image of a car (or multiple cars) and we will detect the brand!")
|
10 |
+
|
11 |
+
# 2. Initialize YOLO with our custom trained model name
|
12 |
+
model = YOLO("best.pt")
|
13 |
+
|
14 |
+
# 3. Add a file uploader widget to the sidebar
|
15 |
+
uploaded_file = st.sidebar.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
|
16 |
+
uploaded_image_path = ''
|
17 |
+
|
18 |
+
# Check if a file was uploaded
|
19 |
+
if uploaded_file is not None:
|
20 |
+
|
21 |
+
col1, col2 = st.columns((1, 2))
|
22 |
+
|
23 |
+
with col1:
|
24 |
+
# Display the uploaded image
|
25 |
+
image = Image.open(uploaded_file)
|
26 |
+
|
27 |
+
# Calculate the new width based on the aspect ratio
|
28 |
+
aspect_ratio = image.width / image.height
|
29 |
+
new_width = int(400 * aspect_ratio)
|
30 |
+
|
31 |
+
# Resize the image while maintaining aspect ratio
|
32 |
+
image_resized = image.resize((new_width, 400))
|
33 |
+
|
34 |
+
st.image(image_resized, caption='Uploaded Image')
|
35 |
+
|
36 |
+
# Save the uploaded file to a directory named 'uploaded_files'
|
37 |
+
save_dir = 'uploaded_files'
|
38 |
+
os.makedirs(save_dir, exist_ok=True) # Create the directory if it doesn't exist
|
39 |
+
file_path = os.path.join(save_dir, uploaded_file.name)
|
40 |
+
with open(file_path, 'wb') as f:
|
41 |
+
f.write(uploaded_file.getvalue())
|
42 |
+
|
43 |
+
uploaded_image_path = file_path
|
44 |
+
|
45 |
+
with col2:
|
46 |
+
# Create a button for brand detection
|
47 |
+
if st.button('Detect Brand'):
|
48 |
+
# Perform brand detection when the button is clicked
|
49 |
+
# Use the pre-trained YOLO model to detect the car brand from the uploaded image
|
50 |
+
results = model.predict(source=uploaded_image_path, save=True, save_txt=True)
|
51 |
+
|
52 |
+
# Display the uploaded image
|
53 |
+
image = Image.open('runs/detect/predict/'+uploaded_file.name)
|
54 |
+
|
55 |
+
# Calculate the new width based on the aspect ratio
|
56 |
+
aspect_ratio = image.width / image.height
|
57 |
+
new_width = int(400 * aspect_ratio)
|
58 |
+
|
59 |
+
# Resize the image while maintaining aspect ratio
|
60 |
+
image_resized = image.resize((new_width, 400))
|
61 |
+
|
62 |
+
classes = []
|
63 |
+
for detection in results:
|
64 |
+
# Assuming detection.boxes.cls and detection.boxes.conf are tensors with multiple elements
|
65 |
+
for class_id_tensor, confidence_tensor in zip(detection.boxes.cls, detection.boxes.conf):
|
66 |
+
class_id = int(class_id_tensor.item())
|
67 |
+
confidence = round(float(confidence_tensor.item()), 2)
|
68 |
+
class_name = results[0].names[class_id].capitalize()
|
69 |
+
classes.append((class_name, confidence))
|
70 |
+
|
71 |
+
caption = 'Brand(s) detected: '+str(classes)
|
72 |
+
st.image(image_resized, caption=caption)
|
73 |
+
|
74 |
+
for c in classes:
|
75 |
+
st.write("Brand:", c[0],"--- Confidence:", c[1])
|
76 |
+
|
77 |
+
# Delete the 'uploaded_files' directory after brand detection
|
78 |
+
if os.path.exists('uploaded_files'):
|
79 |
+
shutil.rmtree('uploaded_files')
|
80 |
+
|
81 |
+
# Delete the 'runs' directory after brand detection
|
82 |
+
if os.path.exists('runs'):
|
83 |
+
shutil.rmtree('runs')
|
best.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:d2eef332bcfda1146dfeba4860ca25fd060601a7d4613f9eafd6590390d7d917
|
3 |
+
size 22553251
|
requirements.txt
ADDED
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
1 |
+
Pillow
|
2 |
+
streamlit
|
3 |
+
opencv-python
|
4 |
+
ultralytics
|