import cv2
import os
import glob
import numpy as np
import base64


def tratar_imagens(pasta_origem, pasta_destino='captcha_tratado'):
    arquivos = glob.glob(f'{pasta_origem}/*')
    for arquivo in arquivos:
        imagem = cv2.imread(arquivo)
        imagem_cinza = cv2.cvtColor(imagem, cv2.COLOR_RGB2GRAY)
        _, imagem_tratada = cv2.threshold(imagem_cinza, 127, 255, cv2.THRESH_BINARY or cv2.THRESH_OTSU)

        image = imagem_tratada
        _, binary_image = cv2.threshold(image, 128, 255, cv2.THRESH_BINARY_INV)
        kernel = np.ones((2, 2), np.uint8)
        opened_image = cv2.morphologyEx(binary_image, cv2.MORPH_OPEN, kernel)
        result_image = cv2.bitwise_not(opened_image)
        
        nome_arquivo = os.path.basename(arquivo)
        cv2.imwrite(f'{pasta_destino}/{nome_arquivo}', result_image)

def tratar_base64(imagem_base64):
    imagem_bytes = base64.b64decode(imagem_base64)
    imagem_np = np.frombuffer(imagem_bytes, np.uint8)
    imagem = cv2.imdecode(imagem_np, cv2.IMREAD_COLOR)
    imagem_cinza = cv2.cvtColor(imagem, cv2.COLOR_BGR2GRAY)
    
    _, imagem_tratada = cv2.threshold(imagem_cinza, 127, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    _, imagem_binaria = cv2.threshold(imagem_tratada, 128, 255, cv2.THRESH_BINARY_INV)
    
    kernel = np.ones((2, 2), np.uint8)
    imagem_aberta = cv2.morphologyEx(imagem_binaria, cv2.MORPH_OPEN, kernel)
    
    imagem_resultado = cv2.bitwise_not(imagem_aberta)
    
    return imagem_resultado

if __name__ == "__main__":
    tratar_imagens('captcha_images')