Captcha / app.py
Reaumur's picture
Update app.py
57fe102 verified
raw
history blame
2.43 kB
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()