| """ |
| Step 2 — Split crops into train/val/test BY CROP (so overlapping patches from |
| one crop never straddle two splits), copying them into |
| <out>/dataset/{split}/{features,masks}/. |
| """ |
|
|
| import os |
| import glob |
| import shutil |
| import random |
|
|
|
|
| def split_dataset(cfg): |
| out = cfg["out_path"] |
| feat_dir = os.path.join(out, "features") |
| mask_dir = os.path.join(out, "masks") |
| dataset_dir = os.path.join(out, "dataset") |
| tr_r, va_r, seed = cfg["train_ratio"], cfg["val_ratio"], cfg["seed"] |
|
|
| ids = [] |
| for fp in sorted(glob.glob(os.path.join(feat_dir, "feature_*.tif"))): |
| idx = os.path.basename(fp).split("_")[1].split(".")[0] |
| if os.path.exists(os.path.join(mask_dir, f"mask_{idx}.tif")): |
| ids.append(idx) |
| print(f"[split] matched pairs: {len(ids)}") |
|
|
| random.seed(seed) |
| random.shuffle(ids) |
| n = len(ids) |
| n_tr = int(n * tr_r) |
| n_va = int(n * (tr_r + va_r)) - n_tr |
| splits = {"train": ids[:n_tr], |
| "val": ids[n_tr:n_tr + n_va], |
| "test": ids[n_tr + n_va:]} |
|
|
| for name, sub in splits.items(): |
| fo = os.path.join(dataset_dir, name, "features") |
| mo = os.path.join(dataset_dir, name, "masks") |
| os.makedirs(fo, exist_ok=True) |
| os.makedirs(mo, exist_ok=True) |
| for idx in sub: |
| shutil.copy2(os.path.join(feat_dir, f"feature_{idx}.tif"), fo) |
| shutil.copy2(os.path.join(mask_dir, f"mask_{idx}.tif"), mo) |
| print(f"[split] {name}: {len(sub)}") |
|
|
| tr, va, te = (set(splits[k]) for k in ("train", "val", "test")) |
| assert not (tr & va) and not (tr & te) and not (va & te), "SPLIT LEAK!" |
| print("[split] leak check passed — no crop in two splits.") |
|
|
