File size: 1,717 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 | import os
import glob
import shutil
import random
# Set a fixed seed for reproducibility so the split is always the same
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')))
# Shuffle files
random.shuffle(feature_files)
# Calculate split indices (70% train, 15% val, 15% test)
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}")
|