This repository contains PyTorch code train segmentation models on chest-xray images.
View different models, results, and sample images here.
.
├── data
├── train_seg.py
└── utils
├── load_data.py
├── load_extra.py
├── load_loss.py
├── load_model.py
└── load_train_loop.py
Dataloaders for following datasets are available. Images should be saved in ./data
.
- Montgomery
- Shenzhen
- JSRT
Segmentation can be run on the following configurations.
- Full: Runs on complete dataset.
- Normal: Runs only on normal chest x-rays.
- Abnormal: Runs only on abnormal chest x-rays.
The following segmentation models are available. Please note that these are custom implementations and don't exeactly match the actual papers.
- FCN8
- DeconvNet
python train_seg.py --dataset [str] --subset [str] --model_name [str] --batch_size [int] --epochs [int] --learning_rate [float]
python train_seg.py --dataset "montgomery" --subset "full" --model_name "custom_model_name" --batch_size 16 --epochs 100 --learning_rate 0.01