File size: 2,938 Bytes
80e070e | 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 | """
Split harmonized feature/mask pairs into train/val/test sets.
Copies matched pairs into output/{train,val,test}/{features,masks}/.
Usage:
python split_dataset.py
"""
import os
import glob
import yaml
import shutil
import random
def load_config():
config_path = os.path.join(os.path.dirname(__file__), 'config.yaml')
with open(config_path, 'r') as f:
return yaml.safe_load(f)
def main():
config = load_config()
base_dir = os.path.join(os.path.dirname(__file__), config['paths']['output_dir'])
feature_dir = os.path.join(base_dir, 'features')
mask_dir = os.path.join(base_dir, 'masks')
dataset_dir = os.path.join(base_dir, 'dataset')
train_r = config['splitting']['train_ratio']
val_r = config['splitting']['val_ratio']
seed = config['splitting']['random_seed']
# Find all matched pairs
feature_files = sorted(glob.glob(os.path.join(feature_dir, 'feature_*.tif')))
# Extract IDs and verify masks exist
paired_ids = []
for fp in feature_files:
idx_str = os.path.basename(fp).split('_')[1].split('.')[0]
mp = os.path.join(mask_dir, f"mask_{idx_str}.tif")
if os.path.exists(mp):
paired_ids.append(idx_str)
else:
print(f"Warning: No mask for feature_{idx_str}.tif, skipping.")
print(f"Found {len(paired_ids)} matched feature-mask pairs.")
# Shuffle and split
random.seed(seed)
random.shuffle(paired_ids)
n = len(paired_ids)
n_train = int(n * train_r)
n_val = int(n * (train_r + val_r)) - n_train
train_ids = paired_ids[:n_train]
val_ids = paired_ids[n_train:n_train + n_val]
test_ids = paired_ids[n_train + n_val:]
splits = {'train': train_ids, 'val': val_ids, 'test': test_ids}
for split_name, ids in splits.items():
feat_out = os.path.join(dataset_dir, split_name, 'features')
mask_out = os.path.join(dataset_dir, split_name, 'masks')
os.makedirs(feat_out, exist_ok=True)
os.makedirs(mask_out, exist_ok=True)
for idx_str in ids:
src_feat = os.path.join(feature_dir, f"feature_{idx_str}.tif")
src_mask = os.path.join(mask_dir, f"mask_{idx_str}.tif")
shutil.copy2(src_feat, feat_out)
shutil.copy2(src_mask, mask_out)
print(f" {split_name}: {len(ids)} pairs")
print(f"\nSplit complete! Dataset saved to {dataset_dir}")
print(f" Train: {len(train_ids)} | Val: {len(val_ids)} | Test: {len(test_ids)}")
# Leak check
tr, va, te = set(train_ids), set(val_ids), set(test_ids)
assert len(tr & va) == 0, "LEAK: train ∩ val is not empty!"
assert len(tr & te) == 0, "LEAK: train ∩ test is not empty!"
assert len(va & te) == 0, "LEAK: val ∩ test is not empty!"
print("Leak check PASSED — no overlap between splits.")
if __name__ == '__main__':
main()
|