KCAPTCHA prediction

predict.py

import tensorflow as tf
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications import densenet
import numpy as np

def decode_prediction(predicted, char_set, length=2):
    l = len(char_set)
    chars = []
    for i in range(length):
        chars.append(char_set[np.argmax(predicted[i * l : (i + 1) * l])])

    return "".join(chars)

def main():
    image_height = 60
    image_width = 160
    img_path = "test_img.png"

    model_path = "l2.model.densenet121.h5"
    char_set = "0123456789"
    captcha_length = 2
    
    img = image.load_img(
        img_path,
        target_size=(image_height, image_width),
    )
    img = image.img_to_array(img)
    img = densenet.preprocess_input(img)

    net = tf.keras.models.load_model(model_path)
    predictions = net.predict(np.expand_dims(img, 0))

    prediction = predictions[0]
    prediction_decoded = decode_prediction(prediction, char_set, captcha_length)
    print(f"Predicted: {prediction_decoded}")


if __name__ == "__main__":
    main()