import os, random
from shutil import copy, rmtree
import fnmatch

TARGET_NUM = # number of jpg files for each digit to extract from the source folder (integer)
SOURCE_DIR = "" # path to the source folder with jpg files
DEST_DIR = "" # path to the destination folder, where the TARGET_NUM amount of jpg files will be stored
DEST_DIR = DEST_DIR + str(TARGET_NUM) + "/"
CATEGORIES = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]

if os.path.exists(DEST_DIR):
    rmtree(DEST_DIR)
    
os.mkdir(DEST_DIR)
    
    

for category in CATEGORIES:
    source_category = os.path.join(SOURCE_DIR, category)    
    if len(fnmatch.filter(os.listdir(source_category), '*.jpg')) < TARGET_NUM:
        print("The category " + category + " contains less than " + str(TARGET_NUM) + " images!\n\r")
        exit()
        
for category in CATEGORIES:
    source_category = os.path.join(SOURCE_DIR, category)
    dest_category = os.path.join(DEST_DIR, category) 
    os.mkdir(dest_category)
    
    for i in range(0, TARGET_NUM):
        file = random.choice(os.listdir(source_category))
        source_path = os.path.join(source_category, file)
        dest_path = os.path.join(dest_category, file)
        copy(source_path, dest_path)
    
    while len(fnmatch.filter(os.listdir(dest_category), '*.jpg')) < TARGET_NUM:
        file = random.choice(os.listdir(source_category))
        source_path = os.path.join(source_category, file)
        dest_path = os.path.join(dest_category, file)
        copy(source_path, dest_path)