ENet.PyTorch

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.

Requirement

• PyTorch
• TensorboardX

Performance

Quantitative Redults

	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)

Qualitative Redults

	exemplars
encoder
encoder + decoder (step-by-step training)

Usage

Data Preparation

1. Downloading the CityScapes dataset.
2. Unzip the leftImg8bit_trainvaltest.zip and the gtFine_trainvaltest.zip in a certain folder (root path).
3. Preprocess the dataset: 1). modify the root path and the processed the path in ./script/preprocess.py; 2). run python ./script/preprocess.py.
4. 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 in config.py

Train model

Train the encoder

1. modify config.py:

   __C.TRAIN.STAGE = 'encoder' # encoder or all
   __C.TRAIN.PRETRAINED_ENCODER = '' # 
   

2. python train.py The loss curve line:
3. Here, we provide the trained Pytorch model.

Train the ENet based on the encoder

1. 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
   

2. python train.py The loss curve line:
3. Here, we provide the trained Pytorch model.