import os
import re
import sys
from glob import glob
import cv2
import numpy as np
import tensorflow as tf
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../")
from model import u_net_denoise
from utils import pil_resize

package_dir = os.path.abspath(os.path.dirname(__file__))
image_shape = (32, 192, 1)
model_path = package_dir + "/models/denoise_loss_53.97.h5"


def denoise(image_np, model=None, sess=None):
    if sess is None:
        sess = tf.compat.v1.Session(graph=tf.Graph())

    if model is None:
        with sess.as_default():
            with sess.graph.as_default():
                model = u_net_denoise(input_shape=image_shape, class_num=image_shape[2])
                model.load_weights(model_path)

    X = []
    img = pil_resize(image_np, image_shape[0], image_shape[1])
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = np.expand_dims(img, axis=-1)
    img = img / 255.

    X.append(img)
    X = np.array(X)
    with sess.as_default():
        with sess.graph.as_default():
            pred = model.predict(X)
    pred = np.uint8(pred[0]*255.)

    # cv2.imshow("origin", image_np)
    # cv2.imshow("pred", pred)
    # cv2.waitKey(0)
    return pred


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