import tensorflow as tf
from tensorflow.keras.models import load_model
from imutils import paths
import numpy as np
import cv2
import pickle
from tratar_captcha import tratar_imagens


tamanho_imagem = (35, 50)

def quebrar_captcha():
    # Importar o modelo e o tradutor
    with open("rotulos_modelo.dat", "rb") as arquivo_tradutor:
        lb = pickle.load(arquivo_tradutor)
    
    modelo = load_model("modelo_treinado.keras")
    tratar_imagens("captcha_teste", pasta_destino="captcha_teste")

    arquivos = paths.list_images("captcha_teste")

    # Posições fixas para o corte de cada caractere
    posicoes_corte = [(6, 40), (56, 90), (111, 144), (161, 197), (214, 250), (265, 297)]

    resultados = []  # Lista para armazenar os resultados de todos os captchas

    for arquivo in arquivos:
        imagem = cv2.imread(arquivo, cv2.IMREAD_GRAYSCALE)

        # Aplicar um limiar para binarizar a imagem
        _, imagem_bin = cv2.threshold(imagem, 128, 255, cv2.THRESH_BINARY_INV)

        imagem_final = cv2.merge([imagem_bin] * 3)  # Convertendo de volta para RGB para desenhar os retângulos coloridos

        # Processar cada região de corte definida
        previsao = []
        for x1, x2 in posicoes_corte:
            imagem_caractere = imagem_bin[20:-15, x1:x2]  # Cortar a região do caractere
            imagem_caractere = cv2.bitwise_not(imagem_caractere)
            
            imagem_caractere = cv2.resize(imagem_caractere, tamanho_imagem)

            imagem_caractere = np.expand_dims(imagem_caractere, axis=2)
            imagem_caractere = np.expand_dims(imagem_caractere, axis=0)

            letra_prevista = modelo.predict(imagem_caractere)
            letra_prevista = lb.inverse_transform(letra_prevista)[0]
            previsao.append(letra_prevista)

        texto_previsao = "".join(previsao)
        resultados.append(texto_previsao)  # Armazenar resultado na lista

    for resultado in resultados:
        print(resultado)  # Imprimir todos os resultados

    return resultados  # Retornar a lista de resultados

if __name__ == "__main__":
    quebrar_captcha()