This repository is a PyTorch implementation of ENet: A Deep Neural Network Architecture for Real-Time Semantic Segmentation. The partial codes are based on
- The official Torch code by the authors of the paper;
- The unofficial Caffe code by TimoSaemann;
- The unofficial PyTorch code by vietdoan;
- pytorch-semantic-segmentation by ZijunDeng.
- PyTorch
- TensorboardX
mean IoU | Model Size | |
---|---|---|
encoder | 53.66% | 1.38M |
encoder + decoder (step-by-step training) | 56.04% | 1.49M |
The best mean IoU of 56.04% on the val set is close to the 58.3% in the paper (58.3% is reported on the test set)
exemplars | |
---|---|
encoder | ![]() |
encoder + decoder (step-by-step training) | ![]() |
- Downloading the CityScapes dataset.
- Unzip the
leftImg8bit_trainvaltest.zip
and thegtFine_trainvaltest.zip
in a certain folder (root path
). - Preprocess the dataset:
1). modify the
root path
and theprocessed the path
in./script/preprocess.py
; 2). runpython ./script/preprocess.py
. - Calculate the label weights:
python ./script/cal_label_weighting_Enet.py --trainDataPath=your_label_path --num_classes=19
, and repalce the__C.TRAIN.LABEL_WEIGHT
inconfig.py
- modify config.py:
__C.TRAIN.STAGE = 'encoder' # encoder or all __C.TRAIN.PRETRAINED_ENCODER = '' #
python train.py
The loss curve line:- Here, we provide the trained Pytorch model.
- modify config.py:
__C.TRAIN.STAGE = 'all' # encoder or all __C.TRAIN.PRETRAINED_ENCODER = './ckpt/encoder_ep_497_mIoU_0.5098.pth' # Path of the pretrained encoder
python train.py
The loss curve line:- Here, we provide the trained Pytorch model.