This repository contains code for a multiple classification image segmentation model based on UNet and UNet++

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 from here 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