A python script to split an image dataset

split.py

#! /usr/bin/env python3

import argparse
from pathlib import Path
from random import shuffle
import shutil

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="""
        Split a dataset in two given a ratio.
    """)
    parser.add_argument('input', type=Path)
    parser.add_argument('output', type=Path)
    parser.add_argument('--ratio', type=float, default=0.66)
    args = parser.parse_args()

    for class_path in args.input.iterdir():
        samples = [p.relative_to(args.input) for p in class_path.iterdir()]
        split = int(len(samples) * args.ratio)
        shuffle(samples)

        for name, dataset in [('train', samples[:split]),
                              ('test', samples[split:])]:
            output_path = args.output / name
            for p in dataset:
                q = output_path / p
                q.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy(str(args.input / p), str(output_path / p))