Spaces:
Running on Zero
Running on Zero
| """DepthWizard v3 — configuration. | |
| One dataclass, every knob CLI-overridable (same pattern as v1/v2). | |
| What changed vs v2 (see README.md §1 for the full post-mortem): | |
| * scale augmentation is *crop-then-resample*, so a tile is never zero-padded | |
| * data is materialised once into memmap shards instead of streamed per step | |
| * encoder is fully unfrozen after a short warmup, with layer-wise LR decay | |
| * heads run at 1/2 resolution (DPT-standard) instead of full res | |
| * height-stratum balanced loss weighting + weight EMA | |
| * the train-time preprocessing contract is serialised into the checkpoint so | |
| `infer/predict.py` cannot drift from it | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import os | |
| from dataclasses import dataclass, fields | |
| from pathlib import Path | |
| _V3_DIR = Path(__file__).resolve().parent | |
| _DEFAULT_OUTPUT_DIR = os.environ.get("DW_OUTPUT_DIR", str(_V3_DIR / "outputs" / "v3")) | |
| _DEFAULT_DATA_ROOT = os.environ.get("DW_DATA_ROOT", str(_V3_DIR / "data")) | |
| # GAMUS ships integer class ids in classes/<split>/<STEM>_CLS.h5. v3/v4 used | |
| # neutral names ("class0".."class6") until the ids were pinned from evidence; | |
| # v5 pins them from the measured per-id heights (v5 config.py, 2861 GAMUS test | |
| # tiles): building reads ~9.4 m, tree ~10.2 m, water is the flattest class. This | |
| # is the published GAMUS order. The Space reports the class map and extracts 3D | |
| # objects (`infer/objects.py`) under these names, so a checkpoint trained on a | |
| # different id space must change them here. | |
| CLASS_IDS: tuple[int, ...] = (0, 1, 2, 3, 4, 5, 6) | |
| CLASS_NAMES: tuple[str, ...] = ("other", "ground", "low_veg", "building", | |
| "water", "road", "tree") | |
| N_SEG_CLASSES = 8 # 7 real ids + 7 == ignore | |
| SEG_IGNORE_INDEX = 7 | |
| # Per-height-stratum edges (m) for the stratified / balanced RMSE table. | |
| HEIGHT_STRATA_M: tuple[float, ...] = (0.0, 2.0, 5.0, 10.0, 20.0, 1e9) | |
| class Config: | |
| # ----- data ------------------------------------------------------ | |
| datasets: str = "gamus,synrs3d" # comma list of {gamus,geonrw,synrs3d} | |
| # (a source that was not prepared is skipped) | |
| data_root: str = _DEFAULT_DATA_ROOT # where prepare_data.py wrote the shards | |
| sampler_weights: str = "gamus:1,geonrw:1,synrs3d:1" | |
| gamus_repo: str = "earthflow/GAMUS" | |
| geonrw_repo: str = "torchgeo/geonrw" | |
| synrs3d_repo: str = "JTRNEO/SynRS3D" | |
| # An "epoch" is a fixed number of random crops, decoupled from tile count so | |
| # the LR schedule and the wall-clock budget stay predictable. | |
| crops_per_epoch: int = 12000 | |
| val_tiles: int = 400 | |
| max_valid_height_m: float = 200.0 # AGL above this is LiDAR noise, not a building | |
| # ----- geometry / GSD ------------------------------------------- | |
| tile_size: int = 512 # model input, multiple of patch(16) | |
| canonical_gsd_m: float = 0.5 # the GSD the model works at | |
| gsd_jitter_lo_m: float = 0.30 # train-time scale augmentation range | |
| gsd_jitter_hi_m: float = 1.20 | |
| gsd_jitter_p: float = 0.9 | |
| # ----- radiometry / domain gap ---------------------------------- | |
| # v2 trained on raw GAMUS bytes with zero photometric augmentation and then | |
| # hallucinated ~4 m of height on flat ground in a different sensor's imagery | |
| # (Inria Austin). Both knobs below exist to close that gap, and the stretch | |
| # is part of the *inference contract* — it runs identically in predict.py. | |
| radiometric_stretch: bool = True # per-scene 2/98-percentile stretch | |
| stretch_lo_pct: float = 2.0 | |
| stretch_hi_pct: float = 98.0 | |
| photo_p: float = 0.9 # prob. a train crop gets photometric jitter | |
| photo_brightness: float = 0.25 | |
| photo_contrast: float = 0.25 | |
| photo_saturation: float = 0.35 | |
| photo_gamma: float = 0.35 # gamma ~ exp(U(-g, g)) | |
| photo_channel_gain: float = 0.12 # per-channel multiplicative gain | |
| photo_blur_p: float = 0.2 | |
| photo_noise_std: float = 0.015 | |
| # ----- model ----------------------------------------------------- | |
| encoder_model_id: str = "facebook/dinov3-vitl16-pretrain-sat493m" | |
| encoder_feature_indices: tuple = (6, 12, 18, 24) | |
| decoder_dim: int = 256 | |
| detail_branch: bool = False # restored from the checkpoint | |
| detail_dim: int = 64 | |
| n_bins: int = 96 | |
| bin_min_m: float = 0.0 | |
| bin_max_m: float = 120.0 | |
| # ----- schedule -------------------------------------------------- | |
| epochs: int = 26 | |
| freeze_epochs: int = 2 # encoder frozen for this many epochs | |
| max_minutes: float = 330.0 # hard wall-clock cap for the whole run | |
| learning_rate: float = 3e-4 # decoder + heads | |
| encoder_lr: float = 6e-5 # top encoder block; decayed downward | |
| llrd: float = 0.80 # layer-wise LR decay per block | |
| weight_decay: float = 0.05 | |
| warmup_frac: float = 0.05 | |
| grad_clip: float = 1.0 | |
| # Sized for a 48 GB L40S running ViT-L/16 at 512 px with the encoder | |
| # UNFROZEN and no gradient checkpointing: ~1.05 GB of activations per | |
| # sample plus ~5 GB of weights/grads/AdamW state. `train_L40S.sh` probes | |
| # the real number on the real card at startup and overrides this; the | |
| # default is the conservative value that is known to fit. | |
| batch_size: int = 24 | |
| grad_accum: int = 2 # effective batch 48 | |
| eval_batch_mult: int = 2 # eval keeps no activations -> 2x batch | |
| num_workers: int = 12 | |
| prefetch_factor: int = 6 | |
| ema_decay: float = 0.9995 # 0 -> disable weight EMA | |
| # ----- precision / perf ----------------------------------------- | |
| amp: bool = True | |
| amp_dtype: str = "bf16" # bf16 on Hopper | |
| channels_last: bool = True | |
| grad_checkpoint_encoder: bool = False # trades ~35% throughput for VRAM | |
| grad_checkpoint_decoder: bool = False | |
| compile_model: bool = False # recompiles on unfreeze; opt-in | |
| # Photometric jitter + encoder normalisation run on the GPU over the whole | |
| # batch instead of per crop in a DataLoader worker, and the worker hands | |
| # over uint8 HWC. ~3x the loader throughput and 4x less H2D traffic; see | |
| # dwdata/gpu_aug.py for the measurements. Turn it off to fall back to the | |
| # all-CPU path (identical maths, much slower). | |
| gpu_augment: bool = True | |
| sdp_flash: bool = True # prefer the fused SDPA kernels | |
| cudnn_benchmark: bool = True | |
| tf32: bool = True | |
| matmul_precision: str = "high" | |
| seed: int = 42 | |
| # ----- loss ------------------------------------------------------ | |
| w_l1: float = 1.0 | |
| w_grad: float = 0.5 | |
| w_silog: float = 1.0 | |
| silog_lambda: float = 0.85 | |
| silog_shift: float = 1.0 | |
| w_head_a: float = 0.3 # auxiliary; `fused` is the product | |
| w_head_b: float = 0.3 | |
| w_fused: float = 1.0 | |
| w_bin_ce: float = 0.5 | |
| w_seg: float = 0.2 | |
| w_normal: float = 0.3 # surface-normal loss -> keeps flat things flat | |
| w_flat: float = 0.2 # planarity penalty on ground/road/water | |
| stratum_balance_beta: float = 0.5 # 0 -> plain pixel-uniform loss | |
| stratum_weight_clip: float = 5.0 | |
| # ----- eval / io ------------------------------------------------- | |
| eval_every: int = 2 | |
| tta: bool = True | |
| tta_scales: tuple = (1.0,) # dihedral only by default; add 1.25 for +scale | |
| final_sliding_eval: bool = True # full-tile sliding-window eval at the end | |
| per_class_metrics: bool = True | |
| dump_class_stats: bool = True | |
| class_names: tuple = CLASS_NAMES | |
| output_dir: str = _DEFAULT_OUTPUT_DIR | |
| n_qualitative: int = 12 | |
| # ----- misc ------------------------------------------------------ | |
| hf_token: str = "" | |
| smoke: bool = False | |
| resume: str = "" # checkpoint to resume/warm-start from | |
| make_zip: bool = True | |
| # ----------------------------------------------------------------- | |
| def apply_smoke(self) -> "Config": | |
| if not self.smoke: | |
| return self | |
| self.crops_per_epoch = 24 | |
| self.val_tiles = 4 | |
| self.epochs = 2 | |
| self.freeze_epochs = 1 | |
| self.max_minutes = 20.0 | |
| self.batch_size = 2 | |
| self.grad_accum = 1 | |
| self.eval_batch_mult = 1 | |
| self.num_workers = 0 | |
| self.compile_model = False | |
| self.cudnn_benchmark = False | |
| self.eval_every = 1 | |
| self.n_qualitative = 2 | |
| self.final_sliding_eval = False | |
| self.ema_decay = 0.0 | |
| return self | |
| def validate(self) -> "Config": | |
| assert self.tile_size % 16 == 0, "tile_size must be a multiple of 16" | |
| assert self.amp_dtype in ("bf16", "fp16") | |
| assert self.gsd_jitter_lo_m <= self.gsd_jitter_hi_m | |
| assert 0 <= self.freeze_epochs <= self.epochs | |
| try: | |
| Path(self.output_dir).mkdir(parents=True, exist_ok=True) | |
| except OSError: | |
| pass | |
| return self | |
| def dataset_list(self) -> list[str]: | |
| return [d.strip() for d in self.datasets.split(",") if d.strip()] | |
| def sampler_weight_map(self) -> dict[str, float]: | |
| out: dict[str, float] = {} | |
| for part in self.sampler_weights.split(","): | |
| if ":" in part: | |
| k, v = part.split(":") | |
| out[k.strip()] = float(v) | |
| return out | |
| def parse_config(argv: list[str] | None = None) -> Config: | |
| cfg = Config() | |
| p = argparse.ArgumentParser(description="DepthWizard v3") | |
| for f in fields(cfg): | |
| cur = getattr(cfg, f.name) | |
| if isinstance(cur, tuple): | |
| continue | |
| t = type(cur) | |
| if t is bool: | |
| p.add_argument( | |
| f"--{f.name}", nargs="?", const=True, default=None, | |
| type=lambda x: str(x).lower() in ("1", "true", "yes", "y"), | |
| ) | |
| else: | |
| p.add_argument(f"--{f.name}", type=t) | |
| args, _ = p.parse_known_args(argv) | |
| for k, v in vars(args).items(): | |
| if v is not None: | |
| setattr(cfg, k, v) | |
| return cfg.apply_smoke().validate() | |
| def safe_config_dict(cfg: Config) -> dict: | |
| from dataclasses import asdict | |
| d = asdict(cfg) | |
| d.pop("hf_token", None) | |
| return d | |