import cv2
import os
import numpy as np
import pickle
from imutils import paths
from sklearn.preprocessing import LabelBinarizer
from sklearn.model_selection import train_test_split
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense

dados = []
rotulos = []
pasta_base_imagens = 'base_letras'
tamanho_imagem = (35, 50)

imagens = paths.list_images(pasta_base_imagens)

for arquivo in imagens:
    rotulo = arquivo.split(os.path.sep)[-2]
    imagem = cv2.imread(arquivo)
    imagem = cv2.cvtColor(imagem, cv2.COLOR_BGR2GRAY)

    imagem = cv2.resize(imagem, tamanho_imagem)

    # Adicionar uma dimensao para o Keras ler a Imagem
    imagem = np.expand_dims(imagem, axis=2)

    # Adicionar as listas de Dados e Rotulos
    rotulos.append(rotulo)
    dados.append(imagem)


dados = np.array(dados, dtype="float") / 255
rotulos = np.array(rotulos)

# Separação em treino e teste

(x_train, x_test, y_train, y_test) = train_test_split(dados, rotulos, test_size=0.25, random_state=0)

lb = LabelBinarizer().fit(y_train)
y_train = lb.transform(y_train)
y_test = lb.transform(y_test)

# Salvar o LabelBinarizer em um arquivo com Pickle
with open('rotulos_modelo.dat', 'wb') as arquivo_pickle:
    pickle.dump(lb, arquivo_pickle)

# Criar IA
modelo = Sequential()

modelo.add(Conv2D(20, (5, 5), padding="same", input_shape=(50, 35, 1), activation="relu"))
modelo.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

modelo.add(Conv2D(50, (5, 5), padding="same", activation="relu"))
modelo.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

modelo.add(Flatten())
modelo.add(Dense(500, activation="relu"))

modelo.add(Dense(26, activation="softmax"))

modelo.compile(loss="categorical_crossentropy", optimizer="adam", metrics=['accuracy'])

# Treinar
modelo.fit(x_train, y_train, validation_data=(x_test, y_test), batch_size=26, epochs=3, verbose=1)

# Salvar o Modelo
modelo.save("modelo_treinado.keras")