[Paper Review] Rethinking Rotated Object Detection with Gaussian Wasserstein Distance Loss

Xue Yang, Junchi Yan+4 more|arXiv (Cornell University)|Jan 28, 2021

Advanced Neural Network Applications82 references204 citations

TL;DR

The paper replaces angle regression-based rotation losses with a Gaussian Wasserstein Distance (GWD) loss by modeling rotated boxes as 2-D Gaussian distributions, yielding differentiable, metric-consistent regression that mitigates boundary discontinuities and square-like issues across various bounding box definitions.

ABSTRACT

Boundary discontinuity and its inconsistency to the final detection metric have been the bottleneck for rotating detection regression loss design. In this paper, we propose a novel regression loss based on Gaussian Wasserstein distance as a fundamental approach to solve the problem. Specifically, the rotated bounding box is converted to a 2-D Gaussian distribution, which enables to approximate the indifferentiable rotational IoU induced loss by the Gaussian Wasserstein distance (GWD) which can be learned efficiently by gradient back-propagation. GWD can still be informative for learning even there is no overlapping between two rotating bounding boxes which is often the case for small object detection. Thanks to its three unique properties, GWD can also elegantly solve the boundary discontinuity and square-like problem regardless how the bounding box is defined. Experiments on five datasets using different detectors show the effectiveness of our approach. Codes are available at https://github.com/yangxue0827/RotationDetection and https://github.com/open-mmlab/mmrotate.

Motivation & Objective

• Identify core flaws in existing rotation detectors: metric-loss inconsistency, boundary discontinuity, and square-like problems.
• Propose a unified, differentiable loss based on Gaussian Wasserstein distance to approximate rotating IoU between boxes.
• Show that GWD-based regression is robust to bounding box definitions and improves detection across datasets.
• Provide open-source code to facilitate adoption and replication.

Proposed method

• Convert a rotated bounding box B(x, y, w, h, theta) into a 2-D Gaussian distribution with mean m=(x,y) and covariance Sigma derived from w, h, theta.
• Define Gaussian Wasserstein distance d between two Gaussians via d^2 = ||m1 - m2||^2 + Tr(Sigma1 + Sigma2 - 2(Sigma1^1/2 Sigma2 Sigma1^1/2)^1/2).
• Transform d^2 through a nonlinear function f and hyperparameter tau to obtain a differentiable affinity 1/(tau + f(d^2)) akin to IoU-based loss L_gwd = 1 - 1/(tau + f(d^2)).
• Integrate L_gwd into a multi-task RetinaNet framework with regression targets and focal loss for classification.
• Show that Sigma^1/2(w,h,theta) satisfies properties making D_oc and D_le representations equivalent under GWD, reducing box-definition sensitivity.
• Provide ablation studies across datasets (HRSC2016, UCAS-AOD, DOTA, MLT, ICDAR2015, etc.) demonstrating improvements when using GWD.

Experimental results

Research questions

• RQ1 Can Gaussian Wasserstein distance provide a differentiable and IoU-consistent surrogate for rotating IoU losses?
• RQ2 Does modeling rotated boxes as 2-D Gaussians resolve boundary discontinuity and square-like issues independent of bounding box definitions?
• RQ3 How does GWD-based regression perform across diverse datasets and detectors compared with traditional Smooth L1 losses?
• RQ4 Is the proposed GWD loss robust to different bounding box parameterizations (OpenCV D_oc vs long-edge D_le)?

Key findings

• GWD-based regression improves mAP on multiple datasets when replacing Smooth L1 loss (e.g., on DOTA, HRSC2016, UCAS-AOD).
• GWD provides a differentiable surrogate that remains informative even for non-overlapping boxes (IoU=0).
• GWD naturally mitigates boundary discontinuity and square-like problems and is largely agnostic to the bounding box definition used.
• Ablation shows gains across detectors (RetinaNet, R3Det) and backbones (R-50, R-101, R-152) with various training tricks enhancing performance.
• Extensive experiments across five public datasets and two detectors validate the effectiveness of the approach.
• Code is publicly available to reproduce results.