Rotation Correction with Diffusion

基于 Marigold 扩散模型深度估计框架，改造用于图像旋转矫正任务。通过微调 Stable Diffusion v2 的 UNet，利用深度信息辅助判断和矫正图像旋转角度。

项目概述

本项目在 Marigold (CVPR 2024 Oral) 的基础上进行改造：

1. UNet 输入通道扩展：从 4 通道扩展到 8 通道，同时处理 RGB 和深度潜特征
2. 动态文本条件：根据文件名中的旋转角度生成特定提示词（如 "rotated X degrees clockwise..."），引导模型学习旋转-深度的对应关系
3. RotationDataset：新增自定义旋转矫正数据集类，支持按旋转角度组织数据
4. Multi-resolution noise：引入多分辨率噪声增强训练稳定性

环境配置

# 创建环境 (推荐 mamba)
mamba env create -n marigold --file environment.yaml
conda activate marigold

# 或用 pip
python -m venv venv/marigold && source venv/marigold/bin/activate
pip install -r requirements.txt

# 训练额外依赖
pip install -r requirements++.txt -r requirements+.txt -r requirements.txt

核心依赖：torch==2.4.0, diffusers>=0.25.0, accelerate>=0.22.0, transformers>=4.32.1

快速开始

推理

# 标准推理 (LCM 快速模式)
python run.py --input_rgb_dir input/ --output_dir output/

# DDIM 高精度推理
python run.py \
    --checkpoint prs-eth/marigold-v1-0 \
    --denoise_steps 50 \
    --ensemble_size 10 \
    --input_rgb_dir input/ \
    --output_dir output/

# 加载自定义微调 checkpoint
python app.py \
    --custom_ckpt_path output/<run>/checkpoint/latest/unet \
    --input_rgb_dir input/ \
    --output_dir output/

# 多 GPU 平衡模式
python run_multi.py --input_rgb_dir input/ --output_dir output/

Gradio Web Demo

python app1.py

推理参数

参数	说明	默认值
--checkpoint	预训练模型路径或 HuggingFace hub 名称	prs-eth/marigold-lcm-v1-0
--custom_ckpt_path	自定义微调 UNet checkpoint 路径	None
--input_rgb_dir	输入图片目录	必填
--output_dir	输出目录	必填
--denoise_steps	扩散去噪步数，越多精度越高	None (自动)
--ensemble_size	集成推理次数，越多精度越高	5
--processing_res	处理分辨率，0 表示原始分辨率	None (自动)
--half_precision / --fp16	半精度推理，减少显存占用	False
--seed	随机种子，用于可复现结果	None
--color_map	深度图着色方案	Spectral
--apple_silicon	Apple Silicon MPS 加速	False

训练

准备

pip install -r requirements++.txt -r requirements+.txt -r requirements.txt

export BASE_DATA_DIR=<训练数据目录>
export BASE_CKPT_DIR=<预训练模型目录>

将 Stable Diffusion v2 checkpoint 下载至 ${BASE_CKPT_DIR}。

运行训练

python train.py --config config/train_marigold.yaml

从 checkpoint 恢复训练

python train.py --resume_run output/marigold_base/checkpoint/latest

训练参数

python train.py --help

常用参数：

• --config：配置文件路径
• --resume_run：恢复训练的 checkpoint 路径
• --base_data_dir：训练数据目录
• --base_ckpt_dir：预训练模型目录
• --no_wandb：禁用 WandB 日志
• --add_datetime_prefix：输出目录添加时间戳前缀

评估

# 安装额外依赖
pip install -r requirements+.txt

# 示例：NYU 数据集
bash script/eval/11_infer_nyu.sh    # 推理
bash script/eval/12_eval_nyu.sh     # 计算指标

# 其他数据集
bash script/eval/21_infer_kitti.sh && bash script/eval/22_eval_kitti.sh
bash script/eval/31_infer_eth3d.sh && bash script/eval/32_eval_eth3d.sh
bash script/eval/41_infer_scannet.sh && bash script/eval/42_eval_scannet.sh
bash script/eval/51_infer_diode.sh && bash script/eval/52_eval_diode.sh

或直接使用评估脚本：

python eval.py \
    --prediction_dir output/<run>/evaluation/ \
    --output_dir eval_results/ \
    --dataset_config config/dataset/data_nyu_test.yaml \
    --base_data_dir ${BASE_DATA_DIR}

项目结构

├── marigold/              # 核心流水线
│   ├── marigold_pipeline.py  # MarigoldPipeline (VAE + UNet + Text Encoder)
│   ├── inference.py          # 共享推理模块
│   └── util/                 # 工具函数 (批处理、集成、图像处理)
├── src/                   # 训练与数据处理
│   ├── trainer/              # MarigoldTrainer
│   ├── dataset/              # 数据集类 (Hypersim, KITTI, RotationDataset 等)
│   └── util/                 # 损失函数、指标、配置、噪声等工具
├── config/                # YAML 配置文件
├── script/                # 工具脚本
│   ├── eval/                 # 评估脚本
│   ├── heatmap.py            # 图像差异热力图
│   ├── visualize_rotation.py # 旋转数据集可视化
│   └── dataset_preprocess/   # 数据预处理
├── run.py                 # 标准推理入口
├── app.py                 # 自定义 checkpoint 推理入口
├── run_multi.py           # 多 GPU 推理入口
├── app1.py                # Gradio Web Demo
├── train.py               # 训练入口
└── eval.py                # 评估入口

数据集

支持的数据集：Hypersim, Virtual KITTI 2, NYU v2, KITTI, ETH3D, DIODE, ScanNet，以及自定义 RotationDataset。

数据集配置文件位于 config/dataset/，数据划分文件位于 data_split/。

致谢

本项目基于 Marigold (CVPR 2024 Oral, Best Paper Award Candidate) 进行改造。感谢原作者 Bingxin Ke, Anton Obukhov 等的工作。

License

This work is licensed under the Apache License, Version 2.0 (as defined in the LICENSE).