utils/dataset_list_generator.py

"""Generate train/val/test list files for list-based change-detection loaders."""

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

IMG_EXTS = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"}


def generate_list_files(
    data_root: str,
    splits: list[str] | None = None,
    img_subdir: str = "A",
    overwrite: bool = False,
) -> dict[str, int]:
    if splits is None:
        splits = ["train", "val", "test"]

    root = Path(data_root)
    list_dir = root / "list"
    list_dir.mkdir(parents=True, exist_ok=True)

    results: dict[str, int] = {}
    for split in splits:
        list_file = list_dir / f"{split}.txt"
        if list_file.exists() and not overwrite:
            count = sum(1 for _ in list_file.open("r", encoding="utf-8"))
            print(f"  [SKIP] {list_file} already exists ({count} entries)")
            results[split] = count
            continue

        img_dir = None
        for candidate in [img_subdir, "A", "T1", "t1", "img", "images", "image"]:
            candidate_path = root / split / candidate
            if candidate_path.is_dir():
                img_dir = candidate_path
                break

        if img_dir is None:
            print(f"  [WARN] No image directory found under {root / split} - skipping {split}")
            results[split] = 0
            continue

        filenames = sorted(
            p.name for p in img_dir.iterdir() if p.is_file() and p.suffix.lower() in IMG_EXTS
        )
        if not filenames:
            print(f"  [WARN] No image files found in {img_dir}")
            results[split] = 0
            continue

        list_file.write_text("\n".join(filenames) + "\n", encoding="utf-8")
        print(f"  [OK] {list_file}: {len(filenames)} entries")
        results[split] = len(filenames)

    return results


def generate_for_dataset(dataset_name: str, overwrite: bool = False) -> None:
    from utils.config_loader import load_dataset_config

    cfg = load_dataset_config(dataset_name)
    data_root = cfg["data_root"]
    if not os.path.isdir(data_root):
        print(f"[SKIP] Dataset root does not exist: {data_root}")
        return
    img_subdir = cfg.get("image_a_folder", "A")
    splits = list((cfg.get("splits") or {"train": "train", "val": "val", "test": "test"}).keys())
    print(f"[LIST_GEN] Generating list files for {dataset_name} at {data_root}")
    results = generate_list_files(data_root, splits=splits, img_subdir=img_subdir, overwrite=overwrite)
    print(f"[LIST_GEN] Done: {results}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", type=str, default=None)
    parser.add_argument("--all", action="store_true")
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args()

    from utils.config_loader import list_available_datasets

    datasets = list_available_datasets() if args.all else ([args.dataset] if args.dataset else [])
    for ds in datasets:
        generate_for_dataset(ds, overwrite=args.overwrite)