C³Net: Coherent Context Complement Network for Polyp Segmentation

Official PyTorch implementation

[Architecture] [Environment] [Data Preparation] [Training] [Evaluation] [Notebook Demo]

Overview

Accurate polyp segmentation requires a model to preserve fine boundary details while capturing long-range context across polyps with substantial variations in size, shape, and appearance. C³Net addresses this problem with a U-shaped encoder-decoder that combines convolutional local modeling and visual state-space global modeling.

The framework contains three principal components:

• Multi-Granularity Modeling Block (MGMB) combines a Multi-Kernel Filter (MKF) with a Visual State Space Block (VSSB) to model local and global context.
• Referring Scale Context Aggregation (RSCA) uses decoder features as region priors to aggregate multi-scale encoder features.
• Dynamic Context Fusion (DCF) adaptively aligns and fuses encoder and decoder representations.

Boundary-aware deep supervision is applied at multiple decoding stages to improve region prediction and boundary delineation. The repository provides the model implementation, dataset loader, boundary generation utility, training and evaluation scripts, configuration presets, and an executable evaluation notebook.

Architecture

The left panel shows the complete encoder-decoder architecture. The right panels detail RSCA and MGMB. Intermediate region and boundary predictions provide deep supervision throughout decoding.

Environment

We recommend Python 3.10 or newer. Install a PyTorch build compatible with your CUDA toolkit before installing the remaining dependencies.

conda create -n c3net python=3.10 -y
conda activate c3net

python -m pip install torch torchvision \
  --index-url https://download.pytorch.org/whl/cu118
python -m pip install -r requirements.txt

GPU training is best performed with the CUDA implementation provided by mamba-ssm. A selective-scan fallback is included for environments where the CUDA extension is unavailable.

The default encoder initialization uses the VMamba-S ImageNet checkpoint. Place it at the following path:

pre_trained_weights/vmamba_small_e238_ema.pth

Data Preparation

C³Net is evaluated on Kvasir-SEG, CVC-ClinicDB, CVC-ColonDB, and ETIS-LaribPolypDB. Organize the datasets as follows and pass data_root to the training or evaluation command.

data_root/
|-- TrainDataset/
|   |-- images/
|   |-- masks/
|   `-- boundary3/
`-- TestDataset/
    |-- CVC-ClinicDB/
    |   |-- images/
    |   `-- masks/
    |-- Kvasir/
    |   |-- images/
    |   `-- masks/
    |-- CVC-ColonDB/
    |   |-- images/
    |   `-- masks/
    `-- ETIS-LaribPolypDB/
        |-- images/
        `-- masks/

Image and mask filenames must match. The training images, masks, and boundary maps must also share the same filenames.

Boundary Generation

Generate the boundary targets used for deep supervision with boundary.py. The default command uses the original 3 x 3 max-pooling boundary operator.

python boundary.py \
  /path/to/data_root/TrainDataset/masks \
  /path/to/data_root/TrainDataset/boundary3

Training

The global batch size must be divisible by the number of processes.

OMP_NUM_THREADS=1 CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
torchrun --standalone --nproc_per_node=8 --master_port=29531 \
  train_c3net.py --config latest \
  --batch_size 40 --epochs 200 --num_workers 0 \
  --data_path /path/to/data_root \
  --work-dir results/c3net_latest

Before a full run, validate the configuration and dataset paths:

python train_c3net.py --config latest --check-only \
  --batch_size 2 --data_path /path/to/data_root

python train_c3net.py --config latest --smoke-test \
  --batch_size 2 --num_workers 0 --data_path /path/to/data_root

Training outputs are written to --work-dir, including logs, TensorBoard summaries, configuration snapshots, the latest checkpoint, and the best validation checkpoints.

Evaluation

The released C³Net checkpoint is available from Google Drive. Model weights are loaded with strict parameter-name and tensor-shape checking.

Evaluate a trained checkpoint on the four benchmark datasets with one GPU:

CUDA_VISIBLE_DEVICES=0 python test_c3net.py \
  --checkpoint /path/to/checkpoint.pth \
  --data-root /path/to/data_root \
  --input-size 256 \
  --output-dir results/test_c3net

For distributed evaluation:

OMP_NUM_THREADS=1 CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
torchrun --standalone --nproc_per_node=8 --master_port=29532 \
  test_c3net.py \
  --checkpoint /path/to/checkpoint.pth \
  --data-root /path/to/data_root \
  --input-size 256 \
  --output-dir results/test_c3net

Useful evaluation options include:

• --check-only: validate dataset paths, model construction, and strict checkpoint loading without inference.
• --max-samples N: evaluate only the first N samples of each dataset.
• --save-predictions: save binary prediction masks at their original sizes.
• --datasets ...: select a subset of the supported test datasets.

Notebook Demo

notebooks/C3Net_demo.ipynb provides a reproducible checkpoint-loading and test-set evaluation workflow. It downloads the released weights from Google Drive, verifies strict model compatibility, runs the official evaluator on selected benchmark datasets, and displays the resulting segmentation metrics.

Repository Structure

C3Net/
|-- assets/                         # README figures
|-- configs/                        # Training presets
|-- datasets/                       # Training data loader
|-- models/c3net/
|   |-- backbone/                   # Encoder and decoder feature stages
|   |-- blocks/                     # MGMB, MKF, VSSB, and DCF modules
|   |-- heads/                      # Segmentation and auxiliary heads
|   `-- layers/                     # Hierarchical encoder layers
|-- notebooks/C3Net_demo.ipynb      # Checkpoint evaluation notebook
|-- utils/                          # Losses and training utilities
|-- boundary.py                     # Boundary target generation
|-- train_c3net.py                  # Training entry point
|-- test_c3net.py                   # Evaluation entry point
`-- requirements.txt

Citation

If you use this implementation in your research, please cite the C³Net manuscript:

@misc{lv2026c3net,
  title  = {Coherent Context Complement Network for Polyp Segmentation},
  author = {Lv, Yongfeng and Zhang, Pingping and Tang, Tongdan and
            Lv, Yanan and Ma, Yili and Lu, Huchuan},
  year   = {2026},
  note   = {Manuscript under review}
}

Please also retain the citations and licenses of the upstream VMamba and VM-UNet projects when redistributing derived code.