Graded Image Classification is a multi-task learning approach for improving image classification performance when working with graded data, such as the stage of cancer.
In this method, a Vision Transformer (ViT) network is trained on graded training data, learning both image classification and comparison of grades.
At the test phase, the approach utilizes a voting method based on comparison information to make the final prediction, resulting in improved performance.
The approach has been tested and shows improved performance compared to traditional image classification methods using Convolutional Neural Networks (CNNs).
Install dependencies
# clone project
git clone https://github.com/OZOOOOOH/colon_compare
cd colon_compare
# [OPTIONAL] create conda environment
conda create -n myenv python=3.8
conda activate myenv
# install pytorch according to instructions
# https://pytorch.org/get-started/
# install requirements
pip install -r requirements.txt
Train model with default configuration
# train on CPU
python train.py trainer.gpus=0
# train on GPU
python train.py trainer.gpus=1
Train model with chosen experiment configuration from configs/experiment/
python train.py experiment=experiment_name.yaml
You can override any parameter from command line like this
python train.py trainer.max_epochs=20 datamodule.batch_size=64