| import os |
| import glob |
| import shutil |
| import random |
|
|
| PATCH_DIR = "dataset/extracted_patches" |
| DATASET_DIR = "dataset/patches_10m" |
|
|
| def main(): |
| random.seed(42) |
| |
| features = sorted(glob.glob(os.path.join(PATCH_DIR, "features", "*.npy"))) |
| if not features: |
| print(f"No patches found in {PATCH_DIR}") |
| return |
| |
| random.shuffle(features) |
| |
| total = len(features) |
| train_end = int(total * 0.70) |
| val_end = int(total * 0.85) |
| |
| train_files = features[:train_end] |
| val_files = features[train_end:val_end] |
| test_files = features[val_end:] |
| |
| for split in ['train', 'val', 'test']: |
| os.makedirs(os.path.join(DATASET_DIR, split, 'features'), exist_ok=True) |
| os.makedirs(os.path.join(DATASET_DIR, split, 'masks'), exist_ok=True) |
| |
| print(f"Total patches: {total}") |
| print(f"Train set: {len(train_files)} patches") |
| print(f"Validation set: {len(val_files)} patches") |
| print(f"Test set: {len(test_files)} patches") |
| |
| print("Copying to split directories...") |
| |
| def copy_split(files, split_name): |
| for f in files: |
| name = os.path.basename(f) |
| m = os.path.join(PATCH_DIR, "masks", name) |
| |
| shutil.copy(f, os.path.join(DATASET_DIR, split_name, 'features', name)) |
| if os.path.exists(m): |
| shutil.copy(m, os.path.join(DATASET_DIR, split_name, 'masks', name)) |
| |
| copy_split(train_files, 'train') |
| copy_split(val_files, 'val') |
| copy_split(test_files, 'test') |
| |
| print(f"Split successful. Dataset is at {DATASET_DIR}") |
|
|
| if __name__ == '__main__': |
| main() |
|
|
