Implementation of Segformer in Keras
This repository implements Segformer, introduced in the paper: SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers using Keras. The original implementation is in Pytorch and is available on the open-mmlab/mmsegmentation repo. The implementation is heavily guided by Implementing SegFormer in PyTorch where the author goes thru the implementation step by and step and in great detail.
Pretrained weights is only available for:
- People segmentation (trained on 64k images from COCO on only a single class, person)
- Clone this repository
- Install the required packages using the requirements.txt.
- Accepted ground truth is in the format of a mask png, where the pixel values corresponds to the class. The size of the mask will have to be W/4;H/4, where W and H are the width and the height of the input shape respectively. (A helper notebook to convert COCO anns to mask is provided)
- The model, dataset and training parameters are controlled by the config file
In the root folder, run
python train.py
- Tensorboard logs will be saved in the logs directory.
- Model weights will be saved in the weights directory
Only validation accuracy is being evaluated in this repo. mIOU is not implemented yet.
A helper notebook is provided for visualization of the results.
Segmentation annotations in COCO have overlapping annotations which causes some issues in some images when deciding which classes should preside over which classes. Nevertheless, I have tried training the model on the full 90 classes with COCO but due to computational limitations, I have stopped at 10 epochs. The model performance leaves much to be desired. If anyone could suggest anything (optimizer, LR, etc.) to improve training please drop me a message.
I have requested for cityscapes and ADE20k datasets but still pending... (as of 3 Dec 2022)
The code is published under the Apache License 2.0.