Update app.py

app.py

@@ -2,77 +2,84 @@ import streamlit as st
 from PIL import Image
 import tensorflow as tf
 import numpy as np
+from keras.preprocessing.image import img_to_array
+from tensorflow.keras.models import load_model
 import os
 
-# Definisi CustomLayer
-class CustomLayer(tf.keras.layers.Layer):
-    def __init__(self, units=32, activation=None):
-        super(CustomLayer, self).__init__()
-        self.units = units
-        self.activation = activation
-        self.dense = tf.keras.layers.Dense(units)
+# Load custom CTC Layer if necessary
+class CTCLayer(tf.keras.layers.Layer):
+    def __init__(self, name=None):
+        super().__init__(name=name)
+        self.loss_fn = tf.keras.backend.ctc_batch_cost
 
-    def call(self, inputs):
-        x = self.dense(inputs)
-        if self.activation:
-            x = self.activation(x)
-        return x
+    def call(self, y_true, y_pred, input_length, label_length):
+        # Compute the training-time loss value and add it
+        # to the layer using `self.add_loss()`.
+        loss = self.loss_fn(y_true, y_pred, input_length, label_length)
+        self.add_loss(loss)
 
-    def get_config(self):
-        config = super(CustomLayer, self).get_config()
-        config.update({
-            'units': self.units,
-            'activation': self.activation,
-        })
-        return config
+        # On test time, just return the computed loss
+        return loss
 
-# Caching the model loading function to optimize performance
+# Load the trained model with a custom CTC layer if needed
 @st.cache_resource
 def load_model():
-    model_path = "captcha.keras"  # Update with the actual model path
-    return tf.keras.models.load_model(model_path, custom_objects={'CustomLayer': CustomLayer})
+    model_path = "model_ocr.h5"  # Update with the correct model file path
+    model = tf.keras.models.load_model(model_path, custom_objects={"CTCLayer": CTCLayer})
+    return model
 
-# Load the model
 model = load_model()
 
-# Function to prepare the image for model prediction
+# Function to preprocess the image
 def prepare_image(img):
-    try:
-        # Resize image to the input shape required by the 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=-1)  # Add channel dimension
-        img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
+    img = img.resize((img_width, img_height))  # Resize to the expected input size for the model
+    img_array = img_to_array(img)
+    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
 
-        # Predict the output using the loaded model
-        predictions = model.predict(img_array)
+    # The input_length and label_length need to be set according to your data
+    input_length = np.ones((img_array.shape[0], 1)) * (img_width // 4)  # Example input length
+    label_length = np.ones((img_array.shape[0], 1)) * max_length  # Example label length
 
-        # Decode predictions assuming the model outputs probabilities
-        decoded_captcha = ''.join([chr(np.argmax(pred) + ord('A')) for pred in predictions])
+    # Make prediction
+    preds = model.predict([img_array, input_length, label_length])
+    
+    # Decode predictions (use your custom decoding function)
+    pred_texts = decode_batch_predictions(preds)
 
-        return decoded_captcha, predictions
+    return pred_texts
 
-    except Exception as e:
-        st.error(f"Error preparing image: {e}")
-        return None, None
+# Define a simple batch decoder (adjust as needed)
+def decode_batch_predictions(pred):
+    # This function should convert the predictions (logits) to text
+    # Modify this function based on your specific character map
+    pred_texts = []
+    for i in range(pred.shape[0]):
+        pred_text = ''.join([characters[int(c)] for c in pred[i] if c != -1])  # Map to characters
+        pred_texts.append(pred_text)
+    return pred_texts
 
-# 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"])
+    st.title("OCR Model Deployment")
+    
+    # Upload image
+    img_file = st.file_uploader("Choose an Image", type=["jpg", "png"])
 
     if img_file is not None:
-        img = Image.open(img_file)
-        st.image(img, caption="Uploaded CAPTCHA", use_column_width=True)
+        img = Image.open(img_file).convert('L')  # Convert to grayscale if needed
+        st.image(img, use_column_width=True)
 
-        # Predict the CAPTCHA
-        predicted_captcha, score = prepare_image(img)
-        if predicted_captcha:
-            st.success(f"**Predicted CAPTCHA: {predicted_captcha}**")
-        else:
-            st.error("Failed to predict CAPTCHA.")
+        # Save the uploaded image
+        upload_dir = './upload_images/'
+        os.makedirs(upload_dir, exist_ok=True)
+        save_image_path = os.path.join(upload_dir, img_file.name)
+        with open(save_image_path, "wb") as f:
+            f.write(img_file.getbuffer())
+
+        # Process the image and make prediction
+        pred_texts = prepare_image(img)
+        
+        # Show predicted text
+        st.success(f"**Predicted Text: {pred_texts[0]}**")
 
-# Run the Streamlit app
 if __name__ == "__main__":
     run()