Universal Weakly Supervised Segmentation by Pixel-to-Segment Contrastive Learning

By Tsung-Wei Ke, Jyh-Jing Hwang, and Stella X. Yu

Weakly supervised segmentation requires assigning a label to every pixel based on training instances with partial annotations such as image-level tags, object bounding boxes, labeled points and scribbles. This task is challenging, as coarse annotations (tags, boxes) lack precise pixel localization whereas sparse annotations (points, scribbles) lack broad region coverage. Existing methods tackle these two types of weak supervision differently: Class activation maps are used to localize coarse labels and iteratively refine the segmentation model, whereas conditional random fields are used to propagate sparse labels to the entire image.

We formulate weakly supervised segmentation as a semi-supervised metric learning problem, where pixels of the same (different) semantics need to be mapped to the same (distinctive) features. We propose 4 types of contrastive relationships between pixels and segments in the feature space, capturing low-level image similarity, semantic annotation, co-occurrence, and feature affinity They act as priors; the pixel-wise feature can be learned from training images with any partial annotations in a data-driven fashion. In particular, unlabeled pixels in training images participate not only in data-driven grouping within each image, but also in discriminative feature learning within and across images. We deliver a universal weakly supervised segmenter with significant gains on Pascal VOC and DensePose.

Code Base

This release of code is based on SegSort in ICCV 2019.

Prerequisites

1. Linux
2. Python3 (>=3.5)
3. Cuda >= 9.2 and Cudnn >= 7.6

Required Python Packages

1. pytorch >= 1.6
2. numpy
3. scipy
4. tqdm
5. easydict == 1.9
6. PyYAML
7. PIL
8. opencv
9. pydensecrf

Data Preparation

Pascal VOC 2012

1. Augmented Pascal VOC training set by SBD. Download link provided by SegSort. Please unzip it and put it besides the VOC2012/ folder as sbd/dataset/clsimg/.
2. Ground-truth semantic segmentation masks are reformatted as grayscale images. Download link provided by SegSort.
3. The over-segmentation masks are generated by combining contour detectors with gPb-owt-ucm. HED-owt-ucm masks provided by SegSort.
4. Scribble and point annotations by ScribbleSup. We follow "Regularized Losses for Weakly-supervised CNN Segmentation" to process scribble annotations. You can download the processed ground-truth scribble(100%, 80%, 50% and 30%) and point annotations, and put them under VOC2012/scribble folder. For scribbles, we set dilation size to 3; For points, we set dilation size to 6.
5. Cam annotations by SEAM. You can download the processed CAMs for image tag and bounding box annotations, and put them under VOC2012/cam. For image tags, we set alpha to 6 and probability threshold to 0.2. Pixels with probability less than 0.2 are considered as unlabeled regions. For bounding boxes, we normalize CAM to the range of 0 and 1 within each box, and set the probability threshold to 0.5. The pixel labels outside the boxes as background class.
6. Dataset layout:

   $DATAROOT/
       |-------- sbd/
       |          |-------- dataset/
       |                       |-------- clsimg/
       |
       |-------- VOC2012/
                  |-------- JPEGImages/
                  |-------- segcls/
                  |-------- hed/
                  |-------- scribble/
                  |            |-------- dilate_3/
                  |            |-------- dilate_6_0.0/
                  |
                  |-------- cam/
                               |-------- seam_a6_th0.2/
                               |-------- seambox_a6_th0.5/

DensePose

1. Images from MSCOCO. Download 2014 Train and Val Images (here).
2. Ground-truth semantic segmentation by DensePose. You can download the processed ground-truth segmentation here, and put them under segcls/ folder.
3. Ground-truth point annotation here. Put them under segcls/ folder.
4. The over-segmentation masks by PMI and gPb-owt-ucm. Download PMI-owt-ucm masks and put them under seginst/ folder.
5. Dataset layout:

   $DATAROOT/
       |-------- images/
       |          |-------- train2014/
       |          |-------- val2014/
       |
       |-------- segcls/
       |          |-------- densepose/gray/
       |          |-------- densepose_points/gray/
       |
       |-------- seginst/
                  |-------- pmi_0.1_256/

ImageNet Pre-Trained Models

We use the same ImageNet pretrained ResNet101 as EMANet. You can download the pretrained models here and put it under a new directory SPML/snapshots/imagenet/trained/. Note: we do not use MSCOCO pretrained ResNet.

Pascal VOC 2012 Trained Models.

We provide the download links for our SPML models trained using image-level tag/bounding box/scribble/point annotations on PASCAL VOC, and summarize the performance as follows. Note: we report the performance with denseCRF post-processing.

DensePose Trained Models.

We provide the download link for our SPML models trained using point annotations on DensePose here. We achieve 44.15% of mIoU on minival2014 set.

Bashscripts to Get Started

Pascal VOC 2012

• SPML with image-level tag.

source bashscripts/voc12/train_spml_tag.sh

• SPML with bounding box.

source bashscripts/voc12/train_spml_box.sh

• SPML with scribbles.

source bashscripts/voc12/train_spml_scribble.sh

• SPML with points.

source bashscripts/voc12/train_spml_point.sh

DensePose

• SPML with points.

source bashscripts/densepose/train_spml_point.sh

Citation

If you find this code useful for your research, please consider citing our paper Universal Weakly Supervised Segmentation by Pixel-to-Segment Contrastive Learning.

@inproceedings{ke2021spml,
  title={Universal Weakly Supervised Segmentation by Pixel-to-Segment Contrastive Learning},
  author={Ke, Tsung-Wei and Hwang, Jyh-Jing and Yu, Stella X},
  booktitle={International Conference on Learning Representations},
  pages={},
  year={2021}
}

License

SPML is released under the MIT License (refer to the LICENSE file for details).