The AI City 2024 Challenge-Track 5 contains the ground truth for 9 classes:
motorbike
DHelmet
DNoHelmet
P1Helmet
P1NoHelmet
P2Helmet
P2NoHelmet
P0Helmet
P0NoHelmet
This repository merges the 9 ground truth classes into 1 single motocycle class for custom object detection training.
gt_display: shows videos with the original 9-class bounding boxes (OPTIONAL).
extract_frms: extracts and saves frames from the videos.
concat.py: contains all the merging code.
tvt_split.py: splits the data into train, validation and test subsets after merging (OPTIONAL).
test.py: show videos with the motorcycle bounding boxes after merging is done (OPTIONAL).
You can comment out the optional functionalities in main.py.
How the merging works:
- Group the 9 classes into 3 main classes: motorcycle, driver and passenger.
- Assign and merge the 2 closest (if any)** driver and passenger** bounding box into 1 driver bounding box.
- Assign and merger the 2 closest (if any)** driver and motorcycle** bounding box into 1 motorcycle bounding box.
- Apply Non-Maximum Supression to filter out duplicate or remaining bounding boxes.
Main function in concat.py:
iou(box1, box2): calculate the Intersection of Union between 2 bounding box, which will later be used as thresholds in bb_assign and nms.
bb_assign(bboxes1, bboxes2, threshold): assign the 2 closest bounding boxes into a pair via calculating their center distance. threshold is the maximum distance between 2 assigned boxes.
concatenate_bb(bboxes): merge the assigned bounding boxes into 1 larger box with the x_min, y_min, x_max, y_max taken from the original bounding box coordinates.
nsm(bboxes, iou_threshold): Non-Maximum Supression algorithm for filtering out the duplicate bounding boxes. iou_threshold is the maximum iou value for the bounding box to be kept.
How to run:
Run main.py:
python main.py --video_dir "YOUR_VIDEO_PATH" --root_path "YOUR_ROOT_PATH" --gt_path "YOUR_GROUND_TRUTH"
--root-path is where the images and labels will be stored.
Merging visualization:
See more images examples: Notebook