File size: 1,700 Bytes
820fab0 | 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 | """
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.")
|