Update app.py

app.py

@@ -1,40 +1,48 @@
 import streamlit as st
 from PIL import Image
-from tensorflow.keras.models import load_model
 import tensorflow as tf
 import numpy as np
-from keras.preprocessing.image import img_to_array
 import os
 
+# Caching the model loading function to optimize performance
 @st.cache_resource
-def load_model():
+def load_captcha_model():
     model_path = "captcha_ocr_model.h5"  # Update with the actual CAPTCHA model path
     return tf.keras.models.load_model(model_path)
 
-model = load_model()
+# Load the model
+model = load_captcha_model()
 
+# Function to prepare the image for model prediction
 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)
+    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)
+        # 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])
+        # Decode predictions assuming the model outputs probabilities
+        # Modify this part based on your specific model's output
+        decoded_captcha = ''.join([chr(np.argmax(pred) + ord('A')) for pred in predictions])
 
-    return decoded_captcha, 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"])
+    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, use_column_width=False)
+        st.image(img, caption="Uploaded CAPTCHA", use_column_width=True)
 
         # Create the directory if it doesn't exist
         upload_dir = './upload_images/'
@@ -47,6 +55,9 @@ def run():
 
         # Predict the CAPTCHA
         predicted_captcha, score = prepare_captcha_image(img)
-        st.success(f"**Predicted CAPTCHA: {predicted_captcha}**")
+        if predicted_captcha:
+            st.success(f"**Predicted CAPTCHA: {predicted_captcha}**")
+        else:
+            st.error("Failed to predict CAPTCHA.")
 
 run()