simran0608 commited on
Commit
edd6baf
·
verified ·
1 Parent(s): 83a4bb1

Upload 3 files

Browse files
Files changed (3) hide show
  1. best.pt +3 -0
  2. main.py +153 -0
  3. requirements.txt +4 -0
best.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b81444d786fee3c4325e135c4f6c63c6df64402f517458b3186ea032e7edacc7
3
+ size 87596350
main.py ADDED
@@ -0,0 +1,153 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from PIL import Image, ImageDraw
2
+ from ultralytics import YOLO
3
+ import streamlit as st
4
+ import tempfile
5
+ import cv2
6
+ import numpy as np
7
+ import base64
8
+
9
+ # Initialize YOLO model
10
+ model = YOLO("best.pt")
11
+
12
+ # Function to perform object detection on an image
13
+ def detect_objects_image(image):
14
+ results = model(image)
15
+ result = results[0]
16
+ output = []
17
+ num_potholes_detected = 0
18
+ num_cracks_detected = 0
19
+ num_alligator_cracks_detected = 0
20
+
21
+ for box in result.boxes:
22
+ x1, y1, x2, y2 = [round(x) for x in box.xyxy[0].tolist()]
23
+ class_id = box.cls[0].item()
24
+ prob = round(box.conf[0].item(), 2)
25
+ class_name = result.names[class_id]
26
+ output.append([x1, y1, x2, y2, class_name, prob])
27
+
28
+ # Count detections by class
29
+ if class_name == "pothole":
30
+ num_potholes_detected += 1
31
+ elif class_name == "crack":
32
+ num_cracks_detected += 1
33
+ elif class_name == "alligator-crack":
34
+ num_alligator_cracks_detected += 1
35
+
36
+ return output, num_potholes_detected, num_cracks_detected, num_alligator_cracks_detected
37
+
38
+ # Function to process and annotate a video
39
+ def process_video(video_path, output_path, frame_interval):
40
+ cap = cv2.VideoCapture(video_path)
41
+ fps = int(cap.get(cv2.CAP_PROP_FPS))
42
+ frame_interval_count = int(fps * frame_interval)
43
+ width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
44
+ height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
45
+ fourcc = cv2.VideoWriter_fourcc(*'mp4v')
46
+ out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
47
+ frame_count = 0
48
+ detections_summary = {
49
+ 'potholes': 0,
50
+ 'cracks': 0,
51
+ 'alligator_cracks': 0
52
+ }
53
+
54
+ while cap.isOpened():
55
+ ret, frame = cap.read()
56
+ if not ret:
57
+ break
58
+
59
+ if frame_count % frame_interval_count == 0:
60
+ image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
61
+ detections, num_potholes, num_cracks, num_alligator_cracks = detect_objects_image(image)
62
+ detections_summary['potholes'] += num_potholes
63
+ detections_summary['cracks'] += num_cracks
64
+ detections_summary['alligator_cracks'] += num_alligator_cracks
65
+
66
+ draw = ImageDraw.Draw(image)
67
+ for detection in detections:
68
+ x1, y1, x2, y2, class_name, prob = detection
69
+ draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
70
+ text = f"{class_name} {prob:.2f}"
71
+ draw.text((x1, y1), text, fill="red")
72
+
73
+ annotated_frame = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
74
+ else:
75
+ annotated_frame = frame
76
+
77
+ out.write(annotated_frame)
78
+ frame_count += 1
79
+
80
+ cap.release()
81
+ out.release()
82
+
83
+ return detections_summary
84
+
85
+ # Function to generate a download link for a file
86
+ def get_download_link(file_path, text, file_type):
87
+ with open(file_path, 'rb') as f:
88
+ file_bytes = f.read()
89
+ file_b64 = base64.b64encode(file_bytes).decode()
90
+ download_link = f'<a href="data:{file_type};base64,{file_b64}" download="{text}">{text}</a>'
91
+ return download_link
92
+
93
+ # Streamlit app
94
+ def main():
95
+ st.title("Road Condition Inspection")
96
+ st.subheader("Upload an image or video to detect objects")
97
+
98
+ # File uploader for image and video
99
+ uploaded_file = st.file_uploader("Choose a file...", type=["jpg", "jpeg", "png", "mp4"])
100
+
101
+ if uploaded_file is not None:
102
+ file_type = uploaded_file.type
103
+ temp_file = tempfile.NamedTemporaryFile(delete=False, suffix="." + uploaded_file.name.split('.')[-1])
104
+ temp_file.write(uploaded_file.read())
105
+ temp_file.close()
106
+
107
+ if file_type.startswith("image"):
108
+ image = Image.open(temp_file.name)
109
+ st.image(image, caption='Uploaded Image', use_column_width=True)
110
+
111
+ if st.button('Detect Objects (Image)'):
112
+ detections, num_potholes, num_cracks, num_alligator_cracks = detect_objects_image(image)
113
+ draw = ImageDraw.Draw(image)
114
+ for detection in detections:
115
+ x1, y1, x2, y2, class_name, prob = detection
116
+ draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
117
+ text = f"{class_name} {prob:.2f}"
118
+ draw.text((x1, y1), text, fill="red")
119
+
120
+ st.image(image, caption='Annotated Image', use_column_width=True)
121
+ st.subheader("Detection Summary")
122
+ if num_potholes > 0:
123
+ st.write(f"Potholes Detected: {num_potholes}")
124
+ if num_cracks > 0:
125
+ st.write(f"Cracks Detected: {num_cracks}")
126
+ if num_alligator_cracks > 0:
127
+ st.write(f"Alligator Cracks Detected: {num_alligator_cracks}")
128
+
129
+ annotated_image_path = temp_file.name.replace(".", "_annotated.")
130
+ image.save(annotated_image_path)
131
+ st.markdown(get_download_link(annotated_image_path, "Download Annotated Image", "image/png"), unsafe_allow_html=True)
132
+
133
+ elif file_type.startswith("video"):
134
+ video_bytes = open(temp_file.name, 'rb').read()
135
+ st.video(video_bytes)
136
+
137
+ if st.button('Detect Objects (Video)'):
138
+ annotated_video_path = temp_file.name.replace(".", "_annotated") + ".mp4"
139
+ detections_summary = process_video(temp_file.name, annotated_video_path, frame_interval=1)
140
+
141
+ st.subheader("Annotated Video Download")
142
+ st.markdown(get_download_link(annotated_video_path, "Download Annotated Video", "video/mp4"), unsafe_allow_html=True)
143
+
144
+ st.subheader("Detection Summary")
145
+ if detections_summary['potholes'] > 0:
146
+ st.write(f"Total Potholes Detected: {detections_summary['potholes']}")
147
+ if detections_summary['cracks'] > 0:
148
+ st.write(f"Total Cracks Detected: {detections_summary['cracks']}")
149
+ if detections_summary['alligator_cracks'] > 0:
150
+ st.write(f"Total Alligator Cracks Detected: {detections_summary['alligator_cracks']}")
151
+
152
+ if __name__ == '__main__':
153
+ main()
requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ ultralytics
2
+ gtts
3
+ streamlit
4
+ pillow