Repository files navigation

This repository contains code for a multiple classification image segmentation model based onUNet andUNet++

make sure to put the files as the following structure:

data├── images|   ├── 0a7e06.jpg│   ├── 0aab0a.jpg│   ├── 0b1761.jpg│   ├── ...|└── masks    ├── 0a7e06.png    ├── 0aab0a.png    ├── 0b1761.png    ├── ...

mask is a single-channel category index. For example, your dataset has three categories, mask should be 8-bit images with value 0,1,2 as the categorical value, this image looks black.

You can download the demo dataset fromhere to data/

python train.py

python inference.py -m ./data/checkpoints/epoch_10.pth -i ./data/test/input -o ./data/test/output

If you want to highlight your mask with color, you can

python inference_color.py -m ./data/checkpoints/epoch_10.pth -i ./data/test/input -o ./data/test/output

You can visualize in real time the train and val losses, along with the model predictions with tensorboard:

tensorboard --logdir=runs