├── .gitignore ├── README.md ├── requirements.txt └── yolov8tracker.py /.gitignore: -------------------------------------------------------------------------------- 1 | ByteTrack 2 | data 3 | env 4 | results 5 | yolov8x.pt 6 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Tracking and counting of object using YOLO v8 2 | 3 | This repository contains Python code for tracking vehicles (such as cars, buses, and bikes) as they enter and exit the road, thereby incrementing the counters for incoming and outgoing vehicles. 4 | 5 | ## Installation 6 | 7 | ```bash 8 | 1. git clone https://github.com/sankalpvarshney/Track-And-Count-Object-using-YOLO.git 9 | 2. cd Track-And-Count-Object-using-YOLO 10 | 3. conda create --prefix ./env python=3.8 -y 11 | 4. conda activate ./env 12 | 5. pip install ultralytics 13 | 6. git clone https://github.com/ifzhang/ByteTrack.git 14 | 7. cd ByteTrack 15 | 8. sed -i 's/onnx==1.8.1/onnx==1.9.0/g' requirements.txt 16 | 9. pip install -q -r requirements.txt 17 | 10. python setup.py -q develop 18 | 11. pip install -q cython_bbox 19 | 12. pip install -q onemetric 20 | 13. pip install -q loguru lap 21 | 14. pip install numpy==1.22.4 22 | 15. pip install supervision==0.1.0 23 | ``` 24 | 25 | ## Usage 26 | 27 | Firstly set the crossing line co-ordinates inside the code i.e yolov8tracker.py for the incoming and outgoing vehicles. And then execute the python code as mentioned below. 28 | ### Linux 29 | 30 | ```bash 31 | python yolov8tracker.py -i -o 32 | ``` 33 | 34 | ### Python 35 | 36 | ```python 37 | from yolov8tracker import TrackObject 38 | obj = TrackObject(,) 39 | obj.process_video() 40 | ``` 41 | 42 | https://github.com/sankalpvarshney/Track-And-Count-Object-using-YOLO/assets/41926323/bbeb35b4-3f0f-49cd-b222-2bf92ac001f7 43 | 44 | 45 | 46 | ## Contributing 47 | 48 | Pull requests are welcome. For major changes, please open an issue first 49 | to discuss what you would like to change. 50 | 51 | Please make sure to update tests as appropriate. 52 | 53 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | torchvision==0.13.0 2 | ultralytics 3 | supervision==0.1.0 -------------------------------------------------------------------------------- /yolov8tracker.py: -------------------------------------------------------------------------------- 1 | from yolox.tracker.byte_tracker import BYTETracker, STrack 2 | from onemetric.cv.utils.iou import box_iou_batch 3 | from dataclasses import dataclass 4 | from supervision.draw.color import ColorPalette 5 | from supervision.geometry.dataclasses import Point 6 | from supervision.video.dataclasses import VideoInfo 7 | from supervision.video.source import get_video_frames_generator 8 | from supervision.video.sink import VideoSink 9 | from supervision.tools.detections import Detections, BoxAnnotator 10 | from supervision.tools.line_counter import LineCounter, LineCounterAnnotator 11 | from typing import List 12 | import numpy as np 13 | from ultralytics import YOLO 14 | from tqdm import tqdm 15 | import argparse 16 | 17 | parser = argparse.ArgumentParser( 18 | prog='yolov8', 19 | description='This program help to track object and maintain in & out count', 20 | epilog='Text at the bottom of help') 21 | 22 | parser.add_argument('-i', '--input',required=True) # option that takes a value 23 | parser.add_argument('-o', '--output',required=True) 24 | 25 | args = parser.parse_args() 26 | 27 | @dataclass(frozen=True) 28 | class BYTETrackerArgs: 29 | track_thresh: float = 0.25 30 | track_buffer: int = 30 31 | match_thresh: float = 0.8 32 | aspect_ratio_thresh: float = 3.0 33 | min_box_area: float = 1.0 34 | mot20: bool = False 35 | 36 | # converts Detections into format that can be consumed by match_detections_with_tracks function 37 | def detections2boxes(detections: Detections) -> np.ndarray: 38 | return np.hstack(( 39 | detections.xyxy, 40 | detections.confidence[:, np.newaxis] 41 | )) 42 | 43 | 44 | # converts List[STrack] into format that can be consumed by match_detections_with_tracks function 45 | def tracks2boxes(tracks: List[STrack]) -> np.ndarray: 46 | return np.array([ 47 | track.tlbr 48 | for track 49 | in tracks 50 | ], dtype=float) 51 | 52 | 53 | # matches our bounding boxes with predictions 54 | def match_detections_with_tracks( 55 | detections: Detections, 56 | tracks: List[STrack] 57 | ) -> Detections: 58 | if not np.any(detections.xyxy) or len(tracks) == 0: 59 | return np.empty((0,)) 60 | 61 | tracks_boxes = tracks2boxes(tracks=tracks) 62 | iou = box_iou_batch(tracks_boxes, detections.xyxy) 63 | track2detection = np.argmax(iou, axis=1) 64 | 65 | tracker_ids = [None] * len(detections) 66 | 67 | for tracker_index, detection_index in enumerate(track2detection): 68 | if iou[tracker_index, detection_index] != 0: 69 | tracker_ids[detection_index] = tracks[tracker_index].track_id 70 | 71 | return tracker_ids 72 | 73 | class TrackObject(): 74 | 75 | def __init__(self,input_video_path, output_video_path) -> None: 76 | 77 | self.model = YOLO("yolov8x.pt") 78 | self.model.fuse() 79 | # dict maping class_id to class_name 80 | self.CLASS_NAMES_DICT = self.model.model.names 81 | # class_ids of interest - car, motorcycle, bus and truck 82 | self.CLASS_ID = [2, 3, 5, 7] 83 | 84 | # settings 85 | self.LINE_START = Point(50, 1500) 86 | self.LINE_END = Point(3840-50, 1500) 87 | 88 | self.input_video_path = input_video_path 89 | self.output_video_path = output_video_path 90 | 91 | # create BYTETracker instance 92 | self.byte_tracker = BYTETracker(BYTETrackerArgs()) 93 | # create VideoInfo instance 94 | self.video_info = VideoInfo.from_video_path(self.input_video_path) 95 | # create frame generator 96 | self.generator = get_video_frames_generator(self.input_video_path) 97 | # create LineCounter instance 98 | self.line_counter = LineCounter(start=self.LINE_START, end=self.LINE_END) 99 | # create instance of BoxAnnotator and LineCounterAnnotator 100 | self.box_annotator = BoxAnnotator(color=ColorPalette(), thickness=4, text_thickness=4, text_scale=2) 101 | self.line_annotator = LineCounterAnnotator(thickness=4, text_thickness=4, text_scale=2) 102 | 103 | def process_video(self): 104 | 105 | # open target video file 106 | with VideoSink(self.output_video_path, self.video_info) as sink: 107 | # loop over video frames 108 | for frame in tqdm(self.generator, total=self.video_info.total_frames): 109 | # model prediction on single frame and conversion to supervision Detections 110 | results = self.model(frame) 111 | detections = Detections( 112 | xyxy=results[0].boxes.xyxy.cpu().numpy(), 113 | confidence=results[0].boxes.conf.cpu().numpy(), 114 | class_id=results[0].boxes.cls.cpu().numpy().astype(int) 115 | ) 116 | # filtering out detections with unwanted classes 117 | mask = np.array([class_id in self.CLASS_ID for class_id in detections.class_id], dtype=bool) 118 | detections.filter(mask=mask, inplace=True) 119 | # tracking detections 120 | tracks = self.byte_tracker.update( 121 | output_results=detections2boxes(detections=detections), 122 | img_info=frame.shape, 123 | img_size=frame.shape 124 | ) 125 | tracker_id = match_detections_with_tracks(detections=detections, tracks=tracks) 126 | detections.tracker_id = np.array(tracker_id) 127 | # filtering out detections without trackers 128 | mask = np.array([tracker_id is not None for tracker_id in detections.tracker_id], dtype=bool) 129 | detections.filter(mask=mask, inplace=True) 130 | # format custom labels 131 | labels = [ 132 | f"#{tracker_id} {self.CLASS_NAMES_DICT[class_id]} {confidence:0.2f}" 133 | for _, confidence, class_id, tracker_id 134 | in detections 135 | ] 136 | # updating line counter 137 | self.line_counter.update(detections=detections) 138 | # annotate and display frame 139 | frame = self.box_annotator.annotate(frame=frame, detections=detections, labels=labels) 140 | self.line_annotator.annotate(frame=frame, line_counter=self.line_counter) 141 | sink.write_frame(frame) 142 | 143 | 144 | 145 | if __name__ == "__main__": 146 | 147 | obj = TrackObject(args.input,args.output) 148 | obj.process_video() 149 | 150 | --------------------------------------------------------------------------------