input.py 文件源码

def extract_data(path):
    global CLASSES
    images, labels = traverse_dir(path)
    images = np.array(images)
    # change to ont-hot vector
    one_hot = preprocessing.LabelBinarizer()
    one_hot.fit(labels)
    nb_classes = len(one_hot.classes_)

    with open(path+'\labels.txt', 'w') as f:
        for label in one_hot.classes_:
            f.write(label + '\n')

    one_hots = list(one_hot.transform([i]) for i in labels)
    one_hots = np.array(one_hots)
    one_hots = np.reshape(one_hots, (images.shape[0], nb_classes))

    return images, one_hots, nb_classes