File size: 1,653 Bytes
820fab0 | 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 | 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()
|