PyTorch (v1.0.0) re-implementation of ENet: A Deep Neural Network Architecture for Real-Time Semantic Segmentation, ported from the excellent PyTorch impelentation MapLiteSeg, which was in-turn ported from the lua-torch implementation ENet-training created by the authors.
This implementation has been tailored to suit MapLite LiDAR data.
On an NVIDIA GeForce GTX 1060, ENet
operates at a rate of about 80 Hz! Should run faster on a better GPU.
- Python 3 and pip.
- Set up a virtual environment (recommended).
- Install dependencies using pip:
pip install -r requirements.txt
.
Run main.py
, the main script file used for training and/or testing the model. The code has a lot of options. Make sure to read through most of them, before training/testing.
usage: main.py [-h] [--mode {train,test,inference,full}] [--resume]
[--generate-images] [--arch {rgb,rgbd}]
[--seg-classes {nyu40,scannet20}] [--batch-size BATCH_SIZE]
[--epochs EPOCHS] [--learning-rate LEARNING_RATE]
[--beta0 BETA0] [--beta1 BETA1] [--lr-decay LR_DECAY]
[--lr-decay-epochs LR_DECAY_EPOCHS]
[--weight-decay WEIGHT_DECAY] [--dataset {scannet}]
[--dataset-dir DATASET_DIR] [--trainFile TRAINFILE]
[--valFile VALFILE] [--testFile TESTFILE] [--height HEIGHT]
[--width WIDTH] [--weighing {enet,mfb,none}]
[--class-weights-file CLASS_WEIGHTS_FILE] [--with-unlabeled]
[--workers WORKERS] [--print-step PRINT_STEP] [--imshow-batch]
[--device DEVICE] [--name NAME] [--save-dir SAVE_DIR]
[--validate-every VALIDATE_EVERY]
For help on the optional arguments run: python main.py -h
python main.py -b 4 --epochs 50 --dataset-dir /path/to/prepared/data --trainFile cache/train.txt --valFile cache/test.txt --testFile cache/test.txt --class-weights-file cache/class_weights_maplite.txt --print-step 25 --name ENet --lr-decay-epochs 60 -lr 1e-3 --beta0 0.7 --validate-every 3
Training for 10 epochs on 500 scans should take about 5-6 minutes on an NVIDIA GeForce GTX 1060.
python main.py -b 4 --epochs 50 --dataset-dir /path/to/prepared/data --trainFile cache/train.txt --valFile cache/test.txt --testFile cache/test.txt --class-weights-file cache/class_weights_maplite.txt --print-step 25 --name ENet --lr-decay-epochs 60 -lr 1e-3 --beta0 0.7 --validate-every 3 --resume
Once you're all trained and set, you can use inference.py
to generate the cool-looking qualitative results on top of this README
. A sample inference.py
call would look like
python inference.py --mode inference -b 4 --epochs 1 --dataset-dir /path/to/prepared/data --trainFile cache/train.txt --valFile cache/val.txt --testFile cache/test.txt --print-step 1 --class-weights-file cache/class_weights_maplite.txt --name ENet --generate-images
This will create a directory named ENetDepth_images
in the save
directory.
data
: Contains instructions on how to download the datasets and the code that handles data loading.metric
: Evaluation-related metrics.models
: ENet model definition.save
: By default,main.py
will save models in this folder. The pre-trained models can also be found here.
args.py
: Contains all command-line options.main.py
: Main script file used for training and/or testing the model.test.py
: Defines theTest
class which is responsible for testing the model.train.py
: Defines theTrain
class which is responsible for training the model.transforms.py
: Defines image transformations to convert an RGB image encoding classes to atorch.LongTensor
and vice versa.