import os

import cv2
import numpy as np
from click_captcha.model import mobile_net, cnn_net
from click_captcha.utils import pil_resize

image_shape = (40, 40, 3)
weights_path = "./models/char_f1_0.93.h5"
project_dir = os.path.dirname(os.path.abspath(__file__)) + "/../"


def recognize(image_path):
    model = cnn_net(input_shape=image_shape)
    model.load_weights(weights_path)

    img = cv2.imread(image_path)
    img = pil_resize(img, image_shape[0], image_shape[1])
    cv2.imshow("img", img)
    cv2.waitKey(0)
    img = img / 255.
    X = np.expand_dims(img, 0)

    pred = model.predict(X)
    index = int(np.argmax(pred))
    with open(project_dir + "data/chinese_5710.txt") as f:
        char_str = f.read()
    char = char_str[index]
    print("recognize chinese", char)
    return char


if __name__ == "__main__":
    _path = "../data/test/char_6.jpg"
    # _path = "../data/click/2019_73_1.jpg"
    recognize(_path)