import numpy as np
from pathlib import Path
from os.path import basename
import matplotlib.image as mpimg
import curses

"""
This script read all files recursively in given directory, create a class of
each dir, and put the images (normalized numpy arrays) in a npy binary format
in the form: (num_classes, num_exemple, h, w, c)
"""

def create_db(images_dir_path):
    if not Path(train_dir).is_dir():
        print('Error')
        exit()

    dataset = [] # need (num_classes, num_exemple, h, w, c)
    for d in Path(images_dir_path).glob('*'): # for files/dir in pathdir
        # skipping files
        if not d.is_dir():
            continue;

        classDataSet = []
        for f in Path(d).glob('*.png'):
            img = mpimg.imread(f)
            img = img.astype(np.float)
            img /= 255.0
            img = np.reshape(img, newshape=(img.shape[0], img.shape[1], 1))
            classDataSet.append(img) # list numpy images

        dataset.append(classDataSet)

    dataset.sort(key=len)

    return np.array(dataset)


train_dir = '/path/to/my/files'
data = create_db(train_dir)
np.save('my_database.npy', data)