MagicDrive

✨ 欢迎查看我们关于3D场景生成的新作品 MagicDrive3D！

✨ 如果您需要视频生成功能，请在video分支中查找代码。

MagicDrive生成的视频（点击图片查看视频）。

本仓库包含以下论文的实现：

MagicDrive：具有多样3D几何控制的街景生成 高瑞远1*，陈凯2*，谢恩泽3^，洪兰清3，李震国3，杨意2，徐强1^ 1香港中文大学 2香港科技大学 3华为诺亚方舟实验室 *共同第一作者 ^通讯作者

摘要

<details> <summary>简述 MagicDrive生成高质量的街景图像和视频，具有多样化的3D几何控制和多视图一致性，可作为各种感知任务中的数据引擎。</summary>

近期扩散模型的进展显著提高了2D控制下的数据合成能力。然而，街景生成中的精确3D控制，对3D感知任务至关重要，仍然难以实现。具体来说，使用鸟瞰图（BEV）作为主要条件常常导致几何控制（如高度）的挑战，影响物体形状、遮挡模式和路面高程的表现，这些对感知数据合成至关重要，尤其是对3D物体检测任务。在本文中，我们提出了MagicDrive，一个新颖的街景生成框架，通过定制的编码策略，提供多样化的3D几何控制，包括相机姿态、道路地图和3D边界框，以及文本描述。此外，我们的设计还包含一个跨视图注意力模块，确保多个相机视图之间的一致性。通过MagicDrive，我们实现了高保真度的街景图像和视频合成，捕捉细微的3D几何特征和各种场景描述，增强了BEV分割和3D物体检测等任务。

</details>

新闻

[2024/06/07] MagicDrive现在可以生成60帧视频！我们发布了配置文件：rawbox_mv2.0t_0.4.3_60.yaml。在项目页面查看我们的演示。
[2024/06/07] 我们发布了16帧视频生成的预训练权重。点击查看！
[2024/06/01] 我们在ECCV2024举办W-CODA工作坊。挑战赛赛道2将使用MagicDrive作为基线。我们将在不久的将来发布更多资源。敬请关注！

方法

在MagicDrive中，我们采用两种策略（交叉注意力和加性编码器分支）来注入文本提示、相机姿态、物体边界框和道路地图作为生成条件。我们还提出了一个跨视图注意力模块，以实现多视图一致性。

待办事项

基本分辨率（224x400）的配置文件和预训练权重
基本分辨率（224x400）的演示
用于交互式边界框编辑的图形用户界面
训练和测试代码发布
FID测试代码
高分辨率的配置和预训练权重

入门指南

环境设置

克隆此仓库及其子模块

git clone --recursive https://github.com/cure-lab/MagicDrive.git

代码在V100服务器上使用Pytorch==1.10.2和cuda 10.2进行了测试。要设置Python环境，请按以下步骤操作：

# 选项1：仅运行GUI
pip install -r requirements/gui.txt
# 😍 我们的GUI不需要mm系列包。
# 继续从`third_party`安装diffusers。

# 选项2：运行完整的测试演示（同时在训练前测试您的环境）
cd ${ROOT}
pip install -r requirements/dev.txt
# 继续按如下方式安装`third_party`。

我们选择使用cd ${FOLDER}; pip -vvv install .安装以下包的源代码：

# 安装第三方包
third_party/
├── bevfusion -> 基于db75150
├── diffusers -> 基于v0.17.1 (afcca39)
└── xformers  -> 基于v0.0.19 (8bf59c9)，可选

查看关于我们xformers的说明。如果您在环境设置中遇到问题，请先查看常见问题。

使用以下命令为accelerate设置默认配置：

accelerate config

我们的默认日志目录是${ROOT}/magicdrive-log。请做好准备。

预训练权重

我们的训练基于stable-diffusion-v1-5。我们假设您将它们放在${ROOT}/pretrained/中，如下所示：

{ROOT}/pretrained/stable-diffusion-v1-5/
├── text_encoder
├── tokenizer
├── unet
├── vae
└── ...

使用MagicDrive生成街景

从onedrive下载我们为MagicDrive预训练的权重，并将其放在${ROOT}/pretrained/中

运行我们的演示

👍 我们建议用户首先运行我们的交互式GUI，因为我们已经最小化了GUI演示的依赖项。

cd ${ROOT}
python demo/interactive_gui.py
# 基于gradio的gui，使用您的网络浏览器

根据#37的建议，提示可以通过GUI配置！

gui

运行我们的相机视图生成演示。

cd ${ROOT}
python demo/run.py resume_from_checkpoint=magicdrive-log/SDv1.5mv-rawbox_2023-09-07_18-39_224x400

生成的图像将位于magicdrive-log/test中。更多信息可以在演示文档中找到。

训练MagicDrive

准备数据

我们按照bevfusion的说明准备nuScenes数据集。具体来说，

从网站下载nuScenes数据集并将它们放在./data/中。您应该有这些文件：

data/nuscenes
├── maps
├── mini
├── samples
├── sweeps
├── v1.0-mini
└── v1.0-trainval

[!提示] 您可以从OneDrive下载.pkl文件。它们应该足够用于训练和测试。

通过以下方式生成mmdet3d注释文件：

python tools/create_data.py nuscenes --root-path ./data/nuscenes \
  --out-dir ./data/nuscenes_mmdet3d_2 --extra-tag nuscenes

您应该有这些文件：

data/nuscenes_mmdet3d_2
├── nuscenes_dbinfos_train.pkl (-> ${bevfusion-version}/nuscenes_dbinfos_train.pkl)
├── nuscenes_gt_database (-> ${bevfusion-version}/nuscenes_gt_database)
├── nuscenes_infos_train.pkl
└── nuscenes_infos_val.pkl

注意：如上所示，一些文件可以与bevfusion的原始版本软链接。如果某些文件位于data/nuscenes中，您可以手动将它们移动到data/nuscenes_mmdet3d_2。

（可选）为加快数据加载速度，我们为BEV地图准备了h5格式的缓存文件。它们可以通过tools/prepare_map_aux.py与configs/dataset中的不同配置生成。例如：
```
python tools/prepare_map_aux.py +process=train
python tools/prepare_map_aux.py +process=val
```
您将得到类似./val_tmp.h5和./train_tmp.h5的文件。生成后，您必须正确重命名缓存文件。我们的默认设置是：
```
data/nuscenes_map_aux
├── train_26x200x200_map_aux_full.h5 (42G)
└── val_26x200x200_map_aux_full.h5 (9G)
```

训练模型

使用以下命令启动训练（使用8xV100）：

accelerate launch --mixed_precision fp16 --gpu_ids all --num_processes 8 tools/train.py \
  +exp=224x400 runner=8gpus

在训练过程中，您可以查看tensorboard以获取日志和中间结果。

此外，我们提供了调试配置来测试您的环境和数据加载过程（使用2xV100）：

accelerate launch --mixed_precision fp16 --gpu_ids all --num_processes 2 tools/train.py \
  +exp=224x400 runner=debug runner.validation_before_run=true

测试模型

训练完成后，您可以通过以下方式测试您的驾驶视图生成模型：

python tools/test.py resume_from_checkpoint=${YOUR MODEL}
# 以我们的预训练模型为例
python tools/test.py resume_from_checkpoint=./pretrained/SDv1.5mv-rawbox_2023-09-07_18-39_224x400

请在./magicdrive-log/test/中查看结果。

测试FID

首先，您应该生成完整的验证集：

python perception/data_prepare/val_set_gen.py \
  resume_from_checkpoint=./pretrained/SDv1.5mv-rawbox_2023-09-07_18-39_224x400 \
  task_id=224x400 fid.img_gen_dir=./tmp/224x400 +fid=data_gen +exp=224x400
  # 对于map=zero作为CFG的空条件，添加`runner.pipeline_param.use_zero_map_as_unconditional=true`

对于此脚本，也可以通过accelerate使用多进程/多节点。只需使用类似训练的命令启动即可。

然后，使用以下命令测试FID分数：

# 我们假设您的torch缓存目录位于"../pretrained/torch_cache/"。如果您想
# 使用默认位置，请注释掉"tools/fid_score.py"中倒数第二行。
python tools/fid_score.py cfg \
  resume_from_checkpoint=./pretrained/SDv1.5mv-rawbox_2023-09-07_18-39_224x400 \
  fid.rootb=tmp/224x400

或者，我们在这里提供了预生成的验证集样本。您可以将它们放在./tmp中，并通过以下命令启动测试：

python tools/fid_score.py cfg \
  resume_from_checkpoint=./pretrained/SDv1.5mv-rawbox_2023-09-07_18-39_224x400 \
  fid.rootb=tmp/224x400/samples  # FID=14.46065995481922
  # 或 `fid.rootb=tmp/224x400map0/samples`，FID=16.195992872931697

定量结果

<details> <summary>比较MagicDrive与其他方法的生成质量：</summary>

主要结果

</details> <details> <summary>使用MagicDrive生成的图像进行训练支持：</summary>

可训练性

</details>

更多结果可以在主论文中找到。

定性结果

更多结果可以在主论文中找到。

引用我们

@inproceedings{gao2023magicdrive,
  title={{MagicDrive}: Street View Generation with Diverse 3D Geometry Control},
  author={Gao, Ruiyuan and Chen, Kai and Xie, Enze and Hong, Lanqing and Li, Zhenguo and Yeung, Dit-Yan and Xu, Qiang},
  booktitle = {International Conference on Learning Representations},
  year={2024}
}