Progressively Normalized Self-Attention Network for Video Polyp Segmentation

PNS-Net

This repository provides code for paper"Progressively Normalized Self-Attention Network for Video Polyp Segmentation" published at the MICCAI-2021 conference (arXiv Version | 中文版). If you have any questions about our paper, feel free to contact me. And if you like our PNS-Net or evaluation toolbox for your personal research, please cite this paper (BibTeX).

PNS-Net

Features

Hyper Real-time Speed: Our method, named Progressively Normalized Self-Attention Network (PNS-Net), can efficiently learn representations from polyp videos with real-time speed (~140fps) on a single NVIDIA RTX 2080 GPU without any post-processing techniques (e.g., Dense-CRF).
Plug-and-Play Module: The proposed core module, termed Normalized Self-attention (NS), utilizes channel split,query-dependent, and normalization rules to reduce the computational cost and improve the accuracy, respectively. Note that this module can be flexibly plugged into any framework customed.
Cutting-edge Performance: Experiments on three challenging video polyp segmentation (VPS) datasets demonstrate that the proposed PNS-Net achieves state-of-the-art performance.
One-key Evaluation Toolbox: We release the first one-key evaluation toolbox in the VPS field.

Overview

Introduction

Existing video polyp segmentation (VPS) models typically employ convolutional neural networks (CNNs) to extract features.
However, due to their limited receptive fields, CNNs can not fully exploit the global temporal and spatial information in successive video frames, resulting in false-positive segmentation results.
In this paper, we propose the novel PNS-Net (Progressively Normalized Self-attention Network), which can efficiently learn representations from polyp videos with real-time speed (~140fps) on a single RTX 2080 GPU and no post-processing.

Our PNS-Net is based solely on a basic normalized self-attention block, dispensing with recurrence and CNNs entirely. Experiments on challenging VPS datasets demonstrate that the proposed PNS-Net achieves state-of-the-art performance.
We also conduct extensive experiments to study the effectiveness of the channel split, soft-attention, and progressive learning strategy.
We find that our PNS-Net works well under different settings, making it a promising solution to the VPS task.

Framework Overview

MainFramework

Figure 1: Overview of the proposed PNS-Net, including the normalized self-attention block (see § 2.1) with a stacked (×R) learning strategy.
See § 2 in the paper for details.

Qualitative Results

Qualitive
Figure 2: Qualitative Results.

Proposed Baseline

Training/Testing

The training and testing experiments are conducted using PyTorch with
a single GeForce RTX 2080 GPU of 8 GB Memory.

Configuring your environment (Prerequisites):

Note that PNS-Net is only tested on Ubuntu OS with the following environments.
It may work on other operating systems as well but we do not guarantee that it will.
- Creating a virtual environment in terminal:
conda create -n PNSNet python=3.6.
- Installing necessary packages PyTorch 1.1:
```
conda create -n PNSNet python=3.6
conda activate PNSNet
conda install pytorch=1.1.0 torchvision -c pytorch
pip install tensorboardX tqdm Pillow==6.2.2
pip install git+https://github.com/pytorch/tnt.git@master
```
- Our core design is built on CUDA OP with torchlib.
  Please ensure the base CUDA toolkit
  version is 10.x (not at conda env), and then build the NS Block:
```
cd ./lib/PNS
python setup.py build develop
```
Downloading necessary data:
- downloading the training/testing dataset and move it into ./dataset/*,
  which can be found in this download link (Google Drive, 2GB).
- downloading pre-trained weights and move it into snapshot/ours/PNS.pth,
  which can be found in this download link (Google Drive, 99MB).
Training Configuration:
- First, run python MyTrain_Pretrain.py in the terminal for pretraining, and then, run python MyTrain_finetune.py for finetuning.
- Just enjoy it! Finish it and the snapshot would save in ./snapshot/PNS-Net/*.
Testing Configuration:
- After you download all the pre-trained model and testing dataset,
  just run MyTest_finetune.py to generate the final prediction map in ./res.
- Just enjoy it!
- The prediction results of all competitors and our PNS-Net can be found at Google Drive (7MB).

Evaluating your trained model:

One-key evaluation is written in MATLAB code (link),
please follow this the instructions in ./eval/main_VPS.m and just run it to generate the evaluation results in ./eval-Result/.

Authors: Ge-Peng Ji*, Yu-Cheng Chou*, Deng-Ping Fan, Geng Chen, Huazhu Fu, Debesh Jha, & Ling Shao.