import streamlit as st
from PIL import Image
import tensorflow as tf
import numpy as np
import os
from tensorflow.keras.layers import LSTM
from keras.saving import register_keras_serializable

# Kelas LSTM Kustom
@register_keras_serializable()
class CustomLSTM(LSTM):
    pass

# Caching the model loading function to optimize performance
@st.cache_resource
def load_captcha_model():
    model_path = "model1.keras"  # Update with the actual CAPTCHA model path
    return tf.keras.models.load_model(model_path, custom_objects={'CustomLSTM': CustomLSTM})

# Load the model
model = load_captcha_model()

# Function to prepare the image for model prediction
def prepare_captcha_image(img):
    try:
        # 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.convert('L'))  # Convert to grayscale if necessary
        img_array = img_array / 255.0  # Normalize image
        img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension

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

        # Decode predictions assuming the model outputs probabilities
        decoded_captcha = ''.join([chr(np.argmax(pred) + ord('A')) for pred in predictions])

        return decoded_captcha, predictions

    except Exception as e:
        st.error(f"Error preparing image: {e}")
        return None, None

# Main function to run the Streamlit app
def run():
    st.title("CAPTCHA Prediction")
    img_file = st.file_uploader("Upload a CAPTCHA Image", type=["jpg", "png", "jpeg"])

    if img_file is not None:
        img = Image.open(img_file)
        st.image(img, caption="Uploaded CAPTCHA", use_column_width=True)

        # 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)
        if predicted_captcha:
            st.success(f"**Predicted CAPTCHA: {predicted_captcha}**")
        else:
            st.error("Failed to predict CAPTCHA.")

# Run the Streamlit app
if __name__ == "__main__":
    run()