hzxie/CityDreamer

CityDreamer – 无限3D城市生成的生成模型

是什么 – CityDreamer 是一个研究级的组合式生成模型，能够合成无限且逼真的3D城市景观。它通过组合三个学习组件实现：

1. 无限布局生成器 – 使用 VQ-VAE + 自回归采样器生成全局城市布局（道路网络、街区排列）。
2. 背景元素生成器 – 用地形、植被、天空等非建筑元素填充布局。
3. 建筑实例生成器 – 为每个街区生成详细的建筑几何和纹理。

该系统基于 PyTorch（已测试版本 1.13.1 + CUDA 11.7），推理时可在单个高内存 GPU（如 RTX 3090）上运行，训练则假设使用多 GPU 环境。

---

快速开始（推理）

1. 安装 – 安装所需的 PyTorch 版本，克隆仓库，并运行 pip install -r requirements.txt。
2. 编译 CUDA 扩展 – 位于 extensions/ 中的扩展（每个子文件夹通过 pip install . 安装）。
3. 运行演示：

   python -m demo.run   # 在 http://localhost:3186 打开 Web UI
   

   或使用 CLI：

   python -m scripts.inference   # 输出 output/rendering.mp4
   

   如果 GPU 显存不足 24 GB，请调整 --patch_height / --patch_width。

---

训练（研究用途）

1. 数据集 – CityDreamer 提供两个大规模数据集：
   • OSM（OpenStreetMap）道路网络数据。
   • GoogleEarth 航拍图像及实例分割掩码（通过 SEEM 生成）。 两者均可通过提供的网关链接下载。
2. 准备 GoogleEarth 标注（约 1 TB） – 运行 SEEM 分割脚本，然后执行 scripts/dataset_generator.py。
3. 训练三个模块（每个模块使用通用 run.py 入口点配合 torchrun）：
   • 布局用 VQ-VAE → 4 GPU。
   • 布局用采样器 → 2 GPU（加载 VQ-VAE 检查点）。
   • 背景元素生成器 → 8 GPU（设置 BUILDING_MODE=False）。
   • 建筑实例生成器 → 8 GPU（设置 BUILDING_MODE=True）。 各部分的配置调整详见 README。

---

预训练模型与资源

组件	检查点（下载）
布局生成器（采样器）	LayoutGen.pth
背景元素生成器	CityDreamer-Bgnd.pth
建筑实例生成器	CityDreamer-Fgnd.pth

实时演示托管于 Hugging Face Spaces，短篇亮点视频链接至 YouTube。

---

许可证

代码以 NTU S-Lab License 1.0（参见仓库中的 LICENSE）发布。再分发必须遵守该许可证条款。

---

开发团队

由南洋理工大学 S-Lab 的谢浩哲、陈兆熙、洪方舟、刘子威开发。该工作已发表于 CVPR 2024，并有后续研究（CityDreamer4D）被 TPAMI 接受。

---

总结 – CityDreamer 为对大规模、程序化生成真实 3D 城市环境感兴趣的研究者提供了完整管道（数据准备、训练脚本、推理演示）。