milesial/Pytorch-UNet

Always [0,1] mask on mutil-classes models

g-mero opened this issue · 2 comments

g-mero commented

expect:
124

but result:
res

train script
python train.py -e 5 -b 10 -s 1 -v 10 -c 4 --amp

predict script
python predict.py -m checkpoints/checkpoint_epoch4.pth -s 1 -c 4 -i data/imgs/124.jpg -o res.png

g-mero commented

cause by data under-fitting. improve dataset and increase epochs fix this

I have the same problem! I changed the color mapping to fit multiple classes but I only get 2 classes :(