| """ |
| 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'] |
| |
| |
| feature_files = sorted(glob.glob(os.path.join(feature_dir, 'feature_*.tif'))) |
| |
| |
| 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.") |
| |
| |
| 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)}") |
| |
| |
| 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() |
|
|
