File size: 8,924 Bytes
0e150d6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 | """Shared model names, source locations, and output conventions."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
@dataclass(frozen=True)
class InferenceModel:
"""One released model that accepts the packaged Smoke-Eval layout."""
entrypoint: str
accelerate_backend: str
config: str
metric_variant: str
prediction_directory: str
data_key: tuple[str, ...]
output_key: tuple[str, ...]
produced_directory: str
backend_args: tuple[str, ...] = ("--config", "{config}")
# ``grade`` is the paper name for the complete model. The evaluator keeps the
# historical ``ours_full`` merged-CSV name. Backend placeholders are rendered
# by evaluation/run_inference.py after it creates the run-specific YAML.
INFERENCE_MODELS: dict[str, InferenceModel] = {
# Complete GRADE and its Stage-1/Stage-2 component ablations.
"grade": InferenceModel(
"src/models/grade/inference.py",
"src/GRADE/stage2_diffusion_refinement/inference_full.py",
"src/models/grade/config.yaml", "ours_full", "ours_full",
("data", "smoke_eval_root"), ("inference", "output_root"), "ours_full",
),
"ours_full": InferenceModel(
"src/models/ours_full/inference.py",
"src/GRADE/stage2_diffusion_refinement/inference_full.py",
"src/models/ours_full/config.yaml", "ours_full", "ours_full",
("data", "smoke_eval_root"), ("inference", "output_root"), "ours_full",
),
"ours_diffusion": InferenceModel(
"src/models/ours_diffusion/inference.py",
"src/GRADE/stage2_diffusion_refinement/inference_diffusion.py",
"src/models/ours_diffusion/config.yaml", "ours_diffusion", "ours_diffusion",
("data", "smoke_eval_root"), ("inference", "output_root"), "ours_diffusion",
),
"ours_radar": InferenceModel(
"src/models/ours_radar/inference.py",
"src/GRADE/stage2_diffusion_refinement/inference_radar.py",
"src/models/ours_radar/config.yaml", "ours_radar", "ours_radar",
("data", "smoke_eval_root"), ("inference", "output_root"), "ours_radar",
),
"ours_radar_no_grad": InferenceModel(
"src/models/ours_radar_no_grad/inference.py",
# This is a Stage-1-only ablation. Its wrapper invokes all three
# GRADE stages, so the direct backend must be the released radar-stage
# entry point.
"src/GRADE/stage2_diffusion_refinement/inference_radar.py",
"src/models/ours_radar_no_grad/config.yaml", "ours_radar_no_grad",
"ours_radar_no_grad", ("data", "smoke_eval_root"),
("inference", "output_root"), "ours_radar",
),
"ours_radar_no_doppler": InferenceModel(
"src/models/ours_radar_no_doppler/inference.py",
"src/Ablation/ours_radar_no_doppler/inference.py",
"src/models/ours_radar_no_doppler/config.yaml", "ours_radar_no_doppler",
"ours_radar_no_doppler", ("data", "test_root"),
("inference", "output_dir"), ".",
),
"ours_full_no_3d": InferenceModel(
"src/models/ours_full_no_3d/inference.py",
"src/Ablation/ours_full_no_3d/inference.py",
"src/models/ours_full_no_3d/config.yaml", "ours_full_no_3d",
"ours_full_no_3d", ("data", "smoke_eval_root"),
("inference", "output_root"), "ours_full",
),
# Baselines.
"da3": InferenceModel(
"src/models/da3/inference.py", "src/Baselines/da3/inference.py",
"src/models/da3/config.yaml", "da3", "da3",
("data", "smoke_eval_root"), (), ".",
(
"--data_root", "{data_root}",
"--checkpoint", "{artifact_root}/checkpoints/baselines/da3/da3metric-large.safetensors",
"--output_dir", "{run_root}",
),
),
"grt": InferenceModel(
"src/models/grt/inference.py", "src/Baselines/grt/inference.py",
"src/models/grt/config.yaml", "grt", "grt", ("paths", "data_root"), (), ".",
(
"--config", "{config}",
"--checkpoint", "{artifact_root}/checkpoints/baselines/grt/grt.safetensors",
"--output_dir", "{run_root}",
),
),
"grt_image": InferenceModel(
"src/models/grt_image/inference.py", "src/Baselines/grt_image/inference.py",
"src/models/grt_image/config.yaml", "grt_image", "grt_image",
("paths", "smoke_eval_root"), (), ".",
(
"--config", "{config}",
"--checkpoint", "{artifact_root}/checkpoints/baselines/grt_image/grt_image.safetensors",
"--output_dir", "{run_root}",
),
),
"cafnet": InferenceModel(
"src/models/cafnet/inference.py", "src/Baselines/cafnet/inference.py",
"src/models/cafnet/config.yaml", "cafnet", "cafnet",
("test_base_dir",), ("prediction_dir",), ".",
),
"cafnet_no_smoke": InferenceModel(
"src/models/cafnet_no_smoke/inference.py", "src/Baselines/cafnet_no_smoke/inference.py",
"src/models/cafnet_no_smoke/config.yaml", "cafnet_no_smoke", "cafnet_no_smoke",
("test_base_dir",), ("prediction_dir",), ".",
),
"radarcam-depth": InferenceModel(
"src/models/radarcam-depth/inference.py",
"src/Baselines/radarcam-depth/smoke_eval_inference.py",
"src/models/radarcam-depth/config.yaml", "radarcam-depth", "radarcam-depth",
(), (), ".",
(
"--config", "{config}", "--smoke_root", "{data_root}",
"--rcnet_checkpoint", "{artifact_root}/checkpoints/baselines/radarcam-depth/radarcam-depth_rcnet.safetensors",
"--sml_checkpoint", "{artifact_root}/checkpoints/baselines/radarcam-depth/radarcam-depth_sml.safetensors",
"--output_dir", "{run_root}",
),
),
# Ablations based on GRT or a modified GRADE loss/input.
"grt_refine_freeze": InferenceModel(
"src/models/grt_refine_freeze/inference.py",
"src/Ablation/grt_refine_freeze/inference_control.py",
"src/models/grt_refine_freeze/config.yaml", "grt_refine_freeze",
"grt_refine_freeze", ("data", "smoke_eval_root"),
("inference", "output_dir"), ".",
),
"grt_refine_retrain": InferenceModel(
"src/models/grt_refine_retrain/inference.py",
"src/Ablation/grt_refine_retrain/inference_control.py",
"src/models/grt_refine_retrain/config.yaml", "grt_refine_retrain", "grt_refine_retrain",
("data", "smoke_eval_root"), ("inference", "output_dir"), ".",
),
"grt_no_doppler": InferenceModel(
"src/models/grt_no_doppler/inference.py",
"src/Ablation/grt_no_doppler/inference.py",
"src/models/grt_no_doppler/config.yaml", "grt_no_doppler",
"grt_no_doppler", ("paths", "test_data_root"),
("inference", "output_dir"), ".",
),
}
# Canonical identifiers are used by all folders, raw metric directories, and
# merged CSV filenames. Old evaluation/paper keys are accepted only when a
# reviewer supplies one explicitly.
CANONICAL_METRIC_VARIANTS = frozenset(
{
"da3", "cafnet", "cafnet_no_smoke", "radarcam-depth", "grt",
"grt_refine_freeze", "grt_no_doppler", "grt_cafnet", "grt_image",
"grt_refine_retrain", "ours_full", "ours_full_no_3d",
"ours_diffusion", "ours_radar", "ours_radar_no_doppler",
"ours_radar_no_grad",
}
)
MODEL_NAME_ALIASES = {
"DA3": "da3",
"CafNet": "cafnet",
"CafNet_filter": "cafnet_no_smoke",
"RadarCam_Rice": "radarcam-depth",
"radarcam_depth": "radarcam-depth",
"GRT": "grt",
"GRT_grade_freeze": "grt_refine_freeze",
"grt_grade": "grt_refine_retrain",
"GRT_Image": "grt_image",
"GRT_no_doppler": "grt_no_doppler",
"ours_radar_no_gradient": "ours_radar_no_grad",
}
# Support lowercased historical spelling such as ``cafnet_filter`` too.
MODEL_NAME_ALIASES.update(
{
legacy.lower(): canonical
for legacy, canonical in tuple(MODEL_NAME_ALIASES.items())
}
)
EVALUATION_VARIANTS = frozenset(
{*CANONICAL_METRIC_VARIANTS, *MODEL_NAME_ALIASES}
)
def canonical_model_name(name: str) -> str:
"""Return the canonical artifact identifier for a public/legacy name."""
return MODEL_NAME_ALIASES.get(name, MODEL_NAME_ALIASES.get(name.lower(), name))
def resolve_metric_variant(name: str) -> str:
"""Translate a public or historical name to the canonical metric name."""
canonical = canonical_model_name(name)
if canonical in INFERENCE_MODELS:
return INFERENCE_MODELS[canonical].metric_variant
if canonical in CANONICAL_METRIC_VARIANTS:
return canonical
choices = sorted({*INFERENCE_MODELS, *EVALUATION_VARIANTS})
raise ValueError(f"Unknown model {name!r}. Choose from: {', '.join(choices)}")
def artifact_path(root: Path, relative: str) -> Path:
"""Resolve a repository-relative path while keeping registry data portable."""
return (root / relative).resolve()
|