Spaces:
Running
Running
Upload 3 files
Browse files- app.py +299 -0
- requirements.txt +6 -0
- 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
|