aminaj commited on
Commit
a4dc79a
1 Parent(s): 9fe106a

Upload 3 files

Browse files
Files changed (3) hide show
  1. app.py +83 -0
  2. best.pt +3 -0
  3. 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