import numpy as np
import keras
from keras import layers
import cv2
import os
from sklearn.utils import shuffle
dataset_path = 'C:/Users/Radim/Desktop/Image processing/dataset/numbers-selected/'

"""**Custom dataset load**"""

# Model / data parameters
num_classes = 10
input_shape = (28, 28, 1)

x_dataset = []
y_dataset = []

for dataset in os.listdir(dataset_path):
#for dataset in ['0027_PL3M','0028_PL3F','0029_PL3M']:
  for number in os.listdir(dataset_path + dataset + "/"):
    print("loading: " + str(dataset) + " number " + str(number))
    source_path = dataset_path + dataset + "/" + str(number) + "/"
    for f in os.listdir(source_path):
      img_path = source_path + f
      if os.path.isfile(img_path):
        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
        #img = 255-img
        x_dataset.append(img)
        y_dataset.append(int(number))

# convert to numpy arrays
x_dataset = np.array(x_dataset)
y_dataset = np.array(y_dataset)
# shuffle
x_dataset, y_dataset = shuffle(x_dataset, y_dataset)
# Scale images to the [0, 1] range
x_dataset = x_dataset.astype("float32") / 255
# Make sure images have shape (28, 28, 1)
x_dataset = np.expand_dims(x_dataset, -1)
print("x_dataset shape:", x_dataset.shape)
print(x_dataset.shape[0], "dataset samples")
# convert class vectors to binary class matrices
y_dataset = keras.utils.to_categorical(y_dataset, num_classes)

"""**ML Model**"""

model = keras.Sequential(
    [
        keras.Input(shape=input_shape),
        layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Flatten(),
        layers.Dropout(0.5),
        layers.Dense(num_classes, activation="softmax"),
    ]
)

model.summary()

"""**Train the model**"""

batch_size = 128
epochs = 10
model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
model.fit(x_dataset, y_dataset, batch_size=batch_size, epochs=epochs, validation_split=0.2)

model.save('C:/Users/Radim/Desktop/Image processing/model.keras')