Pixel-wise segmentation on the VOC2012 dataset using pytorch.
See dataset examples here.
Download image archive and extract and do:
mkdir data
mv VOCdevkit/VOC2012/JPEGImages data/images
mv VOCdevkit/VOC2012/SegmentationClass data/classes
rm -rf VOCdevkit
We recommend using pyenv:
pyenv virtualenv 3.6.0 piwise
pyenv activate piwise
then install requirements with pip install -r requirements.txt
.
For latest documentation use:
python main.py --help
Supported model parameters are fcn8
, fcn16
, fcn32
, unet
, segnet1
,
segnet2
, pspnet
.
If you want to have visualization open an extra tab with:
python -m visdom.server -port 5000
Train the SegNet model 30 epochs with cuda support, visualization and checkpoints every 100 steps:
python main.py --cuda --model segnet2 train --datadir data \
--num-epochs 30 --num-workers 4 --batch-size 4 \
--steps-plot 50 --steps-save 100
Then we want to do semantic segmentation on foo.jpg
:
python main.py --model segnet2 --state segnet2-30-0 eval foo.jpg foo.png
The segmented class image can now be found at foo.png
.