car-counter.py

import numpy as np
import supervision as sv
from supervision.draw.utils import draw_text
from supervision.draw.color import Color
import torch
#import cv2
from ultralytics import YOLO
import os
os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] = "0.0"

model = YOLO("yolov8n.pt")
model = model.to("mps")
tracker = sv.ByteTrack(track_buffer=100)
box_annotator = sv.BoundingBoxAnnotator()
label_annotator = sv.LabelAnnotator()
trace_annotator = sv.TraceAnnotator()
#linezone_annotator = sv.LineZoneAnnotator()

# dict maping class_id to class_name
CLASS_NAMES_DICT = model.model.names
# class_ids of interest - person [0], bicycle [1], car, motorcycle, bus and truck
car_classes = [2, 3, 5, 7]

# # settings - eufy
# LINE_START_WEST = sv.Point(1266, 864)
# LINE_END_WEST = sv.Point(1167, 1167)
# 
# LINE_START_SOUTH = sv.Point(700, 790)
# LINE_END_SOUTH = sv.Point(707, 990)

# scaling_factor = 4/9

# def scale_point(point):
#     return sv.Point(int(point.x * scaling_factor), int(point.y * scaling_factor))

LINE_START_WEST = sv.Point(825, 518)
LINE_END_WEST = sv.Point(9, 538)

LINE_START_SOUTH = sv.Point(143, 398)
LINE_END_SOUTH = sv.Point(311, 818)

VIDEO = "monday-morning-trim.mov"

colors = sv.ColorPalette.default()
video_info = sv.VideoInfo.from_video_path(VIDEO)

# mon eve
# polygons = [
# np.array([
# [1152, 1001],[40, 1021],[24, 205],[1892, 409],[1892, 409],[1320, 1037],[1148, 997]
# ]),  #full map
# np.array([
# [392, 626],[392, 462],[572, 554],[516, 670],[388, 622]
# ]),  #westbound
# np.array([
# [26, 593],[294, 661],[298, 833],[30, 857],[22, 585]
# ]),  #southbound
# np.array([
# [815, 629],[815, 509],[987, 509],[999, 649],[815, 629]
# ]),  #turns_onto_westminster
# np.array([
# [827, 690],[1347, 698],[1511, 810],[1351, 954],[815, 750],[827, 690]
# ])   #turns onto dorchester
# ]

polygons = [
np.array([
[284, 920],[16, 916],[16, 520],[584, 652],[696, 652],[740, 560],[1348, 480],[1356, 684],[1200, 732],[1210, 860],[1452, 908],[1432, 992],[1116, 1044],[1092, 940],[1000, 892],[936, 828],[804, 824],[760, 836],[628, 724],[412, 672],[312, 684],[288, 916]
]),  #full map
np.array([
[260, 684],[260, 456],[312, 516],[308, 604],[372, 624],[372, 676],[312, 684],[256, 684]
]),  #westbound
np.array([
[34, 889],[18, 613],[158, 713],[170, 881],[34, 885]
]),  #southbound
np.array([
[681, 511],[689, 695],[901, 719],[893, 507],[677, 507]
]),  #turns_onto_westminster
np.array([
[705, 754],[1101, 754],[1101, 906],[709, 838],[701, 754]
])   #turns onto dorchester
]

# To store timestamps when line_counter increments
previous_in_count_west = 0
previous_in_count_south = 0
previous_in_count_south = 0
previous_in_count_lt = 0
previous_in_count_rt = 0
timestamps_west = []
timestamps_south = []
timestamps_lt = []
timestamps_rt = []
zone_sum_west = 0
unique_tracker_ids_west = set()
unique_tracker_ids_south = set()
unique_tracker_ids_lt = set()
unique_tracker_ids_rt = set()

# line_counter_west = sv.LineZone(start=LINE_START_WEST, end=LINE_END_WEST)
# line_counter_south = sv.LineZone(start=LINE_START_SOUTH, end=LINE_END_SOUTH)

zones = [
    sv.PolygonZone(
        polygon=polygon,
        frame_resolution_wh=video_info.resolution_wh
    )
    for polygon
    in polygons
]

zone_annotator_wb = sv.PolygonZoneAnnotator(zone=zones[1], color=sv.Color.white(), thickness=1, text_thickness=1, text_scale=0.5)
zone_annotator_sb = sv.PolygonZoneAnnotator(zone=zones[2], color=sv.Color.white(), thickness=1, text_thickness=1, text_scale=0.5)
zone_annotator_lt = sv.PolygonZoneAnnotator(zone=zones[3], color=sv.Color.white(), thickness=1, text_thickness=1, text_scale=0.5)
zone_annotator_rt = sv.PolygonZoneAnnotator(zone=zones[4], color=sv.Color.white(), thickness=1, text_thickness=1, text_scale=0.5)
#line_annotator = sv.LineZoneAnnotator(thickness=1, text_thickness=1, text_scale=1)

def callback(frame: np.ndarray, frame_num: int) -> np.ndarray:
    results = model(frame, device="mps")[0]
    detections = sv.Detections.from_ultralytics(results)
    detections = detections[detections.confidence > 0.6]
    detections = tracker.update_with_detections(detections)
    
    global colors, previous_in_count_west, previous_in_count_south, previous_in_count_lt, previous_in_count_rt, timestamps_west, timestamps_south, timestamps_lt, timestamps_rt, zone_sum_west, unique_tracker_ids_west, unique_tracker_ids_south, unique_tracker_ids_lt, unique_tracker_ids_rt

    mask = zones[0].trigger(detections=detections)
    detections = detections[mask]
    #annotated_frame = box_annotator.annotate(scene=frame, detections=detections, skip_label=True)
    #annotated_frame = zone_annotator.annotate(scene=annotated_frame)

    labels = [
        f"#{tracker_id} {results.names[class_id]}"
        for class_id, tracker_id
        in zip(detections.class_id, detections.tracker_id)
    ]

    annotated_frame = box_annotator.annotate(
       frame.copy(), detections=detections)
    annotated_frame = label_annotator.annotate(
        annotated_frame, detections=detections, labels=labels)

    #line_counter_west.trigger(detections=detections[np.isin(detections.class_id, car_classes)])
    zones[1].trigger(detections=detections[np.isin(detections.class_id, car_classes)])
    zones[2].trigger(detections=detections[np.isin(detections.class_id, car_classes)])
    zones[3].trigger(detections=detections[np.isin(detections.class_id, car_classes)])
    zones[4].trigger(detections=detections[np.isin(detections.class_id, car_classes)])
         
    westbound_mask = zones[1].trigger(detections=detections)
    westbound_detections = detections[westbound_mask]

    southbound_mask = zones[2].trigger(detections=detections)
    southbound_detections = detections[southbound_mask]

    left_turn_mask = zones[3].trigger(detections=detections)
    left_turn_detections = detections[left_turn_mask]

    right_turn_mask = zones[4].trigger(detections=detections)
    right_turn_detections = detections[right_turn_mask]
    
    for class_id, tracker_id in zip(westbound_detections.class_id, westbound_detections.tracker_id):
        if class_id in car_classes and tracker_id not in unique_tracker_ids_south: unique_tracker_ids_west.add(tracker_id)
    for class_id, tracker_id in zip(left_turn_detections.class_id, left_turn_detections.tracker_id):
        if class_id in car_classes and tracker_id in unique_tracker_ids_west and tracker_id not in unique_tracker_ids_rt: unique_tracker_ids_lt.add(tracker_id)
    for class_id, tracker_id in zip(right_turn_detections.class_id, right_turn_detections.tracker_id):
        if class_id in car_classes and tracker_id in unique_tracker_ids_south and tracker_id not in unique_tracker_ids_lt: unique_tracker_ids_rt.add(tracker_id)
    for class_id, tracker_id in zip(southbound_detections.class_id, southbound_detections.tracker_id):
        if class_id in car_classes and tracker_id not in (unique_tracker_ids_rt,unique_tracker_ids_west): unique_tracker_ids_south.add(tracker_id)

    zone_sum_west = len(unique_tracker_ids_west)
    zone_sum_south = len(unique_tracker_ids_south)
    zone_sum_lt = len(unique_tracker_ids_lt)
    zone_sum_rt = len(unique_tracker_ids_rt)
    
    # Check for westbound cars
    if len(unique_tracker_ids_west) > previous_in_count_west:
        print(f"Westbound car detected at frame {frame_num}")
        timestamps_west.append(frame_num / 30) # eufy 15fps
        previous_in_count_west = len(unique_tracker_ids_west)

    # Check for southbound cars
    if len(unique_tracker_ids_south) > previous_in_count_south:
        print(f"Southbound car detected at frame {frame_num}")
        timestamps_south.append(frame_num / 30) # eufy 15fps
        previous_in_count_south = len(unique_tracker_ids_south)

    # Check for cars turning left
    if len(unique_tracker_ids_lt) > previous_in_count_lt:
        print(f"Car turning left detected at frame {frame_num}")
        timestamps_lt.append(frame_num / 30) # eufy 15fps
        previous_in_count_lt = len(unique_tracker_ids_lt)

    # Check for cars turning right
    if len(unique_tracker_ids_rt) > previous_in_count_rt:
        print(f"Car turning right detected at frame {frame_num}")
        timestamps_rt.append(frame_num / 30) # eufy 15fps
        previous_in_count_rt = len(unique_tracker_ids_rt)

    straight_major = max(0, zone_sum_west - zone_sum_lt)
    straight_minor = max(0, zone_sum_south - zone_sum_rt)
    
    annotated_frame = zone_annotator_wb.annotate(scene=annotated_frame, label=f"Major volume: {str(zone_sum_west)}")
    annotated_frame = zone_annotator_sb.annotate(scene=annotated_frame, label=f"Minor volume: {str(zone_sum_south)}")
    annotated_frame = zone_annotator_lt.annotate(scene=annotated_frame, label=f"Minor volume straight: {str(straight_minor)}, major volume left turn: {str(zone_sum_lt)}")
    annotated_frame = zone_annotator_rt.annotate(scene=annotated_frame, label=f"Major volume straight: {str(straight_major)}, minor volume right turn: {str(zone_sum_rt)}")

    gaps_west = [timestamps_west[i+1] - timestamps_west[i] for i in range(len(timestamps_west)-1)]
    gaps_over_14_seconds_west = sum(1 for gap in gaps_west if gap > 14)
    average_gap_west = sum(gaps_west) / len(gaps_west) if gaps_west else 0

    text_anchor = sv.Point(199, 962)
    annotated_frame = draw_text(scene=annotated_frame, text=f"Critical gaps (14s+), wb: {gaps_over_14_seconds_west}, avg gap: {int(average_gap_west)}s", text_anchor=text_anchor, background_color=Color(r=255, g=255, b=255))

    #cv2.imshow('YOLO V8 Detection', annotated_frame)

    return trace_annotator.annotate(
        annotated_frame, detections=detections)

sv.process_video(
    source_path=VIDEO,
    target_path="result.mp4",
    callback=callback
)

# Calculate gaps for west
gaps_west = [timestamps_west[i+1] - timestamps_west[i] for i in range(len(timestamps_west)-1)]
# Calculate gaps for south
gaps_south = [timestamps_south[i+1] - timestamps_south[i] for i in range(len(timestamps_south)-1)]

# Count the number of gaps greater than 15 seconds
gaps_over_14_seconds_west = sum(1 for gap in gaps_west if gap > 14)
gaps_over_14_seconds_south = sum(1 for gap in gaps_south if gap > 14)

print(f"Gaps over 14 seconds, westbound: {gaps_over_14_seconds_west}")

for i in range(len(timestamps_west)-1):
    gap = timestamps_west[i+1] - timestamps_west[i]
    print(f"Starting timestamp: {round(timestamps_west[i],2)}s, Gap: {int(gap)}s")
    
print(f"Gaps over 14 seconds, southbound: {gaps_over_14_seconds_south}")

for i in range(len(timestamps_south)-1):
    gap = timestamps_south[i+1] - timestamps_south[i]
    print(f"Starting timestamp: {round(timestamps_south[i],2)}s, Gap: {int(gap)}s")
    
#print(f"Westbound cars: {zone_sum_west}")