Deep Learning for Semantic Segmentation in RoboCup
This set of scripts allows you to train deep convolutional networks for semantic segmentation in Pytorch as described in our paper. This code is finetuned to work well with the networks and the problem (robot soccer) we discussed in the paper, still it should work reasonably fine for other problems as well.
This code is written in python 3.6. The following python packages are required:
- Pytorch
- Torchvision
- Progressbar 2
- PIL Image library
- Numpy
- Visdom
You can install all of these by entering pip install <package name> into the console.
In order to run the code, you also need our datasets, which you can download here. Extract it to the data subfolder. (NOTE: If after extraction you have "data/data/many folders", then you need to move the whole thing one folder up to get rid of one of the "data"-s.)
You can download the datasets used for training and finetuning from our website
If you do not want to train your own models, just use our pretrained models, which you can download from here. Extract them into the "checkpoints" folder.
The pretrain the whole segmentation network on synthetic data.
python train.py
Then finetune the network on the real dataset, (you can use synthetic transfer learning optionally).
python train.py –-finetune
python train.py –-finetune --transfer
The finetune flag also determines if the synthetic or the real test sets are evaluated.
There are several further options you can use for training:
- --noScale This option will run the network on 640x480 resolution instead of 160x120
- --UNet This option will train the standard U-Net architecture, instead of the ROBO-UNet proposed in our paper.
- --topCam, --bottomCam Use top or bottom camera images only. Providing both or neither flags will use both cameras.
- --noBall, --noRobot, --noGoal, --noLine Treat balls, robot, goals or lines as background. You can use any combinations of these flags, as long as there is at least one foreground class.
Finally, you can generate the labeled images using the tester script. You can use the flags to select different neural network versions. Examples:
python test.py
python test.py –-finetune
python tester.py –-finetune –-transfer
python tester.py –-finetune --v2 --topCam
python tester.py –-finetune --v2 --bottomCam --noGoal --noRobot
There are several further options you can use for training:
The training of the label propagation network is quite similar to the segmentation networks, except that the classification pretraining step is omitted.
python labelPropTrain.py
python labelPropTrain.py -–finetune
python labelPropTrain.py –-finetune –-prune
You can also generate the labeled images using the following script:
python validLabelProp.py –-finetuned –-pruned
You can also compare the results with farneback optical flow using:
python validLabelProp.py –-finetuned –-optFlow
Note, that the 'finetuned' flag is needed to set the database to the real one.