Spaces:
Sleeping
Sleeping
File size: 7,036 Bytes
de631bc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 | """
Utility functions for vehicle tracking and counting system
"""
import cv2
import numpy as np
from typing import Tuple, Dict
def get_center_point(bbox: Tuple[int, int, int, int]) -> Tuple[int, int]:
"""
Calculate the center point of a bounding box
Args:
bbox: Bounding box coordinates (x1, y1, x2, y2)
Returns:
Tuple of (center_x, center_y)
"""
x1, y1, x2, y2 = bbox
center_x = int((x1 + x2) / 2)
center_y = int((y1 + y2) / 2)
return center_x, center_y
def check_line_crossing(
curr_pos: int,
prev_pos: int,
line_pos: int,
margin: int = 5
) -> bool:
"""
Check if an object has crossed the counting line
Args:
curr_pos: Current coordinate (X or Y)
prev_pos: Previous coordinate (X or Y)
line_pos: Coordinate of counting line
margin: Margin of error for line crossing
Returns:
True if object crossed the line (in either direction)
"""
# Check crossing in positive direction (e.g., left to right or top to bottom)
if prev_pos < line_pos - margin and curr_pos >= line_pos + margin:
return True
# Check crossing in negative direction (e.g., right to left or bottom to top)
if prev_pos > line_pos + margin and curr_pos <= line_pos - margin:
return True
return False
def draw_counting_line(
frame: np.ndarray,
line_coords: Tuple[int, int, int, int],
color: Tuple[int, int, int],
thickness: int
) -> np.ndarray:
"""
Draw the counting line on the frame
Args:
frame: Video frame
line_coords: Line coordinates (x1, y1, x2, y2)
color: Line color (B, G, R)
thickness: Line thickness
Returns:
Frame with line drawn
"""
x1, y1, x2, y2 = line_coords
cv2.line(frame, (x1, y1), (x2, y2), color, thickness)
# Add text label for the line
cv2.putText(
frame,
"COUNTING LINE",
(x1 + 10, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX,
0.7,
color,
2
)
return frame
def draw_statistics(
frame: np.ndarray,
counts: Dict[str, int],
position: Tuple[int, int],
font_scale: float = 0.8,
color: Tuple[int, int, int] = (255, 255, 255),
bg_color: Tuple[int, int, int] = (0, 0, 0)
) -> np.ndarray:
"""
Draw counting statistics on the frame
Args:
frame: Video frame
counts: Dictionary of vehicle counts by class
position: Position to draw statistics (x, y)
font_scale: Font scale
color: Text color
bg_color: Background color
Returns:
Frame with statistics drawn
"""
x, y = position
line_height = 30
# Draw background rectangle
total_lines = len(counts) + 1
cv2.rectangle(
frame,
(x - 5, y - 25),
(x + 250, y + line_height * total_lines + 5),
bg_color,
-1
)
# Draw total count
total = sum(counts.values())
cv2.putText(
frame,
f"TOTAL: {total}",
(x, y),
cv2.FONT_HERSHEY_SIMPLEX,
font_scale,
(0, 255, 255), # Yellow for total
2
)
# Draw individual class counts
y_offset = y + line_height
for vehicle_class, count in counts.items():
cv2.putText(
frame,
f"{vehicle_class.upper()}: {count}",
(x, y_offset),
cv2.FONT_HERSHEY_SIMPLEX,
font_scale,
color,
2
)
y_offset += line_height
return frame
def draw_bounding_box(
frame: np.ndarray,
bbox: Tuple[int, int, int, int],
track_id: int,
class_name: str,
confidence: float,
color: Tuple[int, int, int],
thickness: int = 2,
show_id: bool = True,
show_confidence: bool = True
) -> np.ndarray:
"""
Draw bounding box with label on the frame
Args:
frame: Video frame
bbox: Bounding box coordinates (x1, y1, x2, y2)
track_id: Tracking ID
class_name: Class name
confidence: Detection confidence
color: Box color
thickness: Box thickness
show_id: Whether to show track ID
show_confidence: Whether to show confidence score
Returns:
Frame with bounding box drawn
"""
x1, y1, x2, y2 = bbox
# Draw bounding box
cv2.rectangle(frame, (x1, y1), (x2, y2), color, thickness)
# Prepare label text
label_parts = [class_name]
if show_id:
label_parts.append(f"ID:{track_id}")
if show_confidence:
label_parts.append(f"{confidence:.2f}")
label = " | ".join(label_parts)
# Calculate label size
(label_width, label_height), baseline = cv2.getTextSize(
label,
cv2.FONT_HERSHEY_SIMPLEX,
0.6,
2
)
# Draw label background
cv2.rectangle(
frame,
(x1, y1 - label_height - baseline - 5),
(x1 + label_width + 5, y1),
color,
-1
)
# Draw label text
cv2.putText(
frame,
label,
(x1 + 2, y1 - baseline - 2),
cv2.FONT_HERSHEY_SIMPLEX,
0.6,
(0, 0, 0), # Black text
2
)
# Draw center point
center_x, center_y = get_center_point(bbox)
cv2.circle(frame, (center_x, center_y), 4, color, -1)
return frame
def format_time(seconds: float) -> str:
"""
Format elapsed time in seconds to HH:MM:SS format
Args:
seconds: Time in seconds
Returns:
Formatted time string
"""
hours = int(seconds // 3600)
minutes = int((seconds % 3600) // 60)
secs = int(seconds % 60)
return f"{hours:02d}:{minutes:02d}:{secs:02d}"
def get_counting_line_coords(
frame_width: int,
frame_height: int,
line_position: float = 0.5,
custom_coords: Tuple[int, int, int, int] = None,
orientation: str = "horizontal"
) -> Tuple[int, int, int, int]:
"""
Get counting line coordinates based on frame dimensions
Args:
frame_width: Width of video frame
frame_height: Height of video frame
line_position: Position as percentage (0.0 to 1.0)
custom_coords: Custom line coordinates (overrides line_position)
orientation: "horizontal" or "vertical"
Returns:
Line coordinates (x1, y1, x2, y2)
"""
if custom_coords is not None:
return custom_coords
if orientation == "vertical":
line_x = int(frame_width * line_position)
return (line_x, 0, line_x, frame_height)
else:
line_y = int(frame_height * line_position)
return (0, line_y, frame_width, line_y)
|