File size: 3,182 Bytes
ce209f5 | 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 | """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)
|