Image segmentation using pytorch.
Make sure to checkout with —recurse-submodules, there’s a submodule ‘tools’ for common bits and pieces)
It loads a folder full of images and masks (pixel value corresponds to class) with a config file.
An example dataset to train on can be found at:
https://drive.google.com/file/d/0B_mOCEqr7usZa2hra0xZUTJscE0/view
Or classes/images from the COCO/Pascal VOC dataset can be imported from scripts in the import/ folder.
python -m dataset.view --input /path/to/dataset --train (or --test)
Useful to check the preprocessing of images.
python view_labels.py some/file.jpg.mask
python main.py --lr 0.1 --batch_size 4 --input /path/to/dataset --model "unet --depth 5" --epoch_size 1024
--load
, load from a previous checkpoint and cointunue training
--model
, specify model and model parameters (use quotes)
--show
, show results of evaluating the model in training (sanity check)
python test.py --batch /path/to/images --model log/model.pth --save results_path
(and/or --show
)
python test.py --image /my/image.jpg --model log/model.pth --show