An Uncertainty-aware DETR Enhancement Framework for Object Detection

TL;DR

This paper introduces an uncertainty-aware enhancement for DETR object detectors that models bounding boxes as distributions, improving localization accuracy and robustness across general and domain-specific tasks.

Contribution

It proposes a novel framework that incorporates distribution modeling, Gromov-Wasserstein loss, and Bayesian filtering to enhance DETR-based detectors with uncertainty estimation.

Findings

Improved detection accuracy on COCO benchmark.

Achieved state-of-the-art leukocyte detection results.

Framework is compatible with existing DETR variants.

Abstract

This paper investigates the problem of object detection with a focus on improving both the localization accuracy of bounding boxes and explicitly modeling prediction uncertainty. Conventional detectors rely on deterministic bounding box regression, ignoring uncertainty in predictions and limiting model robustness. In this paper, we propose an uncertainty-aware enhancement framework for DETR-based object detectors. We model bounding boxes as multivariate Gaussian distributions and incorporate the Gromov-Wasserstein distance into the loss function to better align the predicted and ground-truth distributions. Building on this, we derive a Bayes Risk formulation to filter high-risk information and improve detection reliability. We also propose a simple algorithm to quantify localization uncertainty via confidence intervals. Experiments on the COCO benchmark show that our method can be effectively integrated into existing DETR variants, enhancing their performance. We further extend our framework to leukocyte detection tasks, achieving state-of-the-art results on the LISC and WBCDD datasets. These results confirm the scalability of our framework across both general and domain-specific detection tasks. Code page: https://github.com/ParadiseforAndaChen/An-Uncertainty-aware-DETR-Enhancement-Framework-for-Object-Detection.