Tree Energy Loss: Towards Sparsely Annotated Semantic Segmentation

Introduction

This repository is an official implementation of the CVPR 2022 paper Tree Energy Loss: Towards Sparsely Annotated Semantic Segmentation.

TEL

Abstract. Sparsely annotated semantic segmentation (SASS) aims to train a segmentation network with coarse-grained (i.e., point-, scribble-, and block-wise) supervisions, where only a small proportion of pixels are labeled in each image. In this paper, we propose a novel tree energy loss for SASS by providing semantic guidance for unlabeled pixels. The tree energy loss represents images as minimum spanning trees to model both low-level and high-level pair-wise affinities. By sequentially applying these affinities to the network prediction, soft pseudo labels for unlabeled pixels are generated in a coarse-to-fine manner, achieving dynamic online self-training. The tree energy loss is effective and easy to be incorporated into existing frameworks by combining it with a traditional segmentation loss.

News

(03/03/2022) Tree Energy Loss has been accepted by CVPR 2022.

(15/03/2022) Update codes and models.

Main Results

Requirements

Linux, Python>=3.6, CUDA>=10.0, pytorch == 1.7.1

Installation

This implementation is built upon openseg.pytorch and TreeFilter-Torch. Many thanks to the authors for the efforts.

git clone https://github.com/megvii-research/TreeEnergyLoss
cd TreeEnergyLoss
pip install -r requirements.txt
cd kernels/lib_tree_filter
sudo python3 setup.py build develop

Sparse Annotation Preparation

Please first follow the Getting Started for data preparation.
Download scribble and point annotations provided by GatedCRFLoss.
Download block annotations provided by ours.

and finally, the dataset directory should look like:

$DATA_ROOT
├── cityscapes
│   ├── train
│   │   ├── image
│   │   ├── label
│   │   └── sparse_label
│   │       ├── block10
│   │       ├── block20
│   │       └── block50
│   ├── val
│   │   ├── image
│   │   └── label
├── ade20k
│   ├── train
│   │   ├── image
│   │   ├── label
│   │   └── sparse_label
│   │       ├── block10
│   │       ├── block20
│   │       └── block50
│   ├── val
│   │   ├── image
│   │   └── label
├── voc2012
│   ├── voc_scribbles.zip
│   ├── voc_whats_the_point.json
│   └── voc_whats_the_point_bg_from_scribbles.json

Block-Supervised Setting

(1) To evaluate the released models:

bash scripts/cityscapes/hrnet/demo.sh val block50

(2) To train and evaluate your own models:

bash scripts/cityscapes/hrnet/train.sh train model_name

bash scripts/cityscapes/hrnet/train.sh val model_name

Point-supervised and Scribble-supervised Settings

(1) To evaluate the released models:

bash scripts/voc2012/deeplab/demo.sh val scribble

(2) To train and evaluate your own models:

bash scripts/voc2012/deeplab/train.sh train model_name

bash scripts/voc2012/deeplab/train.sh val model_name

TreeEnergyLoss

Install / Use

README