| import os |
| import glob |
| import shutil |
| import random |
|
|
| |
| random.seed(42) |
|
|
| H_FEATURE_DIR = r"D:\RockGlacier_Project\harmonized_features" |
| H_MASK_DIR = r"D:\RockGlacier_Project\harmonized_masks" |
|
|
| DATASET_DIR = r"D:\RockGlacier_Project\dataset" |
|
|
| splits = ['train', 'val', 'test'] |
| for s in splits: |
| os.makedirs(os.path.join(DATASET_DIR, s, 'features'), exist_ok=True) |
| os.makedirs(os.path.join(DATASET_DIR, s, 'masks'), exist_ok=True) |
|
|
| feature_files = sorted(glob.glob(os.path.join(H_FEATURE_DIR, 'feature_*.tif'))) |
|
|
| |
| random.shuffle(feature_files) |
|
|
| |
| total = len(feature_files) |
| train_end = int(total * 0.70) |
| val_end = int(total * 0.85) |
|
|
| train_files = feature_files[:train_end] |
| val_files = feature_files[train_end:val_end] |
| test_files = feature_files[val_end:] |
|
|
| def copy_files(files, split_name): |
| for f_path in files: |
| idx = os.path.basename(f_path).split('_')[1].split('.')[0] |
| m_path = os.path.join(H_MASK_DIR, f"mask_{idx}.tif") |
| |
| if os.path.exists(m_path): |
| shutil.copy(f_path, os.path.join(DATASET_DIR, split_name, 'features', f"feature_{idx}.tif")) |
| shutil.copy(m_path, os.path.join(DATASET_DIR, split_name, 'masks', f"mask_{idx}.tif")) |
|
|
| print(f"Total files: {total}") |
| print(f"Train set: {len(train_files)} files") |
| print(f"Validation set: {len(val_files)} files") |
| print(f"Test set: {len(test_files)} files") |
| print("Copying files into dataset directories...") |
|
|
| copy_files(train_files, 'train') |
| copy_files(val_files, 'val') |
| copy_files(test_files, 'test') |
|
|
| print(f"Dataset successfully split and organized in {DATASET_DIR}") |
|
|
