| """ |
| Central asset path resolution for the EchMo radar echo segmentation app. |
| |
| Why this module exists |
| ---------------------- |
| Model weights and datasets are not stored in the code repository; they are |
| pulled from Hugging Face. Those paths used to be hard-coded in `app.py` |
| (a `MODEL_PATHS` dict plus several `if dataset_name == ...` branches), so any |
| retrain that renamed a run directory -- or any new dataset -- required editing |
| the application code. |
| |
| Now everything dataset-specific is *discovered from disk*: |
| |
| models/<dataset>/<run>/config.json -> run folder + "model" field |
| models/<dataset>/<run>/results/<model>.pt -> weights (filename == model) |
| datasets/<dataset>/dataset_meta.json -> per-dataset behaviour |
| |
| `datasets/<dataset>/dataset_meta.json` fields (all optional) |
| ----------------------------------------------------------- |
| display_name str label shown in the UI (default: folder name) |
| full_name str paper abbreviation / expansion (informational) |
| scenario str short description (informational) |
| order int sort order of the dataset dropdown |
| mtd_fftshift bool apply fftshift after the Doppler FFT in MTD |
| mtd_amp_dir str sequence sub-folder holding the MTD amplitude .npy |
| valid_classes list classes offered by the class filter in the UI |
| run_overrides dict {"model": "run_folder"} to pin a specific run when a |
| dataset contains several runs of the same model |
| |
| IMPORTANT: `valid_classes` is *not* `nb_classes`. `nb_classes` comes from the |
| model `config.json` and is the number of output channels used for the |
| segmentation colour scale. HiL-RED, for instance, trains with 9 channels while |
| only classes 0-6 actually occur in the annotations. |
| |
| This module uses the standard library only, so it can be imported by scripts |
| without pulling in torch/gradio. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
|
|
| |
| |
| CUSTOM_DATASET = "Custom Data" |
|
|
| DATASETS_ROOT = "datasets" |
| MODELS_ROOT = "models" |
|
|
| INDEX_FILENAME = "light_dataset_frame_oriented.json" |
| META_FILENAME = "dataset_meta.json" |
| COEFS_DIRNAME = "coefs" |
| TRANSMIT_FILENAME = "s_tx.mat" |
|
|
| DEFAULT_META = { |
| "mtd_fftshift": True, |
| "mtd_amp_dir": "range_doppler_MTD", |
| "valid_classes": [0, 1], |
| "run_overrides": {}, |
| } |
|
|
| _meta_cache: dict[str, dict] = {} |
|
|
|
|
| |
| |
| |
| def read_json(path: str, default=None): |
| |
| |
| |
| |
| try: |
| with open(path, "r", encoding="utf-8-sig") as handle: |
| return json.load(handle) |
| except (OSError, ValueError): |
| return default |
|
|
|
|
| |
| |
| |
| def dataset_dir(dataset: str) -> str: |
| return os.path.join(DATASETS_ROOT, dataset) |
|
|
|
|
| def index_path(dataset: str) -> str: |
| return os.path.join(dataset_dir(dataset), INDEX_FILENAME) |
|
|
|
|
| def meta_path(dataset: str) -> str: |
| return os.path.join(dataset_dir(dataset), META_FILENAME) |
|
|
|
|
| def coefs_dir(dataset: str) -> str: |
| return os.path.join(dataset_dir(dataset), COEFS_DIRNAME) |
|
|
|
|
| def transmit_signal_path(dataset: str) -> str: |
| return os.path.join(coefs_dir(dataset), TRANSMIT_FILENAME) |
|
|
|
|
| def load_dataset_meta(dataset: str) -> dict: |
| """Merged dataset_meta.json (cached). Unknown datasets fall back to defaults.""" |
| if dataset not in _meta_cache: |
| meta = dict(DEFAULT_META) |
| loaded = read_json(meta_path(dataset), {}) |
| if isinstance(loaded, dict): |
| meta.update(loaded) |
| meta.setdefault("display_name", dataset) |
| _meta_cache[dataset] = meta |
| return _meta_cache[dataset] |
|
|
|
|
| def clear_meta_cache() -> None: |
| _meta_cache.clear() |
|
|
|
|
| def available_datasets() -> list[str]: |
| """Dataset folders that actually have a frame index, in display order.""" |
| if not os.path.isdir(DATASETS_ROOT): |
| return [] |
| found = [] |
| for name in os.listdir(DATASETS_ROOT): |
| if name == CUSTOM_DATASET or not os.path.isdir(dataset_dir(name)): |
| continue |
| if not os.path.isfile(index_path(name)): |
| continue |
| found.append((load_dataset_meta(name).get("order", 10_000), name)) |
| found.sort(key=lambda item: (item[0], item[1])) |
| return [name for _, name in found] |
|
|
|
|
| def load_dataset_index(dataset: str) -> dict: |
| """{sequence: [frames]} -- {} when missing or not valid strict JSON.""" |
| data = read_json(index_path(dataset), None) |
| return data if isinstance(data, dict) else {} |
|
|
|
|
| def get_sequences(dataset: str) -> list[str]: |
| if dataset == CUSTOM_DATASET: |
| return [] |
| return list(load_dataset_index(dataset).keys()) |
|
|
|
|
| def get_frames(dataset: str, sequence: str) -> list: |
| if dataset == CUSTOM_DATASET or not sequence: |
| return [] |
| return load_dataset_index(dataset).get(sequence, []) |
|
|
|
|
| def echo_path(dataset: str, sequence: str, frame: str) -> str: |
| return os.path.join(dataset_dir(dataset), sequence, "range_doppler_echo", f"{frame}.npy") |
|
|
|
|
| def annotation_path(dataset: str, sequence: str, frame: str) -> str: |
| return os.path.join(dataset_dir(dataset), sequence, "annotations", "dense", frame, "range_doppler.npy") |
|
|
|
|
| def mtd_amp_path(dataset: str, sequence: str, frame: str) -> str: |
| return os.path.join(dataset_dir(dataset), sequence, load_dataset_meta(dataset)["mtd_amp_dir"], f"{frame}.npy") |
|
|
|
|
| |
| |
| |
| def runs_dir(dataset: str) -> str: |
| return os.path.join(MODELS_ROOT, dataset) |
|
|
|
|
| def run_dir(dataset: str, run_folder: str) -> str: |
| return os.path.join(runs_dir(dataset), run_folder) |
|
|
|
|
| def run_config_path(dataset: str, run_folder: str) -> str: |
| return os.path.join(run_dir(dataset, run_folder), "config.json") |
|
|
|
|
| def run_weights_path(dataset: str, run_folder: str, filename: str) -> str: |
| return os.path.join(run_dir(dataset, run_folder), "results", filename) |
|
|
|
|
| def load_run_config(dataset: str, run_folder: str) -> dict: |
| data = read_json(run_config_path(dataset, run_folder), None) |
| return data if isinstance(data, dict) else {} |
|
|
|
|
| def discover_runs(dataset: str) -> dict[str, tuple[str, str]]: |
| """ |
| {model_name: (run_folder, weights_filename)} discovered from disk. |
| |
| A run folder qualifies when its config.json declares `"model"` and |
| `results/<model>.pt` exists. When several runs train the same model the |
| alphabetically first folder wins, unless `run_overrides` in |
| dataset_meta.json pins a specific one. |
| """ |
| runs: dict[str, tuple[str, str]] = {} |
| if not os.path.isdir(runs_dir(dataset)): |
| return runs |
|
|
| for folder in sorted(os.listdir(runs_dir(dataset))): |
| if not os.path.isdir(run_dir(dataset, folder)): |
| continue |
| cfg = load_run_config(dataset, folder) |
| model = cfg.get("model") or cfg.get("model_name") |
| if not model: |
| continue |
| filename = f"{model}.pt" |
| if not os.path.isfile(run_weights_path(dataset, folder, filename)): |
| continue |
| runs.setdefault(model, (folder, filename)) |
|
|
| for model, folder in (load_dataset_meta(dataset).get("run_overrides") or {}).items(): |
| filename = f"{model}.pt" |
| if os.path.isfile(run_weights_path(dataset, folder, filename)): |
| runs[model] = (folder, filename) |
| else: |
| print(f"[warn] run_overrides: '{folder}' for model '{model}' in '{dataset}' is not usable") |
|
|
| return runs |
|
|
|
|
| def model_names(dataset: str) -> list[str]: |
| return sorted(discover_runs(dataset)) |
|
|
