Home

Awesome

TensorFlow-ENet

TensorFlow implementation of ENet: A Deep Neural Network Architecture for Real-Time Semantic Segmentation.

This model was tested on the CamVid dataset with street scenes taken from Cambridge, UK. For more information on this dataset, please visit: http://mi.eng.cam.ac.uk/research/projects/VideoRec/CamVid/.

Requirements: TensorFlow >= r1.2

Visualizations

Note that the gifs may be out of sync if the network doesn't load them together. You can refresh your page to see them in sync.

Test Dataset Output

CamVid Test Dataset Output CamVid Test Dataset Output

TensorBoard Visualizations

Execute tensorboard --logdir=log on your root directory to monitor your training and watch your segmentation output form against the ground truth and the original image as you train your model.

Contents

Code

Folders

Important Notes

  1. As the Max Unpooling layer is not officially available from TensorFlow, a manual implementation was used to build the decoder portion of the network. This was based on the implementation suggested in this TensorFlow github issue.

  2. Batch normalization and 2D Spatial Dropout are still retained during testing for good performance.

  3. Class weights are used to tackle the problem of imbalanced classes, as certain classes appear more dominantly than others. More notably, the background class has weight of 0.0, in order to not reward the model for predicting background.

  4. On the labels and colouring scheme: The dataset consists of only 12 labels, with the road-marking class merged with the road class. The last class is the unlabelled class.

  5. No preprocessing is done to the images for ENet. (see references below on clarifications with author).

  6. Once you've fine-tuned to get your best hyperparameters, there's an option to combine the training and validation datasets together. However, if your training dataset is large enough, this won't make a lot of difference.

Implementation and Architectural Changes

  1. Skip connections can be added to connect the corresponding encoder and decoder portions for better performance.

  2. The number of initial blocks and the depth of stage 2 residual bottlenecks are tunable hyperparameters. This allows you to build a deeper network if required, since ENet is rather lightweight.

  3. Fused batch normalization is used over standard batch normalization for faster computations. See TensorFlow's best practices.

  4. To obtain the class weights for computing the weighted loss, Median Frequency Balancing (MFB) is used by default instead of the custom ENet class weighting function. This is due to an observation that MFB gives a slightly better performance than the custom function, at least on my machine. However, the option of using the ENet custom class weights is still possible.

References

  1. ENet: A Deep Neural Network Architecture for Real-Time Semantic Segmentation
  2. Implementation of Max Unpooling
  3. Implementation of PReLU
  4. Clarifications from ENet author
  5. Original Torch implementation of ENet
  6. ResNet paper for clarification on residual bottlenecks
  7. Colouring scheme

Feedback and Bugs

This implementation may not be entirely correct and may contain bugs. It would be great if the open source community can spot any bugs and raise a github issue/submit a pull request to fix those bugs if any!

Citation

If you are using this work for your research, please consider citing:

@misc{kwot_sin_lee_2017_3403269,
  author       = {Kwot Sin Lee},
  title        = {kwotsin/TensorFlow-ENet: DOI},
  month        = jun,
  year         = 2017,
  doi          = {10.5281/zenodo.3403269},
  url          = {https://doi.org/10.5281/zenodo.3403269}
}

DOI