Skip to content

feat: add VKITTI far-depth post-training and separate Stage 1 training #38

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Haruko386 merged 4 commits into from
Sep 21, 2026
Merged

feat: add VKITTI far-depth post-training and separate Stage 1 training #38

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2310835
feat: add VKITTI far-depth supervision and separate Stage 1 training
Haruko386 Sep 21, 2026
73a3e70
fix: make VKITTI far-depth supervision a post-training step
Haruko386 Sep 21, 2026
7341b34
fix: require safetensors for external checkpoint initialization
Haruko386 Sep 21, 2026
eb085f1
update readme
Haruko386 Sep 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
108 changes: 85 additions & 23 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,18 +7,18 @@ This repository is based on [Marigold](https://marigoldmonodepth.github.io), CVP
[![Hugging Face Model](https://img.shields.io/badge/🤗%20Hugging%20Face-Model-green)](https://huggingface.co/developy/ApDepth)
[![Hugging Face Demo](https://img.shields.io/badge/🤗%20Hugging%20Face-Demo-purple)](https://huggingface.co/spaces/developy/ApDepth)

[**Haruko386**](https://haruko386.github.io/),
<!-- [**Haruko386**](https://haruko386.github.io/),
[Shuai Yuan](https://syjz.teacher.360eol.com/teacherBasic/preview?teacherId=23776),
[Mingbo Lei](https://github.com/Ltohka),
[Yibo Chen](#)
[Yibo Chen](#) -->

![cover](doc/cover.png)

> We present **ApDepth**, a deterministic single-step diffusion framework for monocular depth estimation. Built upon Marigold (derived from Stable Diffusion) and fine-tuned on synthetic datasets (Hypersim and Virtual KITTI), ApDepth overcomes the feature representation bottlenecks of standard diffusion models. Ultimately, our framework achieves **highly competitive geometric accuracy** and **exceptional object edge refinement**, all while delivering significantly accelerated inference speeds.

## 📢 News
- **2026-04-06:** `ApDepth V2-0` is released!
- **2026-04-03:** We officially release the complete code for **ApDepth**! The repository now includes the full coarse-to-fine two-stage training pipeline and evaluation scripts.
- **2026-04-03:** We officially release the code for **ApDepth**! Stage 1 feature-alignment training is maintained in [ApDepth_Stage1](https://github.com/Haruko386/ApDepth_Stage1); this repository provides Stage 2 training, inference and evaluation.
- **2026-01-15:** We successfully introduce a spatial-preserving **Conv Adapter** and a **Cosine Similarity Loss** to enhance feature alignment, alongside a **Pixel-level $L_1$ Loss** to establish an accurate global metric scale.
- **2025-10-25:** Inspired by DepthMaster, we propose a two-stage loss function training strategy based on `ApDepth V1-0`. In the first stage, we perform foundational training using MSE loss. In the second stage, we learn edge structures through FFT loss. Based on this, we introduce `ApDepth V1-1`.
- **2025-10-09:** We propose a novel diffusion-based depth estimation framework guided by pre-trained models.
Expand Down Expand Up @@ -182,6 +182,10 @@ You can get the result under `output/eval`

## 🏋️ Training

The workflow is **Stage 1 feature alignment → Stage 2 main training → optional
[VKITTI far-depth post-training](#post-training)**. Complete Stage 2, including
its FFT refinement, before starting the optional post-training run.

Based on the previously created environment, install extended requirements:

```bash
Expand All @@ -197,56 +201,114 @@ export BASE_CKPT_DIR=YOUR_CHECKPOINT_DIR # directory of pretrained checkpoint

Download Stable Diffusion v2 [checkpoint](https://huggingface.co/sd2-community/stable-diffusion-2-1) into `${BASE_CKPT_DIR}`

Download the checkpoint of [Depth-Anything-V2](https://github.com/DepthAnything/Depth-Anything-V2) into `DA2/checkpoints/`. Stage 1 uses the ViT-G checkpoint by default at `DA2/checkpoints/depth_anything_v2_vitg.pth` as the external semantic encoder.
Download the ViT-G checkpoint of [Depth-Anything-V2](https://github.com/DepthAnything/Depth-Anything-V2) to `DA2/checkpoints/depth_anything_v2_vitg.pth`. The ApDepth pipeline uses it to generate the depth prior during training and inference.

Prepare for [Hypersim](https://github.com/apple/ml-hypersim) and [Virtual KITTI 2](https://europe.naverlabs.com/research/computer-vision/proxy-virtual-worlds-vkitti-2/) datasets and save into `${BASE_DATA_DIR}`. Please refer to [this README](script/dataset_preprocess/hypersim/README.md) for Hypersim preprocessing. Stage 1 uses `config/dataset/dataset_apdepth_train_s1.yaml`; the main training stage uses `config/dataset/dataset_train.yaml`.
Prepare [Hypersim](https://github.com/apple/ml-hypersim) and [Virtual KITTI 2](https://europe.naverlabs.com/research/computer-vision/proxy-virtual-worlds-vkitti-2/) under `${BASE_DATA_DIR}`. Please refer to [this README](script/dataset_preprocess/hypersim/README.md) for Hypersim preprocessing. Configure the training paths in `config/dataset/dataset_train.yaml` and prepare the validation datasets listed in `config/dataset/dataset_val.yaml` and `config/dataset/dataset_vis.yaml`.

------------

**Stage 1: feature alignment pre-training**

Stage 1 has been integrated into this repository. It trains the denoising U-Net and DINOv2 adapter in feature space before the main ApDepth training stage.
Stage 1 is maintained separately in
[Haruko386/ApDepth_Stage1](https://github.com/Haruko386/ApDepth_Stage1).
Follow that repository's setup and training instructions. This repository contains
only the main Stage 2 trainer; it does not include Stage 1 code or dependencies.

------------

**Stage 2: main ApDepth training**

To initialize Stage 2 from an existing Stage 1 checkpoint, start a new run with
`--init_checkpoint`. This loads only the U-Net; optimizer and iteration counters
start fresh using the Stage 2 config. `config/train_apdepth.yaml` keeps far-depth
supervision disabled throughout the main training stage:

```bash
bash script/apdepth_train_s1.sh
python train.py --config config/train_apdepth.yaml \
--init_checkpoint /path/to/stage1/checkpoint/iter_020000 --no_wandb
```

Equivalent direct command:
The initialization directory must contain
`unet/diffusion_pytorch_model.safetensors`. `--init_checkpoint` rejects `.bin`
checkpoints.

To train the main stage directly from the base model without Stage 1 initialization:

```bash
python apdepth_train_s1.py \
--config config/apdepth_train_s1.yaml \
--base_data_dir ${BASE_DATA_DIR} \
--base_ckpt_dir ${BASE_CKPT_DIR} \
--output_dir output/stage1 \
--no_wandb
python train.py --config config/train_apdepth.yaml --no_wandb
```

Resume Stage 1 from a checkpoint, e.g.
Resume from a checkpoint, e.g.

```bash
python apdepth_train_s1.py --resume_run output/stage1/checkpoint/latest --no_wandb
python train.py --resume_run output/train_apdepth/checkpoint/latest --no_wandb
```

------------

**Stage 2: main ApDepth training**
## 🌤️ Optional post-training after Stage 2 <a name="post-training"></a>

**Start this step only after Stage 2 training has finished.** The new VKITTI
far-depth supervision feature is enabled by `config/train_sky_finetune.yaml` for
a separate post-training run initialized from the completed Stage 2 U-Net.

Measured depths in [80, 655.35] m supply an additional clipped far-depth target
before VAE encoding; missing depth remains excluded. Evaluation masks and
normalization quantiles are unchanged. The separately averaged far latent MSE
and pixel L1 terms have weight 0.5. For a post-training ablation, set
`far_depth_supervision.enabled: false` in the post-training config.

Use the same environment, datasets and base checkpoints prepared for Stage 2.
First audit the native VKITTI depth and target coverage:

```bash
python train.py --config config/train_apdepth.yaml --no_wandb
python -m script.audit_far_supervision --config config/train_sky_finetune.yaml \
--base_data_dir "${BASE_DATA_DIR}" --samples 20
```

Resume from a checkpoint, e.g.
Check the RGB, far masks, target images and `coverage.json` in `output/far_audit`.
If preprocessing replaced the far-plane values with zero, restore the original
depth files. Zero depth is not treated as sky. Adjust each dataset config's `dir`
to match the actual layout under `BASE_DATA_DIR`.

Start a new post-training run from the final Stage 2 checkpoint (adjust the path
if Stage 2 was saved elsewhere):

```bash
python train.py --resume_run output/train_apdepth/checkpoint/latest --no_wandb
python train.py --config config/train_sky_finetune.yaml \
--init_checkpoint output/train_apdepth/checkpoint/iter_021000 --no_wandb
```

------------
This resets the optimizer, learning-rate schedule and iteration counter, then
runs 6000 new optimization steps at LR 5e-6 with 100 warmup steps and saves
backups every 1000 steps in `output/train_sky_finetune/checkpoint`. It retains the
90% Hypersim / 10% VKITTI sampling, original batch settings and periodic validation.
These 6000 steps use latent MSE + pixel L1 reconstruction plus the new far-depth
loss; the Stage 2 FFT refinement has already finished. Initialization requires
`unet/diffusion_pytorch_model.safetensors` in the supplied checkpoint directory.

If this post-training run is interrupted, resume its own checkpoint:

```bash
python train.py --resume_run output/train_sky_finetune/checkpoint/latest --no_wandb
```

`--resume_run` restores the saved config and training state. Starting post-training
requires `--init_checkpoint` with the post-training config above; resuming a
Stage 2 checkpoint would continue Stage 2 with its saved settings.
`--resume_run` and `--init_checkpoint` cannot be combined.

Monitor `train/far_loss`, `train/far_pixel_ratio` and `train/far_latent_ratio`.
Compare KITTI and NYU against the completed Stage 2 model after post-training; this change
alone does not establish a metric improvement.

**Evaluating results**
## Evaluating trained checkpoints

Only the U-Net is updated and saved during training. To use the inference pipeline with your training result, replace `unet` folder in `train_apdepth` checkpoints with that in the `checkpoint` output folder. Then refer to [this section](#evaluation) for evaluation.
Stage 2 and optional post-training update and save the U-Net. For inference, copy a complete
ApDepth pipeline checkpoint to a new directory, then replace that copy's `unet/`
folder with the selected training checkpoint's `unet/`. The pipeline's VAE,
tokenizer, text encoder and scheduler are still required; a training checkpoint
alone is not a complete inference pipeline. Then refer to [evaluation](#evaluation).

> [!IMPORTANT]
>
Expand Down
Loading
Loading