Spaces:
Runtime error
Runtime error
File size: 2,644 Bytes
3366f95 | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | import argparse
import random
import shutil
import sys
from pathlib import Path
root = Path(__file__).resolve().parents[1]
sys.path.append(str(root))
def split_dataset(source_dir, output_dir, train_ratio=0.7, val_ratio=0.15, test_ratio=0.15, seed=123):
source_dir = Path(source_dir)
output_dir = Path(output_dir)
random.seed(seed)
if not source_dir.exists():
raise FileNotFoundError(f'Source directory not found: {source_dir}')
if abs(train_ratio + val_ratio + test_ratio - 1.0) > 1e-6:
raise ValueError('Train, validation, and test ratios must sum to 1.0.')
for subset in ['train', 'val', 'test']:
subset_dir = output_dir / subset
subset_dir.mkdir(parents=True, exist_ok=True)
for class_dir in sorted(source_dir.iterdir()):
if not class_dir.is_dir():
continue
class_name = class_dir.name
images = [p for p in class_dir.iterdir() if p.is_file()]
random.shuffle(images)
n = len(images)
train_end = int(n * train_ratio)
val_end = train_end + int(n * val_ratio)
splits = {
'train': images[:train_end],
'val': images[train_end:val_end],
'test': images[val_end:],
}
for split_name, files in splits.items():
target_dir = output_dir / split_name / class_name
target_dir.mkdir(parents=True, exist_ok=True)
for file_path in files:
shutil.copy2(file_path, target_dir / file_path.name)
print(f'Successfully split dataset from {source_dir} into {output_dir}')
print('Train/val/test ratios:', train_ratio, val_ratio, test_ratio)
def parse_args():
parser = argparse.ArgumentParser(description='Split raw image dataset into train/val/test folders')
parser.add_argument('--source', required=True, help='Source folder containing class subfolders')
parser.add_argument('--output', default='dataset', help='Output folder to create train/val/test splits')
parser.add_argument('--train_ratio', type=float, default=0.7)
parser.add_argument('--val_ratio', type=float, default=0.15)
parser.add_argument('--test_ratio', type=float, default=0.15)
parser.add_argument('--seed', type=int, default=123)
return parser.parse_args()
def main():
args = parse_args()
split_dataset(
args.source,
args.output,
train_ratio=args.train_ratio,
val_ratio=args.val_ratio,
test_ratio=args.test_ratio,
seed=args.seed,
)
if __name__ == '__main__':
main()
|