Importance-Aware Semantic Segmentation in Self-Driving with Discrete Wasserstein Training

TL;DR

This paper introduces an importance-aware Wasserstein training framework for semantic segmentation in self-driving, emphasizing critical classes like pedestrians to improve safety-critical accuracy.

Contribution

It proposes a novel importance-aware Wasserstein loss that incorporates class importance into segmentation training, extending ground metrics with flexible functions for better performance.

Findings

Superior segmentation of critical classes on CamVid and Cityscapes

Effective integration with multiple backbone architectures

Outperforms traditional cross entropy-based methods

Abstract

Semantic segmentation (SS) is an important perception manner for self-driving cars and robotics, which classifies each pixel into a pre-determined class. The widely-used cross entropy (CE) loss-based deep networks has achieved significant progress w.r.t. the mean Intersection-over Union (mIoU). However, the cross entropy loss can not take the different importance of each class in an self-driving system into account. For example, pedestrians in the image should be much more important than the surrounding buildings when make a decisions in the driving, so their segmentation results are expected to be as accurate as possible. In this paper, we propose to incorporate the importance-aware inter-class correlation in a Wasserstein training framework by configuring its ground distance matrix. The ground distance matrix can be pre-defined following a priori in a specific task, and the previous importance-ignored methods can be the particular cases. From an optimization perspective, we also extend our ground metric to a linear, convex or concave increasing function $w.r.t.$ pre-defined ground distance. We evaluate our method on CamVid and Cityscapes datasets with different backbones (SegNet, ENet, FCN and Deeplab) in a plug and play fashion. In our extenssive experiments, Wasserstein loss demonstrates superior segmentation performance on the predefined critical classes for safe-driving.