import os
import re
from glob import glob
import cv2
import numpy as np
from click_captcha.model import u_net_denoise
from click_captcha.pre_process import add_contrast
from click_captcha.utils import pil_resize

image_shape = (32, 192, 1)
weights_path = "./models/e153-loss53.97-denoise.h5"
project_dir = os.path.dirname(os.path.abspath(__file__)) + "/../"
model = u_net_denoise(input_shape=image_shape, class_num=image_shape[2])
model.load_weights(weights_path)


def denoise(image_np):
    X = []
    # img = cv2.imread(image_path)
    img = pil_resize(image_np, image_shape[0], image_shape[1])
    # cv2.imshow("img", img)

    # img = add_contrast(img)
    # cv2.imshow("contrast", img)
    # cv2.waitKey(0)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = np.expand_dims(img, axis=-1)
    img = img / 255.
    X.append(img)

    X = np.array(X)
    pred = model.predict(X)
    # print(pred.shape)

    pred = np.uint8(pred[0]*255.)

    # pred = cv2.adaptiveThreshold(pred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,
    #                              11, 20)
    #
    # cv2.imshow("pred", pred)
    # cv2.waitKey(0)

    return pred


if __name__ == "__main__":
    # _path = "../data/test/char_9.jpg"
    _paths = glob("../data/equation/*")
    # _paths = glob("../data/test/FileInfo1021/*")
    for _path in _paths:
        denoise(cv2.imread(_path))