shjo-april / RecurSeed_and_EdgePredictMix

[Under Review] RecurSeed and EdgePredictMix: Single-stage learning is sufficient for Weakly-Supervised Semantic Segmentation
44 stars 8 forks source link
deep-learning pytorch semantic-segmentation

PWC PWC PWC

RecurSeed and EdgePredictMix

This repository is the official implementation of "RecurSeed and EdgePredictMix: Single-stage learning is sufficient for Weakly-Supervised Semantic Segmentation". Please feel free to reach out for any questions or discussions.

Citation

Abstract

Although weakly-supervised semantic segmentation using only image-level labels (WSSS-IL) is potentially useful, its low performance and implementation complexity still limit its application. The main causes are (a) non-detection and (b) false-detection phenomena: (a) The class activation maps refined from existing WSSS-IL methods still only represent partial regions for large-scale objects, and (b) for small-scale objects, over-activation causes them to deviate from the object edges. We propose RecurSeed which alternately reduces non- and false detections through recursive iterations, thereby implicitly finding an optimal junction that minimizes both errors. We also propose a novel data augmentation (DA) approach called EdgePredictMix, which further expresses an object's edge by utilizing the probability difference information between adjacent pixels in combining the segmentation results, thereby compensating for the shortcomings when applying the existing DA methods to WSSS. We achieved new state-of-the-art performances on both the PASCAL VOC 2012 and MS COCO 2014 benchmarks (VOC \emph{val} $74.4\%$, COCO \emph{val} $46.4\%$).

Overview

Update

[04/14/2022] Released "RecurSeed and CertainMix" for testing.

[08/19/2022] Released "RecurSeed and EdgePredictMix" for testing.

[10/18/2022] Released additional results such as pseudo labels and DeepLabv2.

[01/26/2023] Released additional weights trained on the COCO 2014 dataset and updating citation.

Setup

Setting up for this project involves installing dependencies and preparing datasets. The code is tested on Ubuntu 20.04 with NVIDIA GPUs and CUDA installed.

Installing dependencies

To install all dependencies, please run the following:

python3 -m pip install git+https://github.com/lucasb-eyer/pydensecrf.git
python3 -m pip install -r requirements.txt

Preparing datasets

Please download following VOC and COCO datasets. Each dataset has a different directory structure. Therefore, we modify directory structures of all datasets for a comfortable implementation.

1. PASCAL VOC 2012

Download PASCAL VOC 2012 dataset from our [Google Drive].

2. MS COCO 2014

Download MS COCO 2014 dataset from our [Google Drive].

Create a directory "../VOC2012/" for storing the dataset and appropriately place each dataset to have the following directory structure.

    ../                               # parent directory
    ├── ./                            # current (project) directory
    │   ├── core/                     # (dir.) implementation of RecurSeed and CertainMix
    │   ├── data/                     # (dir.) information per dataset (including class names and the number of classes)
    │   ├── tools/                    # (dir.) helper functions
    │   ├── README.md                 # intstruction for a reproduction
    │   └── ... some python files ...
    |
    ├── VOC2012/                      # PASCAL VOC 2012
    │   ├── train/              
    │   │   ├── image/     
    │   │   ├── mask/        
    │   │   └── xml/        
    │   ├── train_aug/
    │   │   ├── image/     
    │   │   ├── mask/        
    │   │   └── xml/   
    │   ├── validation/
    │   │   ├── image/     
    │   │   ├── mask/        
    │   │   └── xml/   
    │   └── test/
    │       └── image/
    |
    └── COCO2014/                     # MS COCO 2014
        ├── train/              
        │   ├── image/     
        │   ├── mask/        
        │   └── xml/
        └── validation/
            ├── image/     
            ├── mask/        
            └── xml/

Visualization

We prepared a jupyter notebook for visualization.

Training

The whole code and commands are under review and will release soon.

Evaluation

Release our weights, predicted masks, and official results (anonymous link).

Stage Backbone Pretrained weight VOC val VOC test
single-stage ResNet-50 weight link link
multi-stage (DeepLabv2) ResNet-101 weight link mask link mask
multi-stage (DeepLabv3+) ResNet-101 weight link mask link mask
multi-stage (DeepLabv3+, COCO) ResNet-101 weight - -

Release pseudo labels on PASCAL VOC 2012 dataset after applying a random walk with our single-stage method. RS+EPM.zip

Below lines are testing commands to reproduce our method.

1. Single-stage Results

# Generate initial seeds produced from our single-stage method.
python3 infer_rsepm.py \
--gpus 2 --root_dir ../VOC2012/ --dataset VOC --domain validation \
--backbone resnet50 --tag "ResNet50@VOC@RS+EPM@Official" --conf_th 0.34

# Convert initial seeds to predicted masks with CRF.
python3 convert_seed_to_pseudo_masks.py \
--gpus 0 --root_dir ../VOC2012/ --dataset VOC --domain validation \
--folder predictions --tag "ResNet50@VOC@RS+EPM@Official" 

# Calculate the mIoU.
# [validation] http://host.robots.ox.ac.uk:8080/anonymous/LKO4IS.html
# [test] http://host.robots.ox.ac.uk:8080/anonymous/9MLLHH.html
python3 evaluate.py \
--root_dir ../VOC2012/ --domain validation \
--folder pseudo-labels --tag "ResNet50@VOC@RS+EPM@Official"

2. Multi-stage Results

# Generate initial seeds produced from final segmentation model.
python3 infer_seg.py \
--gpus 3 --root_dir ../VOC2012/ --dataset VOC --domain validation \
--backbone resnet101 --decoder deeplabv3+ \
--tag "ResNet101@VOC@RS+EPM@MS@Official"

# Convert initial seeds to predicted masks with CRF.
python3 convert_seed_to_pseudo_masks.py \
--gpus 0 --root_dir ../VOC2012/ --dataset VOC --domain validation \
--folder predictions --tag "ResNet101@VOC@RS+EPM@MS@Official"

# Calculate the mIoU.
# [validation] http://host.robots.ox.ac.uk:8080/anonymous/SNDUAQ.html
# [test] http://host.robots.ox.ac.uk:8080/anonymous/HUHADT.html
python3 evaluate.py \
--root_dir ../VOC2012/ --domain validation \
--folder pseudo-labels --tag "ResNet101@VOC@RS+EPM@MS@Official"