Learning Occlusion-Robust Vision Transformers for Real-Time UAV Tracking

TL;DR

This work tackles occlusion challenges in real-time UAV tracking by introducing ORTrack, a Vision Transformer-based single-stream tracker that learns Occlusion-Robust Representations (ORR) by enforcing invariance of target features under random masking modeled by a spatial Cox process. It further enhances efficiency with Adaptive Feature-Based Knowledge Distillation (AFKD), yielding a lighter student model (ORTrack-D) that closely mimics a trained teacher without sacrificing real-time performance. The method comprises a two-stage training pipeline (teacher for ORR, then student with AFKD) and a prediction head that outputs bounding boxes and confidence scores; the training losses combine an occlusion-robust loss with standard localization/classification terms. Extensive experiments across UAVDT, VisDrone2018, DTB70, and UAV123 show state-of-the-art real-time accuracy and speed, with ORTrack-DeiT achieving leading metrics and ORTrack-D-DeiT offering substantial speedups for deployment while preserving most performance benefits.

Abstract

Single-stream architectures using Vision Transformer (ViT) backbones show great potential for real-time UAV tracking recently. However, frequent occlusions from obstacles like buildings and trees expose a major drawback: these models often lack strategies to handle occlusions effectively. New methods are needed to enhance the occlusion resilience of single-stream ViT models in aerial tracking. In this work, we propose to learn Occlusion-Robust Representations (ORR) based on ViTs for UAV tracking by enforcing an invariance of the feature representation of a target with respect to random masking operations modeled by a spatial Cox process. Hopefully, this random masking approximately simulates target occlusions, thereby enabling us to learn ViTs that are robust to target occlusion for UAV tracking. This framework is termed ORTrack. Additionally, to facilitate real-time applications, we propose an Adaptive Feature-Based Knowledge Distillation (AFKD) method to create a more compact tracker, which adaptively mimics the behavior of the teacher model ORTrack according to the task's difficulty. This student model, dubbed ORTrack-D, retains much of ORTrack's performance while offering higher efficiency. Extensive experiments on multiple benchmarks validate the effectiveness of our method, demonstrating its state-of-the-art performance. Codes is available at https://github.com/wuyou3474/ORTrack.

Figures (5)

• Figure 1: Compared to SOTA UAV trackers on UAVDT, our ORTrack-DeiT sets a new record with 83.4% precision and a speed of 236 FPS. Our ORTrack-D-DeiT strikes a better trade-off with 82.5% precision and a speed of about 313 FPS.
• Figure 2: Overview of the proposed ORTrack framework, which includes separate training pipelines for a teacher and a student model. Note that the spatial Cox process-based masking and occlusion-robust representation learning are applied only in the teacher pipeline. Once the teacher is trained, its weights are fixed for training the student model with the proposed adaptive knowledge distillation.
• Figure 3: Attribute-based comparison on the partial occlusion subset of VisDrone2018 wen2018visdrone. ORTrack-DeiT* refers to ORTrack-DeiT without applying the occlusion-robust enhancement.
• Figure 4: Qualitative evaluation on 3 video sequences from, respectively, UAV123 Mueller2016ABA, UAVDT du2018the, and VisDrone2018 wen2018visdrone (i.e., person9, S1607, and uav0000180_00050_s).
• Figure 5: Visualize the attention map (left) and feature map (right) of the target images. The first row displays the search and masked images with masking ratios of 0%, 10%, 30%, and 70%. The second and third rows show the attention and feature maps generated by ORTrack-DeiT, with and without ORR, respectively.