sample.py

import argparse
import os
import random
import shutil

def safe_mkdir(*args, **kwds):
    try:
        return os.mkdir(*args, **kwds)
    except FileExistsError:
        pass
    

parser = argparse.ArgumentParser(description='Choose a random sample of images')
parser.add_argument('--output', default='./sample',
                    help='the destination directory')
parser.add_argument('--stoats', type=int, required=True,
                    help='the number of stoat images to select')
parser.add_argument('--rats', type=int, required=True,
                    help='the number of rat images to select')
parser.add_argument('--possums', type=int, required=True,
                    help='the number of possum images to select')
parser.add_argument('--others', type=int, required=True,
                    help='the number of non-pest images to select')

args = vars(parser.parse_args())
out_dir = args['output']

animals = ['rat', 'stoat', 'possum', 'other']
sets = {}
for animal in animals:
    paths = set()
    top_level_dir = './images_large_{}/'.format(animal)
    for sub_dir in os.listdir(top_level_dir):
        path = '{}{}/'.format(top_level_dir, sub_dir)
        paths.update('{}{}'.format(path, image) for image in os.listdir(path))
    sets[animal] = paths

safe_mkdir(out_dir)
samples = { animal: random.sample(sets[animal], args[animal + 's']) for animal in animals }
for animal, sample in samples.items():
    sub_dir = '{}/{}'.format(out_dir, animal)
    safe_mkdir(sub_dir)
    for path in sample:
        shutil.copy(path, sub_dir)

sample.sh

python3 sample.py --stoats=50 --rats=50 --possums=50 --others=50 --output=train
python3 sample.py --stoats=25 --rats=25 --possums=25 --others=25 --output=test
python3 sample.py --stoats=25 --rats=25 --possums=25 --others=25 --output=valid