This repository provides a pipeline for lane segmentation using a UNet model trained on the BDD100K dataset. It includes modular scripts for training, validation, and inference, with configuration-driven workflows and TensorBoard logging. The project supports exporting PyTorch models for deployment with the NCNN framework, enabling efficient inference on ARM and x86 platforms via C++.
- UNet architecture for lane segmentation
- BDD100K dataset integration
- Configurable training and inference scripts
- Checkpointing and metrics
- TensorBoard support
- Model export for NCNN (TorchScript → pnnx → NCNN)
- C++ deployment example with NCNN and OpenCV
- Overview
- Setup Instructions
- Dataset
- Pretrained Models on Hugging Face
- Training Usage
- Inference Usage
- Exporting Model for NCNN Deployment
- NCNN Library Setup
- NCNN C++ Deployment
- Gradio Lane Detection Demo
- Model Architecture Comparison
- Results Comparison
- Codebase Directory Tree
- Citation
docker run -it --gpus all --shm-size=16g -v /home:/home <image_id>apt update && apt upgrade -y
apt install python3 python3-pip -y
apt install build-essential git cmake wget libprotobuf-dev protobuf-compiler libomp-dev libopencv-dev -ypip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117
pip install -r requirements.txtBDD100K is a large-scale driving video dataset for diverse road scenes, widely used for autonomous driving research. This project uses the lane segmentation subset, which provides pixel-level lane annotations for training semantic segmentation models.
- Official website: https://bdd-data.berkeley.edu/
- Download the images and lane mask annotations from the BDD100K website or via their download scripts.
- Organize the dataset as follows:
bdd100k/
├── images/
│ └── 100k/
│ ├── train/
│ ├── val/
│ └── test/
└── lane/
├── colormaps/
├── masks/
│ ├── train/
│ └── val/
└── polygons/
images/100k/train,val, andtestcontain the raw images.lane/masks/trainandvalcontain binary lane segmentation masks.lane/colormapsandlane/polygonsprovide additional annotation formats.
Pretrained weights for all UNet variants are available on Hugging Face:
| Model | PyTorch (.pt) | NCNN (.param, .bin) | Repo Link |
|---|---|---|---|
| UNet | pt | param, bin | repo |
| UNetDepthwise | pt | param, bin | repo |
| UNetDepthwiseSmall | pt | param, bin | repo |
| UNetDepthwiseNano | pt | param, bin | repo |
Each folder contains both PyTorch (.pt) and NCNN (.param, .bin) files for direct use in Python or C++ inference workflows.
- Update
config/train_config.yamlandconfig/test_config.yamlto point to your local BDD100K image and mask directories. - The dataset loader expects images and masks to be sorted and matched by filename.
Edit config/train_config.yaml to set dataset paths, batch size, epochs, learning rate, image size, and model parameters.
python3 scripts/train.pyCheckpoints and logs will be saved in the checkpoints/ and runs/ folders.
TensorBoard logs are available in runs/. To view training progress:
tensorboard --logdir runs/Edit config/test_config.yaml to set the trained checkpoint, input images directory, output directory, and model parameters.
python3 scripts/test.pyOutput masks will be saved in the specified output directory. Progress is shown with tqdm.
Run the export script to convert your trained UNet model to TorchScript and directly export NCNN .param and .bin files:
python export/export_to_ncnn.pyThis will generate unet_jit.pt, unet_jit.param, and unet_jit.bin in the output directory specified in your config. No manual pnnx command is needed.
Use the generated .param and .bin files for the later NCNN C++ deployment.
Refer to NCNN documentation for integration details.
To build NCNN from source:
git clone https://github.com/Tencent/ncnn.git
cd ncnn
mkdir build && cd build
cmake ..
make -j$(nproc)
make installGo to the deployment folder:
cd ncnn_deploy
mkdir build && cd build
cmake ..
make -j$(nproc)If you see errors about missing NCNN or OpenCV, make sure to set the correct paths in CMakeLists.txt:
set(ncnn_DIR "/path/to/ncnn/install/lib/cmake/ncnn")
include_directories("/path/to/ncnn/install/include")You can now run the C++ inference with flexible command-line options:
./unet_ncnn key=value ...Supported keys (all optional, defaults shown):
input_image=/path/to/image.jpg
model_param=../../checkpoints/exp_20250907_172056/ncnn_models/unet_jit.ncnn.param
model_bin=../../checkpoints/exp_20250907_172056/ncnn_models/unet_jit.ncnn.bin
output_mask=mask.jpg
output_overlay=overlay.jpg
input_w=256
input_h=256
input_layer=in0
output_layer=out0
mask_alpha=0.8
Example:
./unet_ncnn input_image=img.jpg mask_alpha=0.5 output_mask=mask.png output_overlay=overlay.pngThis will save the output mask and overlay images to the specified paths.
output_mask(default: mask.jpg): Lane mask imageoutput_overlay(default: overlay.jpg): Lane overlay visualization
- If CMake cannot find protobuf, try:
cmake .. -DProtobuf_INCLUDE_DIR=/usr/include -DProtobuf_LIBRARIES=/usr/lib/x86_64-linux-gnu/libprotobuf.so
- If CMake cannot find NCNN, set
ncnn_DIRto the folder containingncnnConfig.cmake. - For ARM cross-compilation, set toolchain and paths as needed.
You can try the lane detection models interactively in your browser using Gradio:
- Hugging Face Space: Lane Detection UNet Demo
To run the Gradio app locally:
pip install -r requirements.txt
python3 gradio_app/app.py- Upload a road image and select a model type (UNet, UNetDepthwise, UNetDepthwiseSmall, UNetDepthwiseNano).
- The app will display the lane detection overlay result.
| Model | Parameters | FLOPs (GFLOPs) | Estimated Size (MB) |
|---|---|---|---|
| UNet | 31,043,521 | 54.79 | 700.10 |
| UNetDepthwise | 4,216,799 | 11.57 | 480.08 |
| UNetDepthwiseSmall | 253,919 | 5.88 | 385.58 |
| UNetDepthwiseNano | 52,191 | 3.00 | 279.92 |
- All models use input size (1, 3, 256, 256).
- FLOPs measured with ptflops (GMac).
- Estimated size includes input, activations, and parameters.
america_highway_x86_unet_comparison.mp4
hsinchu_to_taipei_1k_x86_unet_comparison.mp4
taichung_city_x86_unet_comparison.mp4
lane-detection-unet-ncnn/
├── assets/
│ ├── demo.jpg
│ └──gradio_app.jpg
├── checkpoints
│ ├── exp_20250906_223222
│ │ ├── bdd100k_val_outputs
│ │ │ ├── config
│ │ │ └── inference_results
│ │ ├── config
│ │ │ └── train_config.yaml
│ │ ├── ncnn_models
│ │ │ ├── unet_jit.ncnn.bin
│ │ │ └── unet_jit.ncnn.param
│ │ └── weights
│ │ └── unet_best.pt
│ ├── exp_20250907_094745
| └── ...
├── config/
│ ├── eval_config.yaml
│ ├── export_config.yaml
│ ├── test_config.yaml
│ └── train_config.yaml
├── datasets/
│ └── bdd100k.py
├── export/
│ └── export_to_ncnn.py
├── gradio_app/
│ └── app.py
├── models/
│ ├── unet.py
│ ├── unet_depthwise.py
│ ├── unet_depthwise_small.py
│ └── unet_depthwise_nano.py
├── ncnn/
├── ncnn_deploy/
│ ├── build/
| | ├── input.jpg
│ | └── unet_ncnn
│ ├── CMakeLists.txt
│ └── main.cpp
├── runs/
├── scripts/
│ ├── test.py
│ └── train.py
├── utils/
│ ├── checkpoint.py
│ └── metrics.py
├── .gitignore
├── CITATION.cff
├── LICENSE
├── README.md
└── requirements.txt
If you use this repository or models in your research, please cite as follows:
BibTeX:
@software{Pai_UNet_Lane_Segmentation_2025,
author = {Pai, Nick},
license = {MIT},
month = sep,
title = {{UNet Lane Segmentation \& NCNN Deployment}},
url = {https://github.com/nick8592/lane-detection-unet-ncnn},
version = {1.0.0},
year = {2025}
}
