Sanshruth commited on
Commit
3b03261
·
verified ·
1 Parent(s): 6bce17c

Upload 3 files

Browse files
Files changed (3) hide show
  1. app.py +299 -0
  2. requirements.txt +6 -0
  3. yolo11n.pt +3 -0
app.py ADDED
@@ -0,0 +1,299 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import cv2
3
+ import gradio as gr
4
+ import numpy as np
5
+ from PIL import Image, ImageDraw
6
+ from ultralytics import YOLO
7
+ from ultralytics.utils.plotting import Annotator, colors
8
+ import logging
9
+ import math
10
+
11
+ # Set up logging
12
+ logging.basicConfig(level=logging.INFO)
13
+ logger = logging.getLogger(__name__)
14
+
15
+ # Global variables to store line coordinates and line equation
16
+ start_point = None
17
+ end_point = None
18
+ line_params = None # Stores (slope, intercept) of the line
19
+
20
+ def extract_first_frame(stream_url):
21
+ """
22
+ Extracts the first available frame from the IP camera stream and returns it as a PIL image.
23
+ """
24
+ logger.info("Attempting to extract the first frame from the stream...")
25
+ cap = cv2.VideoCapture(stream_url)
26
+ if not cap.isOpened():
27
+ logger.error("Error: Could not open stream.")
28
+ return None, "Error: Could not open stream."
29
+
30
+ ret, frame = cap.read()
31
+ cap.release()
32
+
33
+ if not ret:
34
+ logger.error("Error: Could not read the first frame.")
35
+ return None, "Error: Could not read the first frame."
36
+
37
+ # Convert the frame to a PIL image
38
+ frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
39
+ pil_image = Image.fromarray(frame_rgb)
40
+
41
+ logger.info("First frame extracted successfully.")
42
+ return pil_image, "First frame extracted successfully."
43
+
44
+ def update_line(image, evt: gr.SelectData):
45
+ """
46
+ Updates the line based on user interaction (click and drag).
47
+ """
48
+ global start_point, end_point, line_params
49
+
50
+ # If it's the first click, set the start point and show it on the image
51
+ if start_point is None:
52
+ start_point = (evt.index[0], evt.index[1])
53
+
54
+ # Draw the start point on the image
55
+ draw = ImageDraw.Draw(image)
56
+ draw.ellipse(
57
+ (start_point[0] - 5, start_point[1] - 5, start_point[0] + 5, start_point[1] + 5),
58
+ fill="blue", outline="blue"
59
+ )
60
+
61
+ return image, f"Line Coordinates:\nStart: {start_point}, End: None"
62
+
63
+ # If it's the second click, set the end point and draw the line
64
+ end_point = (evt.index[0], evt.index[1])
65
+
66
+ # Calculate the slope (m) and intercept (b) of the line: y = mx + b
67
+ if start_point[0] != end_point[0]: # Avoid division by zero
68
+ slope = (end_point[1] - start_point[1]) / (end_point[0] - start_point[0])
69
+ intercept = start_point[1] - slope * start_point[0]
70
+ line_params = (slope, intercept, start_point, end_point) # Store slope, intercept, and points
71
+ else:
72
+ # Vertical line (special case)
73
+ line_params = (float('inf'), start_point[0], start_point, end_point)
74
+
75
+ # Draw the line and end point on the image
76
+ draw = ImageDraw.Draw(image)
77
+ draw.line([start_point, end_point], fill="red", width=2)
78
+ draw.ellipse(
79
+ (end_point[0] - 5, end_point[1] - 5, end_point[0] + 5, end_point[1] + 5),
80
+ fill="green", outline="green"
81
+ )
82
+
83
+ # Return the updated image and line info
84
+ line_info = f"Line Coordinates:\nStart: {start_point}, End: {end_point}\nLine Equation: y = {line_params[0]:.2f}x + {line_params[1]:.2f}"
85
+
86
+ # Reset the points for the next interaction
87
+ start_point = None
88
+ end_point = None
89
+
90
+ return image, line_info
91
+
92
+ def reset_line():
93
+ """
94
+ Resets the line coordinates.
95
+ """
96
+ global start_point, end_point, line_params
97
+ start_point = None
98
+ end_point = None
99
+ line_params = None
100
+ return None, "Line reset. Click to draw a new line."
101
+
102
+ def intersect(A, B, C, D):
103
+ """
104
+ Determines if two line segments AB and CD intersect.
105
+ """
106
+ def ccw(A, B, C):
107
+ return (C[1] - A[1]) * (B[0] - A[0]) - (B[1] - A[1]) * (C[0] - A[0])
108
+
109
+ def on_segment(A, B, C):
110
+ if min(A[0], B[0]) <= C[0] <= max(A[0], B[0]) and min(A[1], B[1]) <= C[1] <= max(A[1], B[1]):
111
+ return True
112
+ return False
113
+
114
+ # Check if the line segments intersect
115
+ ccw1 = ccw(A, B, C)
116
+ ccw2 = ccw(A, B, D)
117
+ ccw3 = ccw(C, D, A)
118
+ ccw4 = ccw(C, D, B)
119
+
120
+ if ((ccw1 * ccw2 < 0) and (ccw3 * ccw4 < 0)):
121
+ return True
122
+ elif ccw1 == 0 and on_segment(A, B, C):
123
+ return True
124
+ elif ccw2 == 0 and on_segment(A, B, D):
125
+ return True
126
+ elif ccw3 == 0 and on_segment(C, D, A):
127
+ return True
128
+ elif ccw4 == 0 and on_segment(C, D, B):
129
+ return True
130
+ else:
131
+ return False
132
+
133
+ def is_object_crossing_line(box, line_params):
134
+ """
135
+ Determines if an object's bounding box is fully intersected by the user-drawn line.
136
+ """
137
+ _, _, line_start, line_end = line_params
138
+
139
+ # Get the bounding box coordinates
140
+ x1, y1, x2, y2 = box
141
+
142
+ # Define the four edges of the bounding box
143
+ box_edges = [
144
+ ((x1, y1), (x2, y1)), # Top edge
145
+ ((x2, y1), (x2, y2)), # Right edge
146
+ ((x2, y2), (x1, y2)), # Bottom edge
147
+ ((x1, y2), (x1, y1)) # Left edge
148
+ ]
149
+
150
+ # Count the number of intersections between the line and the bounding box edges
151
+ intersection_count = 0
152
+ for edge_start, edge_end in box_edges:
153
+ if intersect(line_start, line_end, edge_start, edge_end):
154
+ intersection_count += 1
155
+
156
+ # Only count the object if the line intersects the bounding box at least twice
157
+ return intersection_count >= 2
158
+
159
+ def draw_angled_line(image, line_params, color=(0, 255, 0), thickness=2):
160
+ """
161
+ Draws the user-defined line on the frame.
162
+ """
163
+ _, _, start_point, end_point = line_params
164
+ cv2.line(image, start_point, end_point, color, thickness)
165
+
166
+ def process_video(confidence_threshold=0.5, selected_classes=None, stream_url=None):
167
+ """
168
+ Processes the IP camera stream to count objects of the selected classes crossing the line.
169
+ """
170
+ global line_params
171
+
172
+ errors = []
173
+
174
+ if line_params is None:
175
+ errors.append("Error: No line drawn. Please draw a line on the first frame.")
176
+ if selected_classes is None or len(selected_classes) == 0:
177
+ errors.append("Error: No classes selected. Please select at least one class to detect.")
178
+ if stream_url is None or stream_url.strip() == "":
179
+ errors.append("Error: No stream URL provided.")
180
+
181
+ if errors:
182
+ return None, "\n".join(errors)
183
+
184
+ logger.info("Connecting to the IP camera stream...")
185
+ cap = cv2.VideoCapture(stream_url)
186
+ if not cap.isOpened():
187
+ errors.append("Error: Could not open stream.")
188
+ return None, "\n".join(errors)
189
+
190
+ model = YOLO(model="yolo11n.pt")
191
+ crossed_objects = {}
192
+ max_tracked_objects = 1000 # Maximum number of objects to track before clearing
193
+
194
+ logger.info("Starting to process the stream...")
195
+ while cap.isOpened():
196
+ ret, frame = cap.read()
197
+ if not ret:
198
+ errors.append("Error: Could not read frame from the stream.")
199
+ break
200
+
201
+ # Perform object tracking with confidence threshold
202
+ results = model.track(frame, persist=True, conf=confidence_threshold)
203
+
204
+ if results[0].boxes.id is not None:
205
+ track_ids = results[0].boxes.id.int().cpu().tolist()
206
+ clss = results[0].boxes.cls.cpu().tolist()
207
+ boxes = results[0].boxes.xyxy.cpu()
208
+ confs = results[0].boxes.conf.cpu().tolist()
209
+
210
+ for box, cls, t_id, conf in zip(boxes, clss, track_ids, confs):
211
+ if conf >= confidence_threshold and model.names[cls] in selected_classes:
212
+ # Check if the object crosses the line
213
+ if is_object_crossing_line(box, line_params) and t_id not in crossed_objects:
214
+ crossed_objects[t_id] = True
215
+
216
+ # Clear the dictionary if it gets too large
217
+ if len(crossed_objects) > max_tracked_objects:
218
+ crossed_objects.clear()
219
+
220
+ # Visualize the results with bounding boxes, masks, and IDs
221
+ annotated_frame = results[0].plot()
222
+
223
+ # Draw the angled line on the frame
224
+ draw_angled_line(annotated_frame, line_params, color=(0, 255, 0), thickness=2)
225
+
226
+ # Display the count on the frame with a modern look
227
+ count = len(crossed_objects)
228
+ (text_width, text_height), _ = cv2.getTextSize(f"COUNT: {count}", cv2.FONT_HERSHEY_SIMPLEX, 1, 2)
229
+
230
+ # Calculate the position for the middle of the top
231
+ margin = 10 # Margin from the top
232
+ x = (annotated_frame.shape[1] - text_width) // 2 # Center-align the text horizontally
233
+ y = text_height + margin # Top-align the text
234
+
235
+ # Draw the black background rectangle
236
+ cv2.rectangle(annotated_frame, (x - margin, y - text_height - margin), (x + text_width + margin, y + margin), (0, 0, 0), -1)
237
+
238
+ # Draw the text
239
+ cv2.putText(annotated_frame, f"COUNT: {count}", (x, y), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
240
+
241
+ # Yield the annotated frame to Gradio
242
+ yield annotated_frame, ""
243
+
244
+ cap.release()
245
+ logger.info("Stream processing completed.")
246
+
247
+ # Define the Gradio interface
248
+ with gr.Blocks() as demo:
249
+ gr.Markdown("<center><h1><u>Fast Real-time Object Detection & Tracking with High-Res Output</u></h1></center>")
250
+ gr.Markdown("<center><h2> <u>Detect and count objects crossing a line with Yolo11n </u></h2></center>")
251
+
252
+
253
+
254
+ # Step 1: Enter the IP Camera Stream URL
255
+ # gr.Markdown("### Step 0: Enter the IP Camera Stream URL")
256
+ stream_url = gr.Textbox(label="Enter IP Camera Stream URL", value="https://s86.ipcamlive.com/streams/56bajygtsxwuzdmte/stream.m3u8", visible=False)
257
+
258
+ # Step 1: Extract the first frame from the stream
259
+ gr.Markdown("### Step 1: Click on the frame to draw a line, the objects crossing it would be counted in real-time.")
260
+ first_frame, status = extract_first_frame(stream_url.value)
261
+ if first_frame is None:
262
+ gr.Markdown(f"**Error:** {status}")
263
+ else:
264
+ # Image component for displaying the first frame
265
+ image = gr.Image(value=first_frame, label="First Frame of Stream", type="pil")
266
+
267
+
268
+ line_info = gr.Textbox(label="Line Coordinates", value="Line Coordinates:\nStart: None, End: None")
269
+ image.select(update_line, inputs=image, outputs=[image, line_info])
270
+
271
+ # Reset the line (optional)
272
+ # gr.Markdown("### Step 4: Reset the Line (Optional)")
273
+ # reset_button = gr.Button("Reset Line")
274
+ # reset_button.click(reset_line, inputs=None, outputs=[image, line_info])
275
+
276
+ # Step 2: Select classes to detect
277
+ gr.Markdown("### Step 2: Select Classes to Detect")
278
+ model = YOLO(model="yolo11n.pt") # Load the model to get class names
279
+ class_names = list(model.names.values()) # Get class names
280
+ selected_classes = gr.CheckboxGroup(choices=class_names, label="Select Classes to Detect")
281
+
282
+ # Step 3: Adjust confidence threshold
283
+ gr.Markdown("### Step 3: Adjust Confidence Threshold (Optional)")
284
+ confidence_threshold = gr.Slider(minimum=0.0, maximum=1.0, value=0.2, label="Confidence Threshold")
285
+
286
+ #process the stream
287
+ process_button = gr.Button("Process Stream")
288
+
289
+ # Output image for real-time frame rendering
290
+ output_image = gr.Image(label="Processed Frame", streaming=True)
291
+
292
+ # Error box to display warnings/errors
293
+ error_box = gr.Textbox(label="Errors/Warnings", interactive=False)
294
+
295
+ # Event listener for processing the video
296
+ process_button.click(process_video, inputs=[confidence_threshold, selected_classes, stream_url], outputs=[output_image, error_box])
297
+
298
+ # Launch the interface
299
+ demo.launch(debug=True)
requirements.txt ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+
2
+ ultralytics
3
+ opencv-python
4
+ gradio
5
+ torch
6
+ lap>=0.5.12
yolo11n.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0ebbc80d4a7680d14987a577cd21342b65ecfd94632bd9a8da63ae6417644ee1
3
+ size 5613764