147 lines
4.8 KiB
Python
147 lines
4.8 KiB
Python
import cv2
|
|
import numpy as np
|
|
import pymysql
|
|
from datetime import datetime
|
|
|
|
# DB Connection
|
|
conn = pymysql.connect(
|
|
host='localhost',
|
|
user='root',
|
|
password='',
|
|
database='db_traffic2'
|
|
)
|
|
cursor = conn.cursor()
|
|
|
|
# RTSP or video
|
|
cap = cv2.VideoCapture('rtsp://20.20.20.33:8554/')
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
print("Failed to read from source.")
|
|
exit()
|
|
|
|
height2, width2, channels = frame.shape
|
|
|
|
fgbg = cv2.createBackgroundSubtractorMOG2()
|
|
|
|
vehicle_id_counter = 0
|
|
tracker_dict = {}
|
|
|
|
# Speed lines
|
|
speed_line_y1 = 275
|
|
speed_line_y2 = 370
|
|
real_world_distance_m = 0.25
|
|
min_speed_kmh = 10 # Threshold
|
|
|
|
# Track last counted vehicle IDs to avoid spamming the total_vehicle table
|
|
last_vehicle_ids_in_zone = set()
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
|
fgmask = fgbg.apply(gray)
|
|
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
|
|
closing = cv2.morphologyEx(fgmask, cv2.MORPH_CLOSE, kernel)
|
|
opening = cv2.morphologyEx(closing, cv2.MORPH_OPEN, kernel)
|
|
dilation = cv2.dilate(opening, kernel)
|
|
_, bins = cv2.threshold(dilation, 220, 255, cv2.THRESH_BINARY)
|
|
|
|
contours, _ = cv2.findContours(bins, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[-2:]
|
|
cv2.drawContours(frame, [cv2.convexHull(c) for c in contours], -1, (0, 255, 0), 2)
|
|
|
|
min_area = 400
|
|
max_area = 40000
|
|
|
|
vehicle_ids_in_zone = set()
|
|
|
|
for cnt in contours:
|
|
area = cv2.contourArea(cnt)
|
|
if not (min_area < area < max_area):
|
|
continue
|
|
|
|
M = cv2.moments(cnt)
|
|
if M['m00'] == 0:
|
|
continue
|
|
cx = int(M['m10'] / M['m00'])
|
|
cy = int(M['m01'] / M['m00'])
|
|
|
|
x, y, w, h = cv2.boundingRect(cnt)
|
|
cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
|
|
cv2.drawMarker(frame, (cx, cy), (0, 0, 255), cv2.MARKER_CROSS, 10, 1)
|
|
|
|
matched_id = None
|
|
for vid, info in tracker_dict.items():
|
|
old_cx, old_cy = info['pos']
|
|
if abs(cx - old_cx) < 30 and abs(cy - old_cy) < 30:
|
|
matched_id = vid
|
|
break
|
|
|
|
if matched_id is None:
|
|
matched_id = vehicle_id_counter
|
|
vehicle_id_counter += 1
|
|
tracker_dict[matched_id] = {'pos': (cx, cy), 't1': None, 't2': None, 'logged': False}
|
|
else:
|
|
tracker_dict[matched_id]['pos'] = (cx, cy)
|
|
|
|
vehicle = tracker_dict[matched_id]
|
|
if not vehicle['t1'] and cy >= speed_line_y1:
|
|
vehicle['t1'] = datetime.now()
|
|
elif not vehicle['t2'] and cy >= speed_line_y2:
|
|
vehicle['t2'] = datetime.now()
|
|
|
|
# Speed check and log
|
|
if vehicle['t1'] and vehicle['t2'] and not vehicle['logged']:
|
|
delta_time = (vehicle['t2'] - vehicle['t1']).total_seconds()
|
|
if delta_time > 0:
|
|
speed_kmh = (real_world_distance_m / delta_time) * 3.6
|
|
if speed_kmh >= min_speed_kmh:
|
|
print(f"ID {matched_id} SPEED: {speed_kmh:.2f} km/h")
|
|
cursor.execute(
|
|
"INSERT INTO vehicle_log (vehicle_id, direction, timestamp, speed) VALUES (%s, %s, %s, %s)",
|
|
(matched_id, 'down', datetime.now(), round(speed_kmh, 2))
|
|
)
|
|
conn.commit()
|
|
else:
|
|
print(f"ID {matched_id} IGNORED (too slow): {speed_kmh:.2f} km/h")
|
|
vehicle['logged'] = True
|
|
|
|
# Labeling
|
|
cv2.putText(frame, f"ID:{matched_id}", (x, y - 5),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 255), 1)
|
|
|
|
if vehicle['t1'] and vehicle['t2']:
|
|
delta_time = (vehicle['t2'] - vehicle['t1']).total_seconds()
|
|
if delta_time > 0:
|
|
speed_kmh = (real_world_distance_m / delta_time) * 3.6
|
|
cv2.putText(frame, f"{speed_kmh:.1f} km/h", (x, y + h + 15),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 2)
|
|
|
|
# Check if vehicle is between the lines
|
|
if speed_line_y1 <= cy <= speed_line_y2:
|
|
vehicle_ids_in_zone.add(matched_id)
|
|
|
|
# Save count only if it changes
|
|
if vehicle_ids_in_zone != last_vehicle_ids_in_zone:
|
|
last_vehicle_ids_in_zone = vehicle_ids_in_zone.copy()
|
|
count = len(vehicle_ids_in_zone)
|
|
cursor.execute("INSERT INTO total_vehicle (count, timestamp) VALUES (%s, %s)", (count, datetime.now()))
|
|
conn.commit()
|
|
print(f"🚗 Vehicles in zone: {count}")
|
|
|
|
# Draw detection lines
|
|
cv2.line(frame, (0, speed_line_y1), (width2, speed_line_y1), (0, 255, 255), 2)
|
|
cv2.line(frame, (0, speed_line_y2), (width2, speed_line_y2), (255, 0, 255), 2)
|
|
|
|
cv2.imshow("Vehicle Detection", frame)
|
|
if cv2.waitKey(1) & 0xFF == 27:
|
|
break
|
|
|
|
# Cleanup
|
|
cap.release()
|
|
cursor.close()
|
|
conn.close()
|
|
cv2.destroyAllWindows()
|