Contact: Konrad Reuter, konrad.reuter@tuhh.de
This repository contains the source code for PolypNextLSTM. The main components of the model are an encoder based on ConvNext-tiny, a bidirectional convolutional LSTM for fusing information across multiple frames and a UNet-like decoder.
Clone the repository and create a data folder:
git clone https://github.com/KonradReuter/PVS.git
cd PVS
mkdir data
Request and prepare the SUN-SEG dataset as explained here and move it to the created data folder.
For Windows users: The script to reorganize the file structure from the linked repository above will not work. Instead, use our slightly modified script. After that you should manually remove the "SUN-Positive" folder and rename "SUN-SEG-Annotation" into "SUN-SEG".
Python > 3.9 is recommended.
python3 -m venv venv
source venv/bin/activate || ./venv/Scripts/activate
python3 -m pip install pip setuptools wheel
python3 -m pip install -e .
The main script offers functions for training and evaluating a model. The general usage is as follows.
python scripts/main.py train --model-name=<model_name> --file-name=<file_name> --run-name=<run_name>
python scripts/main.py evaluate --model-name=<model_name> --file-name=<file_name> --run-name=<run_name>
Here the model-name defines which model is trained/evaluated. Our proposed model from above can be selected under the name "Conv_LSTM". Several other models used in our experiments can be used. A full list is given here.
The file-name defines the name of the weight file the model is saved to / loaded from.
The run-name defines the name under which the run is logged in WandB.
To configure the hyperparameters, please edit the args.json file. A more detailed explanation of the different hyperparameters is given here.
During our experiments we used shell scripts to chain multiple model trainings. There are four scripts for training and evaluating models:
- train_models.sh: Trains and evaluates all constructed models using 5-fold cross validation.
- train_num_frames.sh: Trains and evaluates our PolypNextLSTM model with varying number of input frames using 5-fold cross validation.
- train_single.sh: Trains a model using a single direction ConvLSTM instead of a bidirectional one using 5-fold cross validation.
- train_sota.sh: Trains and evaluates 10 state-of-the-art models using 5-fold cross validation.
The code for the Hybird2D/3D model is not publicly available. We want to thank the authors for providing it to us. If you want to repeat our experiments, please send them a request as well.
The model will need some adjustments to run in our environment. For more information, contact us once you received the original model code.
The PNSNet and PNSPlusNet are based on the NS-Block module. To be able to use it, a CUDA extension has to be installed. CUDA 11.7 is required.
cd scripts/models/PNSPlus
python setup.py build develop
CASCADE, COSNet and TransFuse use pretrained backbone weights, which are not downloaded automatically. Please download the weights from the following links and move them into the according model folder.
- PVTv2-B2 (CASCADE)
- DeepLabV3 (COSNet)
- ResNet34 (TransFuse)
- DeiT-small (TransFuse)
Pretrained weights for PolypNextLSTM can be downloaded from here. The number at the end of the file name refers to the group of clips that was used for validation during 5-fold cross validation.
Weights for all other models trained in our experiments can be provided on request.
This repository includes code from various other repository. Many thanks to the authors for providing their code: