Class counter in example traffic_analysis

[hugoles]

Hello, i downloaded the example from the 'examples/traffic_analysis' folder and tried to enhance it by adding more features, but I spent a lot of time trying to separate it by class and couldn't achieve a satisfactory result. For some reason, the count is increasing much more than it should (perhaps due to the instability of the bus detector), but I'm not sure if the problem is solely related to that. If possible, could someone help me.

here is the actual code:

import argparse
from typing import Dict, List, Set, Tuple

import os
import csv
import cv2
import numpy as np
from tqdm import tqdm
from ultralytics import YOLO
import datetime

import supervision as sv

COLORS = sv.ColorPalette.default()
'''
ZONE_IN_POLYGONS = [
    np.array([[1100, 700], [1100, 450], [1300, 450], [1300, 700]]),
]

ZONE_OUT_POLYGONS = [
    np.array([[950, 382], [1350, 382], [1350, 182], [950, 182]]),
]
'''

ZONE_IN_POLYGONS = [
    np.array([[1100, 850], [1100, 500], [1300, 500], [1300, 850]]),
]

ZONE_OUT_POLYGONS = [
    np.array([[250, 400], [1350, 400], [1350, 182], [250, 182]]),
]

class DetectionsManager:
    def __init__(self) -> None:
        self.tracker_id_to_zone_id: Dict[int, int] = {}
        self.counts: Dict[int, Dict[int, Set[int]]] = {}
        self.counts_car = 0
        self.counts_bus = 0

        self.unique_detect_ids: Dict[int, Set[int]] = {}


    def update(
        self,
        detections_all: sv.Detections,
        detections_in_zones: List[sv.Detections],
        detections_out_zones: List[sv.Detections],
    ) -> sv.Detections:
        for zone_in_id, detections_in_zone in enumerate(detections_in_zones):
            for tracker_id in detections_in_zone.tracker_id:
                self.tracker_id_to_zone_id.setdefault(tracker_id, zone_in_id)

        for zone_out_id, detections_out_zone in enumerate(detections_out_zones):
            for tracker_id in detections_out_zone.tracker_id:
                if tracker_id in self.tracker_id_to_zone_id:
                    zone_in_id = self.tracker_id_to_zone_id[tracker_id]
                    self.counts.setdefault(zone_out_id, {})
                    self.counts[zone_out_id].setdefault(zone_in_id, set())
                    self.counts[zone_out_id][zone_in_id].add(tracker_id)
                    if hasattr(detections_out_zone, 'tracker_id'):
                            for tracker_id in enumerate(detections_out_zone.tracker_id):
                                #maybe:
                                #  detect_id = detections_out_zone.class_id
                                #  if detect_id not in self.unique_detect_ids.get(tracker_id, set()):
                                #   self.unique_detect_ids.setdefault(tracker_id, set()).add(detect_id)
                                #   if detect_id == 1:
                                #       self.counts_car += 1
                                #   elif detect_id == 0:
                                #       self.counts_bus += 1
                                for detect_id in detections_out_zone.class_id:
                                    if detect_id not in self.unique_detect_ids.get(tracker_id, set()):
                                        self.unique_detect_ids.setdefault(tracker_id, set()).add(detect_id)
                                        if detect_id == 1:
                                            self.counts_car += 1
                                        elif detect_id == 0:
                                            self.counts_bus += 1                        


        detections_all.class_id = np.vectorize(
            lambda x: self.tracker_id_to_zone_id.get(x, -1)
        )(detections_all.tracker_id)
        
        return detections_all[detections_all.class_id != -1]


def initiate_polygon_zones(
    polygons: List[np.ndarray],
    frame_resolution_wh: Tuple[int, int],
    triggering_position: sv.Position = sv.Position.CENTER,
) -> List[sv.PolygonZone]:
    return [
        sv.PolygonZone(
            polygon=polygon,
            frame_resolution_wh=frame_resolution_wh,
            triggering_position=triggering_position,
        )
        for polygon in polygons
    ]


class VideoProcessor:
    def __init__(
        self,
        source_weights_path: str,
        source_video_path: str,
        target_video_path: str = None,
        confidence_threshold: float = 0.3,
        iou_threshold: float = 0.7,
    ) -> None:
        self.conf_threshold = confidence_threshold
        self.iou_threshold = iou_threshold
        self.source_video_path = source_video_path
        self.target_video_path = target_video_path

        self.model = YOLO(source_weights_path)
        self.tracker = sv.ByteTrack()

        self.video_info = sv.VideoInfo.from_video_path(source_video_path)
        self.zones_in = initiate_polygon_zones(
            ZONE_IN_POLYGONS, self.video_info.resolution_wh, sv.Position.CENTER
        )
        self.zones_out = initiate_polygon_zones(
            ZONE_OUT_POLYGONS, self.video_info.resolution_wh, sv.Position.CENTER
        )

        self.box_annotator = sv.BoxAnnotator(color=COLORS)
        self.trace_annotator = sv.TraceAnnotator(
            color=COLORS, position=sv.Position.CENTER, trace_length=100, thickness=2
        )
        self.detections_manager = DetectionsManager()


    def process_video(self):
        frame_generator = sv.get_video_frames_generator(
            source_path=self.source_video_path
        )

        if self.target_video_path:
            with sv.VideoSink(self.target_video_path, self.video_info) as sink:
                for frame in tqdm(frame_generator, total=self.video_info.total_frames):
                    annotated_frame = self.process_frame(frame)
                    sink.write_frame(annotated_frame)
        else:
            for frame in tqdm(frame_generator, total=self.video_info.total_frames):
                annotated_frame = self.process_frame(frame)
                cv2.imshow("Processed Video", annotated_frame)
                if cv2.waitKey(1) & 0xFF == ord("q"):
                    break
            cv2.destroyAllWindows()

    def annotate_frame(
        self, frame: np.ndarray, detections: sv.Detections
    ) -> np.ndarray:
        annotated_frame = frame.copy()
        for i, (zone_in, zone_out) in enumerate(zip(self.zones_in, self.zones_out)):
            annotated_frame = sv.draw_polygon(
                annotated_frame, zone_in.polygon, COLORS.colors[i]
            )
            annotated_frame = sv.draw_polygon(
                annotated_frame, zone_out.polygon, COLORS.colors[i]
            )

        labels = [f"#{tracker_id}" for tracker_id in detections.tracker_id]
        annotated_frame = self.trace_annotator.annotate(annotated_frame, detections)
        annotated_frame = self.box_annotator.annotate(
            annotated_frame, detections, labels
        )

        for zone_out_id, zone_out in enumerate(self.zones_out):
            zone_center = sv.get_polygon_center(polygon=zone_out.polygon)
            if zone_out_id in self.detections_manager.counts:
                counts = self.detections_manager.counts[zone_out_id]
                for i, zone_in_id in enumerate(counts):
                    count = len(self.detections_manager.counts[zone_out_id][zone_in_id])
                    count_car = self.detections_manager.counts_car
                    count_bus = self.detections_manager.counts_bus
                    text_anchor = sv.Point(x=zone_center.x, y=zone_center.y + 40 * i)
                    annotated_frame = sv.draw_text(
                        scene=annotated_frame,
                        text=str(count) + "," + str(count_car) + "," + str(count_bus),
                        text_anchor=text_anchor,
                        background_color=COLORS.colors[zone_in_id],
                    )

        return annotated_frame

    def process_frame(self, frame: np.ndarray) -> np.ndarray:
        results = self.model(
            frame, verbose=False, conf=self.conf_threshold, iou=self.iou_threshold
        )[0]
        detections = sv.Detections.from_ultralytics(results)
        
        detections = self.tracker.update_with_detections(detections)

        detections_in_zones = []
        detections_out_zones = []

        for i, (zone_in, zone_out) in enumerate(zip(self.zones_in, self.zones_out)):
            detections_in_zone = detections[zone_in.trigger(detections=detections)]
            detections_in_zones.append(detections_in_zone)
            detections_out_zone = detections[zone_out.trigger(detections=detections)]
            detections_out_zones.append(detections_out_zone)

        detections = self.detections_manager.update(
            detections, detections_in_zones, detections_out_zones
        )
        return self.annotate_frame(frame, detections)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Traffic Flow Analysis with YOLO and ByteTrack"
    )

    parser.add_argument(
        "--source_weights_path",
        required=True,
        help="Path to the source weights file",
        type=str,
    )
    parser.add_argument(
        "--source_video_path",
        required=True,
        help="Path to the source video file",
        type=str,
    )
    parser.add_argument(
        "--target_video_path",
        default=None,
        help="Path to the target video file (output)",
        type=str,
    )
    parser.add_argument(
        "--confidence_threshold",
        default=0.3,
        help="Confidence threshold for the model",
        type=float,
    )
    parser.add_argument(
        "--iou_threshold", default=0.7, help="IOU threshold for the model", type=float
    )

    args = parser.parse_args()
    processor = VideoProcessor(
        source_weights_path=args.source_weights_path,
        source_video_path=args.source_video_path,
        target_video_path=args.target_video_path,
        confidence_threshold=args.confidence_threshold,
        iou_threshold=args.iou_threshold,
    )
    processor.process_video()

    model = YOLO(args.source_weights_path)
    print(model.model.names)
    
    e = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    
    video = os.path.basename(args.target_video_path)
    model = os.path.basename(args.source_weights_path)
    
    with open('counter.csv', 'a') as f:
        for key, value in processor.detections_manager.counts.items():
            for key2, value2 in value.items():
                f.write(f'{e},{video},{model},{len(value2)}\n')

[SkalskiP]

Hi, @hugoles 👋🏻! I'm very sorry; I honestly won't have time to debug this code. I am happy to answer questions about specific parts.

[hugoles]

Oh, sorry, I just posted the entire code for context. I only changed a few parts from the model in 'supervision/examples/traffic_analysis,' but the main part I modified is :

class DetectionsManager:
def __init__(self) -> None:
    self.tracker_id_to_zone_id: Dict[int, int] = {}
    self.counts: Dict[int, Dict[int, Set[int]]] = {}
    self.counts_car = 0
    self.counts_bus = 0

    self.unique_detect_ids: Dict[int, Set[int]] = {}
def update(
    self,
    detections_all: sv.Detections,
    detections_in_zones: List[sv.Detections],
    detections_out_zones: List[sv.Detections],
) -> sv.Detections:
    for zone_in_id, detections_in_zone in enumerate(detections_in_zones):
        for tracker_id in detections_in_zone.tracker_id:
            self.tracker_id_to_zone_id.setdefault(tracker_id, zone_in_id)

    for zone_out_id, detections_out_zone in enumerate(detections_out_zones):
        for tracker_id in detections_out_zone.tracker_id:
            if tracker_id in self.tracker_id_to_zone_id:
                zone_in_id = self.tracker_id_to_zone_id[tracker_id]
                self.counts.setdefault(zone_out_id, {})
                self.counts[zone_out_id].setdefault(zone_in_id, set())
                self.counts[zone_out_id][zone_in_id].add(tracker_id)
                if hasattr(detections_out_zone, 'tracker_id'):
                        for tracker_id in enumerate(detections_out_zone.tracker_id):
                            #maybe:
                            #  detect_id = detections_out_zone.class_id
                            #  if detect_id not in self.unique_detect_ids.get(tracker_id, set()):
                            #   self.unique_detect_ids.setdefault(tracker_id, set()).add(detect_id)
                            #   if detect_id == 1:
                            #       self.counts_car += 1
                            #   elif detect_id == 0:
                            #       self.counts_bus += 1
                            for detect_id in detections_out_zone.class_id:
                                if detect_id not in self.unique_detect_ids.get(tracker_id, set()):
                                    self.unique_detect_ids.setdefault(tracker_id, set()).add(detect_id)
                                    if detect_id == 1:
                                        self.counts_car += 1
                                    elif detect_id == 0:
                                        self.counts_bus += 1                        


    detections_all.class_id = np.vectorize(
        lambda x: self.tracker_id_to_zone_id.get(x, -1)
    )(detections_all.tracker_id)
    
    return detections_all[detections_all.class_id != -1]

I don't know how to prevent the counter from repeating the same car. I think using a tracker_id could help, but for some reason, when I run it, every time a bus enters the area, there seems to be something wrong because the counter increases significantly (e.g., counting 8 times for a single bus) and then stops when the bus leaves. I'm really stuck here.

[hugoles]

@SkalskiP can u see a resolution to this?

[hugoles]

Update:
I believe my logic was minimally correct, but to differentiate by class_id, I think the most effective approach is using a "line" like a linezone function. I attempted to recreate a traffic analysis filter considering the class_id, but I encountered stability issues.

I resolved this problem by reducing the "zone_out" area with a 10px margin.

and optimizing the code.
Before:

After:

I believe the impact of these changes was counterproductive, resulting in the incorrect counting of cars and buses. The issue seems to be related to "blinks" in detection IDs, causing the addition of all vehicles in the zone when the function detects the class_id. For example, if a bus (class_id == 0) passes the zone out, and there is both a car and a bus, both are increased in the bus count, whereas I intended to count them separately as 1 car AND 1 bus.

Here are some examples (global counter, car counter, bus counter):

The code has a large zone and is relatively slower, causing counting problems.

Here the code is with shorter zone but with still code problem

Here the last code with some space and more accurate