PytorchAutoDrive: 自动驾驶感知的开源框架

pytorch-auto-drive

PytorchAutoDrive简介

PytorchAutoDrive是一个纯Python实现的自动驾驶感知框架,基于PyTorch深度学习框架开发。该项目由Zhengyang Feng和Shaohua Guo主要维护,旨在为自动驾驶领域的研究人员和开发者提供一个功能全面、易于使用的工具集。

主要特点

PytorchAutoDrive具有以下几个突出特点:

支持多种语义分割和车道线检测模型
基于配置文件的实现,易于修改和扩展
代码模块化程度高,易于理解和二次开发
提供数据加载、图像处理、可视化等辅助功能
支持混合精度训练,提高训练效率
集成TensorBoard日志记录
支持ONNX和TensorRT模型部署

这些特性使PytorchAutoDrive成为一个功能完备、易用性强的自动驾驶感知框架。无论是进行学术研究还是实际应用开发,都可以基于该框架快速构建和验证自己的想法。

支持的数据集

PytorchAutoDrive支持多个主流的自动驾驶数据集,包括:

语义分割数据集

PASCAL VOC 2012
Cityscapes
GTAV
SYNTHIA

其中GTAV和SYNTHIA主要用于无监督域适应(UDA)的基线设置,使用Cityscapes验证集作为验证。

车道线检测数据集

CULane
TuSimple
LLAMAS
BDD100K (进行中)

这些数据集涵盖了不同的场景和标注方式,为开发鲁棒的车道线检测算法提供了丰富的训练和测试数据。

CULane数据集示例

支持的模型

PytorchAutoDrive实现了多种经典和最新的语义分割及车道线检测模型:

语义分割模型

FCN (ResNet-101)
DeeplabV2 (ResNet-101)
DeeplabV3 (ResNet-101)
ENet
ERFNet

车道线检测模型

Baseline (支持多种骨干网络)
SCNN
RESA
SAD (开发中)
PRNet (开发中)
LSTR
LaneATT
BézierLaneNet

这些模型涵盖了从轻量级到大型网络的多种架构,可以根据具体需求选择合适的模型。同时,框架的模块化设计也使得添加新模型变得简单。

使用指南

安装

克隆代码仓库:

git clone https://github.com/voldemortX/pytorch-auto-drive.git
cd pytorch-auto-drive

安装依赖:

pip install -r requirements.txt

按照 DATASET.md 的说明准备数据集

训练模型

以车道线检测为例,训练模型的基本命令如下:

python main_landet.py --train \
                      --config=<config file path> \
                      --mixed-precision

其中--config指定配置文件路径,--mixed-precision启用混合精度训练。

测试模型

测试模型的基本命令:

python main_landet.py --test \
                      --config=<config file path> \
                      --mixed-precision

可视化

PytorchAutoDrive提供了丰富的可视化工具,可以对图像和视频进行推理并可视化结果。详细使用方法请参考 VISUALIZATION.md。

车道线检测可视化示例

模型部署

PytorchAutoDrive支持将训练好的模型导出为ONNX格式,并提供了TensorRT加速的支持。这使得模型可以更容易地部署到实际的自动驾驶系统中。详细的部署指南可以参考 DEPLOY.md。

贡献指南

PytorchAutoDrive是一个开源项目,欢迎社区贡献。如果您想为项目做出贡献,请参考 CONTRIBUTING.md 了解具体的贡献流程和规范。

结语

PytorchAutoDrive为自动驾驶感知任务提供了一个全面而强大的开发框架。无论您是研究人员还是开发者,都可以利用这个工具来加速自己的工作。我们期待看到更多基于PytorchAutoDrive的创新成果,推动自动驾驶技术的进步。

如果您在使用过程中有任何问题或建议,欢迎在GitHub仓库中提出issue或pull request。让我们共同努力,为自动驾驶的未来贡献力量! 🚗💨