手把手教你利用算法工具链训练、量化、编译、可视化 征程 6 参考算法 BEVFormer

头像
地平线智驾开发者
     广东
    阅读 18 分钟
    头图

    1 环境部署

    1.1 开发机准备

    为了顺利地使用工具链，地平线建议您选择的开发机应满足以下要求：

    硬件/操作系统要求 CPUCPU I3 以上或者同级别 E3/E5 的处理器内存 16G 或以上级别 GPUCUDA11.8、驱动版本 Linux:>= 510.39.01*

    （推荐驱动版本 Linux:520.61.05）

    适配显卡包括但不限于：

    1. GeForce RTX 3090
    2. GeForce RTX 2080 Ti
    3. NVIDIA TITAN V
    4. Tesla V100S-PCIE-32GB
    5. A100 系统原生Ubuntu 22.04

    1.2 Docker 容器部署

    地平线要求的 Docker 基础环境如下，请提前在您的宿主机上完成安装：

    Docker（20.10.10或 更高版本，建议安装 20.10.10版本），详见 Docker安装手册 。
    NVIDIA Container Toolkit（1.13.5 或更高版本，建议安装 1.15.0），详见 NVIDIA Container Toolkit安装手册。
    完成 Docker 基础环境安装后，还需要将无 root 权限的用户添加到 Docker 用户组中。参考命令如下：

    sudo groupadd docker
sudo gpasswd -a ${USER} docker
sudo service docker restart

    拉取 docker 镜像，并下载 OE 开发包，将 OE 开发包和 Nuscenes 数据集挂载到 docker 容器中。参考命令如下：

    # GPU Docker
docker pull openexplorer/ai_toolchain_ubuntu_20_j6_gpu:{version}
# 手动启动 GPU Docker 镜像
docker run -it --rm
  --network host \ # 调整网络模式为host
  --gpus all \ # 在启动容器时，添加标记以启用GPU资源的访问
  --shm-size=15g \ # 修改共享内存大小，切记修改不然默认64MB
  -v {OE 包路径}:/open_explorer \ # 挂载 OE 包
  -v {数据集路径}:/data/horizon_j6/data \ # 挂载数据集
  openexplorer/ai_toolchain_ubuntu_20_j6_gpu:{version}

    注意，“--gpus”参数一定要设置为 all，代表着在容器内能使用开发机所有的 GPU，“--shm-size”一定要修改，不然默认共享内存仅为 64MB。

    2 数据集准备

    2.1 数据集下载

    进入 nuscenes 官网，根据提示完成账户的注册，下载 Full dataset(v1.0)、CAN bus expansion 和 Map expansion（v1.3）、nuScenes-lidarseg 这四个项目下的文件。下载后的压缩文件为：

    |-- nuScenes-map-expansion-v1.3.zip
|-- nuScenes-lidarseg-all-v1.0.tar.bz2
|-- can_bus.zip
|-- v1.0-mini.tar
|-- v1.0-trainval01_blobs.tar
|-- ...
|-- v1.0-trainval10_blobs.tar
`-- v1.0-trainval_meta.tar

    Full dataset(v1.0)包含多个子数据集，如果不需要进行 v1.0-trainval 数据集的浮点训练和精度验证，可以只下载 v1.0-mini 数据集进行小场景的训练和验证。

    2.2 Full dataset 版本数据集打包和 meta 文件夹构建

    2.2.1 数据集打包

    将下载完成的 v1.0-trainval01_blobs.tar~v1.0-trainval10_blobs.tar、nuScenes-lidarseg-all-v1.0.tar.bz2、v1.0-trainval_meta.tar、nuScenes-map-expansion-v1.3.zip 和can_bus.zip 进行解压，解压后的目录如下所示：

    |--nuscenes
    |-- can_bus #can_bus.zip解压后的目录
    |-- lidarseg #nuScenes-lidarseg-all-v1.0.tar.bz2解压后的目录
    |-- maps #nuScenes-map-expansion-v1.3.zip解压后的数据
    |-- samples #v1.0-trainvalXX_blobs.tar解压后的目录
    |   |-- CAM_BACK
    |   |-- ...
    |   |-- CAM_FRONT_RIGHT
    |   |--  ...
    |   `-- RADAR_FRONT_RIGHT
    |-- sweeps
    |   |-- CAM_BACK
    |   |-- ...
    |   |-- CAM_FRONT_RIGHT
    |   |--  ...
    |   `-- RADAR_FRONT_RIGHT
    |-- v1.0-trainval #v1.0-trainval_meta.tar解压后的数据
        |-- attribute.json
        |    ...
        `-- visibility.json

    进入 horizon_model_train_sample/scripts 目录，使用以下命令将训练数据集和验证数据集打包，格式为lmdb。如果将上述nuscenes文件夹放到上述目录内，下面WORKSPACE可替换为pwd。

    #pack train_Set
python3 tools/datasets/nuscenes_packer.py --src-data-dir /WORKSPACE/nuscenes/ --pack-type lmdb --target-data-dir /WORKSPACE/tmp_data/nuscenes/v1.0-trainval --version v1.0-trainval --split-name train
#pack val_Set
python3 tools/datasets/nuscenes_packer.py --src-data-dir /WORKSPACE/nuscenes/ --pack-type lmdb --target-data-dir /WORKSPACE/tmp_data/nuscenes/v1.0-trainval --version v1.0-trainval --split-name val

    打包结束生成目录如下所示。

    |--scripts
    |-- tmp_data
        |-- nuscenes  
            |-- v1.0-trainval 
                |-- train_lmdb
                |   |--data.mdb
                |   |--lock.mdb    
                `-- train_lmdb
                    |--data.mdb
                    `--lock.mdb

    --src-data-dir 为解压后的 nuscenes 数据集目录；
    --target-data-dir 为打包后数据集的存储目录；
    --version 选项为["v1.0-trainval", "v1.0-test", "v1.0-mini"]，如果进行全量训练和验证设置为v1.0-trainval，如果仅想了解模型的训练和验证过程，则可以使用 v1.0-mini 数据集；v1.0-test数据集仅为测试场景，未提供注释。
    全量的 nuscenes 数据集较大，打包时间巨长，如果仅仅作为学习使用mini版就好。

    2.2.2 meta 文件夹构建

    在 tmp_data/nuscenes 下创建 meta 文件夹，如果使用--version = "v1.0-mini",将解压后的 v1.0-mini 文件夹拷贝到 tmp_data/nuscenes/meta 文件夹内，如果使用--version = "v1.0-trainval",将解压后的 v1.0-trainval 文件夹拷贝到 tmp_data/nuscenes/meta 文件夹内；
    将解压后的 maps 文件夹拷贝到 tmp_data/nuscenes/meta 文件夹内；
    将解压后的 lidarseg 文件夹拷贝到 tmp_data/nuscenes/meta 文件夹内。
    如果使用--version = "v1.0-trainval"，此时 tmp_data 的目录结构为：

    |-- tmp_data 
|   |-- nuscenes 
|   |   |-- meta
|   |   |   |-- maps        
|   |   |   |   |-- 36092f0b03a857c6a3403e25b4b7aab3.png
|   |   |   |   |-- ...
|   |   |   |   |-- 93406b464a165eaba6d9de76ca09f5da.png
|   |   |   |   |-- prediction
|   |   |   |   |-- basemap
|   |   |   |   |-- expansion
|   |   |   |-- lidarseg     
|   |   |   |   |-- v1.0-mini
|   |   |   |   |-- v1.0-test
|   |   |   |   |-- v1.0-trainval
|   |   |   |-- v1.0-trainval  
|   |   |       |-- attribute.json
|   |   |           ...
|   |   |       |-- visibility.json
|   |   `-- v1.0-trainval 
|   |   |   |-- train_lmdb  
|   |   |   `-- val_lmdb

    2.3 Mini dataset 版本数据集打包和 meta 文件夹构建

    如果需要使用 mini 数据集，在 BEVFormer 对应的 config 文件"bevformer_tiny_resnet50_detection_nuscenes"中，需要更改以下配置：

    float_trainer/predictor、calibration_trainer/predictor、qat_trainer/predictor、int_infer_predictor/int_infer_trainer的version：数据集的版本，选项为["v1.0-trainval", "v1.0-test", "v1.0-mini"]，选择"v1.0-mini"。

    2.3.1 数据集打包

    如果下载的是 Full dataset(v1.0)中的 Mini、CAN bus expansion 和 Map expansion（v1.3）、nuScenes-lidarseg 这四个项目下的文件，将下载完成的 v1.0-mini.tgz、nuScenes-lidarseg-all-v1.0.tar.bz2、nuScenes-map-expansion-v1.3.zip 和 can_bus.zip 进行解压，解压后的目录如下所示：

    |--nuscenes
    |-- can_bus #can_bus.zip解压后的目录
    |-- lidarseg #nuScenes-lidarseg-all-v1.0.tar.bz2解压后的目录
    |-- maps #nuScenes-map-expansion-v1.3.zip解压后的数据
    |-- samples #v1.0-mini.tar解压后的目录
    |   |-- CAM_BACK
    |   |-- ...
    |   |-- CAM_FRONT_RIGHT
    |   |--  ...
    |   `-- RADAR_FRONT_RIGHT
    |-- sweeps #v1.0-mini.tar解压后的目录
    |   |-- CAM_BACK
    |   |-- ...
    |   |-- CAM_FRONT_RIGHT
    |   |--  ...
    |   `-- RADAR_FRONT_RIGHT
    |-- v1.0-mini #v1.0-mini.tar解压后的目录
        |-- attribute.json
        |
        `-- visibility.json

    进入 horizon_model_train_sample/scripts 目录，使用以下命令将训练数据集和验证数据集打包，格式为 lmdb。如果将上述 nuscenes 文件夹放到上述目录内，下面 WORKSPACE 可替换为pwd。

    #pack train_Set
python3 tools/datasets/nuscenes_packer.py --src-data-dir /WORKSPACE/nuscenes/ --pack-type lmdb --target-data-dir /WORKSPACE/tmp_data/nuscenes/v1.0-trainval --version v1.0-mini --split-name train
#pack val_Set
python3 tools/datasets/nuscenes_packer.py --src-data-dir /WORKSPACE/nuscenes/ --pack-type lmdb --target-data-dir /WORKSPACE/tmp_data/nuscenes/v1.0-trainval --version v1.0-mini --split-name val

    打包结束生成目录如下所示。

    |--scripts
    |-- tmp_data
        |-- nuscenes  
            |-- v1.0-trainval 
                |-- train_lmdb
                |   |--data.mdb
                |   |--lock.mdb    
                `-- train_lmdb
                    |--data.mdb
                    `--lock.mdb

    --src-data-dir 为解压后的 nuscenes 数据集目录；
    --target-data-dir 为打包后数据集的存储目录；
    --version 选项为["v1.0-trainval", "v1.0-test", "v1.0-mini"]，如果进行全量训练和验证设置为 v1.0-trainval，如果仅想了解模型的训练和验证过程，则可以使用 v1.0-mini 数据集；v1.0-test 数据集仅为测试场景，未提供注释。
    全量的 nuscenes 数据集较大，打包时间巨长，如果仅仅作为学习使用 mini 版就好。

    2.3.2 meta 文件夹构建

    在 tmp_data/nuscenes 下创建 meta 文件夹，如果使用--version = "v1.0-mini",将解压后的v1.0-mini 文件夹拷贝到 tmp_data/nuscenes/meta 文件夹内，如果使用--version = "v1.0-trainval",将解压后的 v1.0-trainval 文件夹拷贝到 tmp_data/nuscenes/meta 文件夹内；
    将解压后的 maps 文件夹拷贝到 tmp_data/nuscenes/meta 文件夹内；
    将解压后的 lidarseg 文件夹拷贝到 tmp_data/nuscenes/meta 文件夹内。
    如果使用--version = "v1.0-mini"，此时 tmp_data 的目录结构为：

    |-- tmp_data 
|   |-- nuscenes 
|   |   |-- meta
|   |   |   |-- maps        
|   |   |   |   |-- 36092f0b03a857c6a3403e25b4b7aab3.png
|   |   |   |   |-- ...
|   |   |   |   |-- 93406b464a165eaba6d9de76ca09f5da.png
|   |   |   |   |-- prediction
|   |   |   |   |-- basemap
|   |   |   |   |-- expansion
|   |   |   |-- lidarseg     
|   |   |   |   |-- v1.0-mini
|   |   |   |   |-- v1.0-test
|   |   |   |   |-- v1.0-trainval
|   |   |   |-- v1.0-mini  
|   |   |       |-- attribute.json
|   |   |           ...
|   |   |       |-- visibility.json
|   |   `-- v1.0-trainval 
|   |   |   |-- train_lmdb  ls
cd 
|   |   |   `-- val_lmdb   ls

    3 权重文件准备

    如果不想重复 BEVFormer 模型的训练、校准、量化、定点过程，可从scripts/configs/bev/README.md 内下载 BEVFormer 的权重文件，用于浮点模型精度验证、Calibration 模型精度验证、量化模型精度验证、定点模型精度验证、仿真上板精度验证，或可视化定点模型对于单帧的检测效果，权重文件的下载命令参考：

    wget -c ftp://openexplorer@vrftp.horizon.ai/horizon_torch_samples/3.0.15/py310/modelzoo/qat_origin_modelzoo/bevformer_tiny_resnet50_detection_nuscenes/* --ftp-password='c5R,2!pG'

    将权重文件拷贝到 scripts/tmp_models/ 文件夹内，此时 tmp_models 的目录结构为：

    |-- tmp_data 
|   |-- float-checkpoint-best.pth.tar
|   |-- calibration-checkpoint-best.pth.tar
|   |-- qat-checkpoint-best.pth.tar
|   |-- qat.bc
|   |-- quantized.bc

    如果想完整复现 BEVFormer 模型的训练、校准、量化、定点过程，可在数据集和 Config 文件配置完成后，按照第4、5、6部分的指导进行。

    4 Config 文件配置

    config 内容较多，对于用户来说，需要关注的主要有以下参数：

    device_ids：配置为本机支持的 GPU，根据开发环境进行配置；
    ckpt_dir：权重路径配置；
    data_rootdir：打包数据集路径配置；
    meta_rootdir：meta 文件所在的路径；
    float_trainer/predictor、calibration_trainer/predictor、qat_trainer/predictor、int_infer_trainer/int_infer_predictor的checkpoint_path：权重路径配置；
    float_trainer/predictor、calibration_trainer/predictor、qat_trainer/predictor、int_infer_predictor/int_infer_trainer的version：数据集的版本，选项为["v1.0-trainval", "v1.0-test", "v1.0-mini"]。
    上述 checkpoint_path 和 version 需要根据数据集的版本和使用需求修改。

    5 浮点模型训练和精度验证

    5.1 浮点模型训练

    config 文件中的参数配置完成后，使用以下命令训练浮点模型（如果训练浮点模型，需要使用完整数据集而非 mini）：

    python3 tools/train.py --config configs/bev/bevformer_tiny_resnet50_detection_nuscenes.py --stage float
    float训练后模型ckpt的保存路径为config配置的ckpt_callback中save_dir的值，默认为ckpt_dir。

    5.2 浮点模型精度验证

    浮点模型训练完成以后，可以使用以下命令验证已经训练好的浮点模型精度:

    python3 tools/predict.py --config configs/bev/bevformer_tiny_resnet50_detection_nuscenes.py --stage float

    验证完成后，会在终端打印浮点模型在验证集上检测精度，如下所示：

    Per-class results:
Object Class    AP      ATE     ASE     AOE     AVE     AAE
car     0.517   0.665   0.169   0.230   0.144   0.091
truck   0.366   0.542   0.215   0.260   0.075   0.000
bus     0.325   0.795   0.176   0.229   1.488   0.017
trailer 0.000   1.000   1.000   1.000   1.000   1.000
construction_vehicle    0.000   1.000   1.000   1.000   1.000   1.000
pedestrian      0.432   0.818   0.278   0.889   0.615   0.286
motorcycle      0.360   0.887   0.313   1.799   0.078   0.002
bicycle 0.071   0.815   0.248   0.523   0.922   0.024
traffic_cone    0.400   0.631   0.389   nan     nan     nan
barrier 0.000   1.000   1.000   1.000   nan     nan
2024-07-18 08:09:25,945 INFO [nuscenes_metric.py:378] Node[0] NDS: 0.3204, mAP:0.2472
...
2023-06-06 18:24:10,513 INFO [mean_iou.py:170] Node[0] ~~~~ MeanIOU Summary metrics ~~~~
car_AP: [0.5]:0.1182  [1.0]:0.3794  [2.0]:0.6097  [4.0]:0.7232
...
2023-12-19 17:47:03,046 INFO [metric_updater.py:360] Node[0] Epoch[0] Validation bev_cft_efficientnetb3_nuscenes: NDS[0.3280] 
2023-12-19 17:47:03,058 INFO [logger.py:176] Node[0] ==================================================END PREDICT==================================================
2023-12-19 17:47:03,058 INFO [logger.py:176] Node[0] =============================
=====================END FLOAT PREDICT=============================================

Loaded ground truth annotations for 81 samples.
Filtering predictions
=> Original number of boxes: 24300
=> After distance based filtering: 20631
=> After LIDAR and RADAR points based filtering: 20631
=> After bike rack filtering: 20584
Filtering ground truth annotations
=> Original number of boxes: 4441
=> After distance based filtering: 3785
=> After LIDAR and RADAR points based filtering: 3393
=> After bike rack filtering: 3393
Accumulating metric data...
Calculating metrics...
Saving metrics to: ./WORKSPACE/resultsbevformer_tiny_resnet50_detection_nuscenes
mAP: 0.2472
mATE: 0.8153
mASE: 0.4789
mAOE: 0.7699
mAVE: 0.6651
mAAE: 0.3024
NDS: 0.3204
Eval time: 5.5s

    6 模型量化和编译

    完成浮点训练后，还需要进行量化训练和编译，才能将定点模型部署到板端。地平线对该模型的量化采用 horizon_plugin 框架，经过 Calibration+QAT 量化训练后，使用 compile 的工具将量化模型编译成可以上板运行的 hbm 文件。

    6.1 Calibration

    模型完成浮点训练后，便可进行 Calibration。calibration在forward 过程中通过统计各处的数据分布情况，从而计算出合理的量化参数。 通过运行下面的脚本就可以开启模型的 Calibration 过程：

    python3 tools/train.py --config configs/bev/bevformer_tiny_resnet50_detection_nuscenes.py --stage calibration

    6.2 Calibration 模型精度验证

    Calibration 完成以后，可以使用以下命令验证经过 calib 后模型的精度:

    python3 tools/predict.py --config configs/bev/bevformer_tiny_resnet50_detection_nuscenes.py --stage calibration

    验证完成后，会在终端输出 calib 模型在验证集上检测精度。

    6.3 量化模型训练

    Calibration 完成后，就可以加载 calib 权重开启模型的量化训练。 量化训练其实是在浮点训练基础上的 finetue，具体配置信息在 config的qat_trainer 中定义。量化训练的时候，初始学习率设置为浮点训练的十分之一，训练的 epoch 次数也大大减少。和浮点训练的方式一样，将 checkpoint_path指定为训好的 calibration 权重路径。

    通过运行下面的脚本就可以开启模型的 qat 训练：

    python3 tools/train.py --config configs/bev/bevformer_tiny_resnet50_detection_nuscenes.py --stage qat

    6.4 量化模型精度验证

    Calibration 完成以后，可以使用以下命令验证经过 calib 后模型的精度:

    #qat模型精度验证
python3 tools/predict.py --stage qat --config configs/bev/bevformer_tiny_resnet50_detection_nuscenes.py

    验证完成后，会在终端输出 calib 模型在验证集上检测精度，格式见2.3。

    6.5 导出定点模型

    完成量化训练后，便可以开始导出定点模型。可以通过下面命令来导出：

    python3 tools/export_hbir.py --config configs/bev/bevformer_tiny_resnet50_detection_nuscenes.py

    6.6 定点模型精度验证

    指定 calibration-checkpoint 后，通过运行以下命令进行量化模型的精度验证：

    python3 tools/predict.py --config configs/bev/bevformer_tiny_resnet50_detection_nuscenes.py --stage int_infer
    qat模型的精度验证对象为插入伪量化节点后的模型（float32）；quantize模型的精度验证对象为定点模型（int8），验证的精度是最终的int8模型的真正精度，这两个精度应该是十分接近的。

    6.7 量化模型编译

    在量化训练完成之后，可以使用 compile_perf.py 脚本将量化模型编译成可以板端运行的 hbm 模型，同时该工具也能预估在 BPU 上的运行性能，compile_perf 脚本使用方式如下：

    python3 tools/compile_perf_hbir.py --config configs/bev/bevformer_tiny_resnet50_detection_nuscenes.py --out-dir ./ --opt 3
    opt 为优化等级，取值范围为 0~3，数字越大优化等级越高，编译时间更长，但部署性能更好。
    compile_perf 脚本将生成.html 文件和.hbm 文件（compile 文件目录下），.html 文件为 BPU 上的运行性能，.hbm 文件为上板实测文件。

    运行后，ckpt_dir 的 compile 目录下会产出以下文件：

    |-- compile
|   |-- .html #模型在bpu上的静态性能数据
|   |-- .json 
|   |-- model.hbm  #板端部署的模型
|   |-- model.hbir #编译过程的中间文件
    `-- model.pt   #模型的pt文件

    7 模型推理效果可视化

    如果你希望可以看到训练出来的模型对于单帧的检测效果，我们的 tools 文件夹下面同样提供了预测及可视化的脚本，你只需要运行以下脚本即可，可视化结果将会在 save-path 路径下输出。

    python3 tools/infer_hbir.py --config configs/bev/bevformer_tiny_resnet50_detection_nuscenes.py --save-path ./

    但在此之前你需要在/script 中组织你的 tmp_orig_data 文件夹，文件夹目录为：

    scripts/bev_infer_sample/
└── multi_frames
    ├── homo
    │   ├── cam_intrinsic.npy
    │   ├── ego2global.npy
    │   ├── ego2img.npy
    │   ├── rotation.npy
    │   └── translation.npy
    └── imgs
        ├── n008-2018-08-01-15-16-36-0400__CAM_BACK_LEFT__1533151603547405.jpg
        ├── n008-2018-08-01-15-16-36-0400__CAM_BACK_RIGHT__1533151603528113.jpg
        ├── n008-2018-08-01-15-16-36-0400__CAM_BACK__1533151603537558.jpg
        ├── n008-2018-08-01-15-16-36-0400__CAM_FRONT_LEFT__1533151603504799.jpg
        ├── n008-2018-08-01-15-16-36-0400__CAM_FRONT_RIGHT__1533151603520482.jpg
        └── n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151603512404.jpg

    imgs 为同一时刻下 6 个相机输出的图片；

    其中 homo 文件夹内的 ego2global.npy、ego2img.npy 分别是上述图像对应的自车坐标系到世界坐标系、自车坐标系到像素坐标系的 homo 矩阵，尺寸分别为(4,4)和(6,4,4)。

    上述信息可通过一下脚本生成：

    # tools/gen_infer_data.py
# Usage: python3 tools/gen_infer_data.py
import os
import shutil
from nuscenes.nuscenes import NuScenes 
import numpy as np
from pyquaternion import Quaternion
from PIL import Image

# 注意 >>>>> ：执行前先配置mini_data的位置
mini_data = "/data_set0/v1.0-mini"

def get_folder_from_filename(filename):
    if 'FRONT_LEFT' in filename:
        return 'CAM_FRONT_LEFT'
    elif 'FRONT_RIGHT' in filename:
        return 'CAM_FRONT_RIGHT'
    elif 'FRONT' in filename:
        return 'CAM_FRONT'
    elif 'BACK_LEFT' in filename:
        return 'CAM_BACK_LEFT'
    elif 'BACK_RIGHT' in filename:
        return 'CAM_BACK_RIGHT'
    elif 'BACK' in filename:
        return 'CAM_BACK'
    else:
        raise ValueError(f"Unknown camera type in filename: {filename}")

def get_homography_by_cam(sensor2ego_translation, sensor2ego_rotation, camera_intrinsic):
    #将旋转矩阵转为四元数
    rotation = Quaternion(sensor2ego_rotation).rotation_matrix
    ego2sensor_r = np.linalg.inv(rotation)
    ego2sensor_t = sensor2ego_translation @ ego2sensor_r.T
    ego2sensor = np.eye(4)
    ego2sensor[:3, :3] = ego2sensor_r.T
    ego2sensor[3, :3] = -np.array(ego2sensor_t)
    camera_intrinsic = np.array(camera_intrinsic)
    viewpad = np.eye(4)
    viewpad[
            : camera_intrinsic.shape[0], : camera_intrinsic.shape[1]
        ] = camera_intrinsic
    ego2img = viewpad @ ego2sensor.T
    return ego2img

def save_image(nusc, image_filename, output_folder):
    # 获取图像的绝对路径
    image_path = os.path.join(nusc.dataroot+'/samples',image_filename)
    # 打开并保存图像
    image = Image.open(image_path)
    image.save(os.path.join(output_folder, os.path.basename(image_filename)))

def find_ego2global_from_image(nusc, image_filename):
    # 遍历所有sample_data，找到对应的条目
    for sample_data in nusc.sample_data:
        if image_filename in sample_data["filename"]:
            # 获取对应的ego_pose_token
            ego_pose_token = sample_data['ego_pose_token']
            # 获取ego_pose条目
            ego_pose = nusc.get('ego_pose', ego_pose_token)
            # 提取translation和rotation
            translation = np.array(ego_pose['translation'])
            rotation = np.array(ego_pose['rotation'])
            # 生成ego2global矩阵
            ego2global = get_ego2global_matrix(translation, rotation)
            return ego2global
    return None

def get_ego2global_matrix(translation, rotation):
    """
    从translation和rotation生成ego2global矩阵。
    """
    ego2global = np.eye(4)
    # 设置旋转矩阵
    ego2global[:3, :3] = Quaternion(rotation).rotation_matrix
    # 设置平移矩阵
    ego2global[:3, 3] = translation
    return ego2global

if __name__ == "__main__":
    # 环境初始化
    output_folder = "./bev_infer_sample/multi_frames"
    if os.path.isdir(output_folder):
        shutil.rmtree(output_folder)
    image_output = f"{output_folder}/imgs"
    os.makedirs(image_output, exist_ok=True)
    homo_output = f"{output_folder}/homo"
    os.makedirs(homo_output, exist_ok=True)

    #初始化，这里以v1.0-mini数据集为例
    nusc = NuScenes(version='v1.0-mini', dataroot=mini_data, verbose=True)
    #6张图像的文件名
    #要注意这里的输入顺序必须遵循：
    #FRONT_LEFT,FRONT,FRONT_RIGHT,BACK_LEFT,BACK,BACK_RIGHT
    image_list=[
    'n008-2018-08-01-15-16-36-0400__CAM_FRONT_LEFT__1533151603504799.jpg',
    'n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151603512404.jpg',
    'n008-2018-08-01-15-16-36-0400__CAM_FRONT_RIGHT__1533151603520482.jpg',
    'n008-2018-08-01-15-16-36-0400__CAM_BACK_LEFT__1533151603547405.jpg',
    'n008-2018-08-01-15-16-36-0400__CAM_BACK__1533151603537558.jpg',
    'n008-2018-08-01-15-16-36-0400__CAM_BACK_RIGHT__1533151603528113.jpg'
    ]
    # 获取数据集中所有的sample
    # 包含'token','calibrated_sensor_token','filename','channel'等信息
    sample_data=nusc.sample_data 
    ego_poses = nusc.ego_pose
    # 获取传感器定义：
    # 包含'token','sensor_toker','translation','rotation','camera_intrinsic'等信息
    cali_sensor=nusc.calibrated_sensor
    translation=[]
    rotation=[]
    cam_intrinsic=[]
    ego2imgs=[]
    for i,img in enumerate(image_list):
        print(img) 
        for sample in sample_data:
            if img in sample["filename"]:
                folder = get_folder_from_filename(img)
                save_image(nusc, os.path.join(folder, img), image_output)
                #1. 选择输入图片的name来获取捕获该图片的"calibrated_sensor_token"
                calibrated_sensor_token=(sample["calibrated_sensor_token"])
                print("calibrated_sensor_token:",calibrated_sensor_token)
                #2.根据"calibrated_sensor_token"获取传感器的"translation"、
                # "rotation"、和"camera_intrinsic"等参数
                for sensor in cali_sensor:
                    if sensor["token"]==calibrated_sensor_token: 
                        #获取translation矩阵
                        sensor2ego_translation=np.array(sensor["translation"])
                        # print("sensor2ego_translation:",sensor2ego_translation)
                        translation.append(sensor2ego_translation)
                        #获取rotation矩阵
                        sensor2ego_rotation=np.array(sensor["rotation"])
                        # print("sensor2ego_rotation:",sensor2ego_rotation)
                        rotation.append(sensor2ego_rotation)
                        #获取相机内参矩阵
                        camera_intrinsic=np.array(sensor["camera_intrinsic"])
                        # print("camera_intrinsic:",camera_intrinsic)
                        cam_intrinsic.append(camera_intrinsic)
                        #计算homography矩阵
                        ego2img=get_homography_by_cam(
                        sensor2ego_translation,
                        sensor2ego_rotation, 
                        camera_intrinsic
                        )
                        ego2imgs.append(ego2img)
    
    #导出translation矩阵，旋转矩阵、相机内参矩阵、homo矩阵
    translation=np.array(translation)
    np.save(f"{homo_output}/translation",translation)
    rotation=np.array(rotation)
    np.save(f"{homo_output}/rotation",rotation)
    cam_intrinsic=np.array(cam_intrinsic)
    np.save(f"{homo_output}/cam_intrinsic",cam_intrinsic)
    ego2imgs=np.array(ego2imgs)
    print('ego2imgs',ego2imgs.shape)
    np.save(f"{homo_output}/ego2img",ego2imgs)
    ego2global = find_ego2global_from_image(nusc, image_list[0])
    print('ego2global',ego2global.shape)
    np.save(f"{homo_output}/ego2global",ego2global)

    可视化示例如下：

    算法自动驾驶
    地平线智驾开发者
    1 声望2 粉丝

    地平线智能驾驶开发者社区旨在连接智能驾驶领域的开发者和对相关技术感兴趣的其他行业开发者、从业者。

    下一篇 »
    浅谈 Occupancy

    引用和评论

    推荐阅读
    头像
    征程 6E camera diag sample

    地平线智驾开发者

    头像
    告别卡顿，畅享GitHub：国内开发者必看的五大加速访问与下载技巧

    汀丶阅读 3.2k

    头像
    基础思想总结

    zxl20070701赞 1阅读 877

    头像
    ARTS 打卡 第四周，游刃有余

    三掌柜赞 1阅读 273

    头像
    CUDA Toolkit常见安装问题一览

    techlead_kris赞 1阅读 431

    头像
    Python深度学习股价预测、量化交易策略：LSTM、GRU深度门控循环神经网络|附代码数据

    拓端tecdat阅读 716

    头像
    IOS逆向助手-无需越狱就能逆向IOS软件的强大工具

    Soujer阅读 664

    0 条评论
    得票最新