import streamlit as st
from PIL import Image
import tensorflow as tf
import numpy as np
import os

@st.cache_resource
def load_model():
    model_path = "model.h5"  # Update with the actual CAPTCHA model path
    return tf.keras.models.load_model(model_path)

model = load_model()

def prepare_captcha_image(img):
    # Resize image to the input shape required by the CAPTCHA model
    img = img.resize((200, 50))  # Adjust size according to the trained model
    img_array = np.array(img)
    img_array = img_array / 255.0  # Normalize image
    img_array = np.expand_dims(img_array, axis=0)

    # Predict the CAPTCHA characters
    predictions = model.predict(img_array)

    # Assuming the model outputs one-hot encoded characters, decode the predictions
    decoded_captcha = ''.join([chr(np.argmax(pred) + ord('A')) for pred in predictions])

    return decoded_captcha, predictions

def run():
    st.title("CAPTCHA Prediction")
    img_file = st.file_uploader("Upload a CAPTCHA Image", type=["jpg", "png"])

    if img_file is not None:
        img = Image.open(img_file)
        st.image(img, use_column_width=False)

        # Create the directory if it doesn't exist
        upload_dir = './upload_images/'
        os.makedirs(upload_dir, exist_ok=True)

        # Save the uploaded image
        save_image_path = os.path.join(upload_dir, img_file.name)
        with open(save_image_path, "wb") as f:
            f.write(img_file.getbuffer())

        # Predict the CAPTCHA
        predicted_captcha, score = prepare_captcha_image(img)
        st.success(f"**Predicted CAPTCHA: {predicted_captcha}**")

run()