An implementation of U-Net using TensorFlow framework, with multi-class support.
We design the project to find a balance between encapsulation and the ability of code modification, please follow the steps below to train on your custom dataset.
- Annotate your data using tools such as LabelMe
- Use the following methods in
utilities.pyto obtain image and mask data for model input fromjsonlabels:labelme_json_to_mask()generate_sequence_data()data_preprocessing()
- Follow the pipeline in
train.py - Evaluate model performance using
evaluate.py
Please don't forget to modify n_class in train.py and evaluate.py according to your needs.
Please refer to the docstring of aforementioned methods for more info.
A basic inference pipeline is implemented in inference.py, with a wrapper method also provided, please refer to in-code docs for more info.
The following results are from tiles of an underground parking lot IPM image. The model was trained for detecting road lanes and markers, please refer to best.h5.
The model achieved a Mean_IoU score of 86.25% on our current version of dataset.
Segmentation sample on a full image:
Software:
- Linux, Windows >= 10, macOS (No CUDA support) >= 10.14 operating systems
- Python packages required in
requirements.txt - Python version >= 3.6
Hardware:
- Multicore CPU & CUDA-enabled GPU strongly recommended
- Prefer RAM >= 32 GB, VRAM >= 10 GB
- Prefer SSD for faster data loading
The original U-Net architecture was proposed by Olaf Ronneberger, Philipp Fischer, Thomas Brox in 2015, please refer to their paper.
This repo also uses code from bnsreenu's GitHub repository.
Sincerely thanks to all the people mentioned above.
The code hasn't gone through thorough test, and bugs are almost bound to exist, please contact the author or open an issue if you encountered one.
Also, please notice that this repo is primarily developed for research purposes, with no guarantee on model performance, efficiency, etc. Use at your own risk.