from keras.datasets import mnist
import cv2
import os
from tqdm import tqdm

from config import *

(x_train, y_train), (x_test, y_test) = mnist.load_data()

def extract(x,y,dir_path):
    for index in tqdm(range(x.shape[0]),desc="extracting..."):
        image = x[index]
        label = y[index]
        file_name = str(index) + ".jpg"
        target_dir_path = os.path.join(dir_path, str(label))
        if os.path.exists(target_dir_path) == False: 
            os.makedirs(target_dir_path)
        
        cv2.imwrite(os.path.join(target_dir_path, file_name), image)

train_dir_path = os.path.join(DATASET_PATH, "train")
extract(x_train, y_train, train_dir_path)

test_dir_path = os.path.join(DATASET_PATH, "test")
extract(x_test, y_test, test_dir_path)