UNet-Based Keypoint Regression for 3D Cone Localization in Autonomous Racing

AI summary

Overview

• Research area: Computer vision for autonomous racing — keypoint regression (KPR) and 3D object localization, specifically detecting and locating track cones.
• Technical level: Intermediate. The paper assumes familiarity with convolutional neural networks, UNet encoder–decoder architectures, stereo vision and Kalman filtering, but explains each component in enough detail to follow.
• Scope: A UNet-based keypoint detection model trained on a large custom cone dataset, evaluated both on keypoint metrics and when integrated into a full autonomous-racing perception pipeline.

What This Paper Is About

Autonomous racing cars must know exactly where track cones are in 3D space to plan a safe racing line, but cones are small, far away, easily damaged or dirty, and photographed under constantly changing light and weather. Traditional feature-matching algorithms such as SIFT and SURF are unreliable in these conditions, and existing neural approaches for this task are often trained on small datasets and are too slow to run on the car's onboard computer. The authors build a UNet network that predicts six keypoints per cone, convert those keypoints into 3D positions using stereo camera geometry, and test whether the extra accuracy is worth the extra computation when the model is placed inside a real perception pipeline.

Key Contributions

1. Dataset: The largest publicly available labeled dataset of 25k annotated cone images, captured across varied perspectives and conditions, released on Kaggle. The authors report 24,904 total keypoint annotations, filtered down to a high-quality subset of 20k samples.
2. Novel KPR method: A UNet-based architecture designed specifically for accurately localizing cone keypoints in complex racing scenes.
3. Application in autonomous navigation: The KPR model is integrated into a full autonomous vehicle perception pipeline and evaluated through simulation, including a confusion matrix of detected cone colors and covariance analysis at varying distance bins.
4. Geometry-driven 3D supervision: The system derives 3D supervision from stereo disparity rather than requiring ground-truth 3D annotations, which the authors describe as making the approach scalable to larger datasets and a lightweight alternative to methods like NeRF or monocular depth estimation.

Main Findings

• Keypoint accuracy improvement: On unseen test images, the UNet achieved an MSE of 3.4172, Root MSE of 1.8486, Normalized Mean Error of 0.0263, standard deviation of 3.4550, and mAP of 0.83. The replicated ResNet baseline recorded MSE 6.3165, Root MSE 2.3458, Normalized Mean Error 0.0597, standard deviation 6.4299, and mAP 0.42.
• Custom mAP definition: Mean average precision was adapted to the task so that a prediction within three pixels of the ground truth counts as correct.
• Baseline context: The ResNet comparison model is a replication of the architecture used by Dhall et al., which was originally trained on 900 samples with seven keypoints; the authors retrained it on their 20k samples with six keypoints for a like-for-like comparison.
• Failure modes are rare and specific: Roughly 3% of cases, based on ROS2 bag tests, involved inaccurate keypoints — typically cones partially out of view or densely clustered in a single image, where poor detection can propagate into depth and racing-line errors downstream.
• Real-time feasibility: Enabling the KPR network increased load across all 12 CPU cores, with memory and swap usage showing only a slight increase and a peak disparity of approximately 7%. Peak GTX 1060 GPU usage rose by 3% from a baseline of 14%. The authors state these figures lie within their real-time requirements.
• End-to-end pipeline performance: The authors report high-quality performance across all metrics and present a per-color cone detection confusion matrix plus birds-eye-view covariance ellipses by distance bin for the integrated pipeline.
• Color estimation enabled: The six predicted keypoints provide a foundation for algorithmic color estimation (identifying the cone stripe and base), which the paper notes is not possible with SIFT because it cannot reliably outline the cone's six keypoints.
• Comparison to prior detection work (cited, not their own result): Katsamenis et al. applied YOLOv5 to real-time cone detection on roadwork images with an IoU score up to 91.31%, and a KPR-modified YOLOv4 reached 45.6% average precision on COCO2017 versus 45.4% for YOLOv4.

Methodology in Plain English

Building the data. The team wrote a custom annotation tool in Flask that connects to AWS S3 buckets and stores labels as JSON. Annotators placed six keypoints on each cone — points around the cone's reflective strip and around its base. Images were drawn from regions matching YOLOv8 bounding-box detections, based on the FSOCO dataset. After discarding images with the wrong number of points or imprecise labels, they kept 20k samples. The authors explain that one keypoint would be enough for basic position and four for a simple color decision, but six gives robustness to individual keypoint errors and the option to cross-validate.

Training setup. They used a standard UNet with encoder, decoder and bottleneck sections. Each block stacks two 2D convolutional layers with batch normalization and ReLU. Convolutions use 3x3 kernels, stride 2 when downsampling and stride 1 in the decoder, with padding of 1. Channels double down the encoder (64, 128, 256, 512) and halve on the way back up. Data was augmented with rotations and randomized boundary cropping using four transformations (NONE, ROTATE_90, ROTATE_180, ROTATE_270), with rotation matrices applied to both images and keypoint coordinates. The split was 70% training, 20% validation, 10% testing. The network takes three channels, six keypoints and 80x80 images, with a dropout rate of 0.3, trained with AdamW and an exponential learning rate scheduler that reduces the rate by a factor of 0.99 per epoch. The loss combines heatmap-based and position-based terms, with a CustomLoss supporting both L1 and smooth L1. Evaluation used MSE, MAPE and Average Confidence.

Comparing against alternatives. To benchmark, they rebuilt the ResNet-based architecture used by Dhall et al. but trained it on their 20k samples with six keypoints instead of the original 900 samples with seven keypoints. They also implemented SIFT, matching descriptors with FLANN nearest-neighbour search and filtering with log likelihood, mainly to support triangulation.

From keypoints to 3D. The authors load the checkpoint with the lowest validation loss. For each cone they take the mean x position of the six keypoints in the left and right stereo frames, compute the disparity, and derive depth as Z = fT/D, where f is the focal length from the camera intrinsic matrix, T is the baseline from the camera specification, and D is disparity. Depth and the mean x and y keypoint positions are then converted into a 3D point using the standard pinhole equations involving principal points and focal lengths. Because their KPR produces better keypoint correspondences than other methods, the resulting localization is more accurate.

System context. Frames from a Stereolab ZED2 camera enter the pipeline through YOLOv8, which produces bounding boxes. The KPR model runs alongside several other position estimation methods, and their outputs are combined with experimentally derived covariances inside an Extended Kalman Filter before feeding path planning. The integrated system was assessed in a dedicated simulator with known ground-truth poses and on a recorded ROS bag of real ZED2 stereo data.

Why This Matters

Impact on research. The work shows that a purpose-built keypoint regression network can roughly halve the error metrics of a ResNet-based baseline used previously in autonomous racing, while also producing a large open dataset that was previously unavailable. It also demonstrates a practical way to obtain 3D localization supervision from stereo geometry instead of expensive ground-truth 3D annotations, and it argues that traditional feature matchers like SIFT and SURF — and even general-purpose matchers such as SuperGlue, which the paper notes requires offline inference — are unsuitable for on-device racing perception. The qualitative claim the authors emphasize is a positive feedback loop: better perception yields a better racing line, which yields frames containing more useful cones, and so on.

Real-world applications:

• Formula Student and other student autonomous racing competitions, where track boundaries are defined by blue cones on the left and yellow cones on the right.
• Roadworks and highway maintenance zones, where cones mark temporary lanes (the paper cites prior YOLOv5 work on precisely this scenario).
• General autonomous vehicle perception stacks that need small-object detection and low-latency 3D landmark localization.
• Any robotic platform using stereo cameras where a lightweight, embeddable keypoint regressor is preferable to a heavy dense reconstruction method.

Industry relevance. The paper reports the computational cost of adding the network on the actual in-car PC, including how load spreads across 12 CPU cores and how GPU memory and usage change. For teams deploying perception on embedded hardware, that trade-off analysis — a modest increase in compute for a large gain in keypoint accuracy — is directly actionable, and the authors frame the extra demand as a fair trade-off.

Future Directions

• Add the rare failure cases (partially visible and densely clustered cones) to the training data, and introduce cropping and translation transformations to the augmentation set, since the current augmentations cover only rotations and boundary cropping.
• Introduce an interpretable KPR confidence score that could inform the covariance size used by the downstream Kalman filter.
• Extend the framework beyond keypoints toward full 3D scene understanding by regressing depth-aware object maps or semantic voxel grids.
• Explore image-text supervision and ways to mitigate occlusion scenarios in a computationally efficient manner.

Target Audience

This paper is most useful to autonomous racing and Formula Student teams building perception pipelines, computer vision engineers working on keypoint detection or small-object localization, and researchers interested in stereo-based 3D supervision without ground-truth 3D labels. It is also relevant to embedded and real-time systems engineers who need to judge the computational cost of adding a neural network stage to an existing perception stack. Readers without a background in convolutional networks or stereo geometry will find the paper's methodology sections dense, though the system-level framing and discussion are broadly accessible.

Authors’ abstract

Accurate cone localization in 3D space is essential in autonomous racing for precise navigation around the track. Approaches that rely on traditional computer vision algorithms are sensitive to environmental variations, and neural networks are often trained on limited data and are infeasible to run in real time. We present a UNet-based neural network for keypoint detection on cones, leveraging the largest custom-labeled dataset we have assembled. Our approach enables accurate cone position estimation and the potential for color prediction. Our model achieves substantial improvements in keypoint accuracy over conventional methods. Furthermore, we leverage our predicted keypoints in the perception pipeline and evaluate the end-to-end autonomous system. Our results show high-quality performance across all metrics, highlighting the effectiveness of this approach and its potential for adoption in competitive autonomous racing systems.