Reaumur commited on
Commit
7496a0d
·
verified ·
1 Parent(s): 27c5f2d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +58 -51
app.py CHANGED
@@ -2,77 +2,84 @@ import streamlit as st
2
  from PIL import Image
3
  import tensorflow as tf
4
  import numpy as np
 
 
5
  import os
6
 
7
- # Definisi CustomLayer
8
- class CustomLayer(tf.keras.layers.Layer):
9
- def __init__(self, units=32, activation=None):
10
- super(CustomLayer, self).__init__()
11
- self.units = units
12
- self.activation = activation
13
- self.dense = tf.keras.layers.Dense(units)
14
 
15
- def call(self, inputs):
16
- x = self.dense(inputs)
17
- if self.activation:
18
- x = self.activation(x)
19
- return x
20
 
21
- def get_config(self):
22
- config = super(CustomLayer, self).get_config()
23
- config.update({
24
- 'units': self.units,
25
- 'activation': self.activation,
26
- })
27
- return config
28
 
29
- # Caching the model loading function to optimize performance
30
  @st.cache_resource
31
  def load_model():
32
- model_path = "captcha.keras" # Update with the actual model path
33
- return tf.keras.models.load_model(model_path, custom_objects={'CustomLayer': CustomLayer})
 
34
 
35
- # Load the model
36
  model = load_model()
37
 
38
- # Function to prepare the image for model prediction
39
  def prepare_image(img):
40
- try:
41
- # Resize image to the input shape required by the model
42
- img = img.resize((200, 50)) # Adjust size according to the trained model
43
- img_array = np.array(img.convert('L')) # Convert to grayscale if necessary
44
- img_array = img_array / 255.0 # Normalize image
45
- img_array = np.expand_dims(img_array, axis=-1) # Add channel dimension
46
- img_array = np.expand_dims(img_array, axis=0) # Add batch dimension
47
 
48
- # Predict the output using the loaded model
49
- predictions = model.predict(img_array)
 
50
 
51
- # Decode predictions assuming the model outputs probabilities
52
- decoded_captcha = ''.join([chr(np.argmax(pred) + ord('A')) for pred in predictions])
 
 
 
53
 
54
- return decoded_captcha, predictions
55
 
56
- except Exception as e:
57
- st.error(f"Error preparing image: {e}")
58
- return None, None
 
 
 
 
 
 
59
 
60
- # Main function to run the Streamlit app
61
  def run():
62
- st.title("CAPTCHA Prediction")
63
- img_file = st.file_uploader("Upload a CAPTCHA Image", type=["jpg", "png", "jpeg"])
 
 
64
 
65
  if img_file is not None:
66
- img = Image.open(img_file)
67
- st.image(img, caption="Uploaded CAPTCHA", use_column_width=True)
68
 
69
- # Predict the CAPTCHA
70
- predicted_captcha, score = prepare_image(img)
71
- if predicted_captcha:
72
- st.success(f"**Predicted CAPTCHA: {predicted_captcha}**")
73
- else:
74
- st.error("Failed to predict CAPTCHA.")
 
 
 
 
 
 
75
 
76
- # Run the Streamlit app
77
  if __name__ == "__main__":
78
  run()
 
2
  from PIL import Image
3
  import tensorflow as tf
4
  import numpy as np
5
+ from keras.preprocessing.image import img_to_array
6
+ from tensorflow.keras.models import load_model
7
  import os
8
 
9
+ # Load custom CTC Layer if necessary
10
+ class CTCLayer(tf.keras.layers.Layer):
11
+ def __init__(self, name=None):
12
+ super().__init__(name=name)
13
+ self.loss_fn = tf.keras.backend.ctc_batch_cost
 
 
14
 
15
+ def call(self, y_true, y_pred, input_length, label_length):
16
+ # Compute the training-time loss value and add it
17
+ # to the layer using `self.add_loss()`.
18
+ loss = self.loss_fn(y_true, y_pred, input_length, label_length)
19
+ self.add_loss(loss)
20
 
21
+ # On test time, just return the computed loss
22
+ return loss
 
 
 
 
 
23
 
24
+ # Load the trained model with a custom CTC layer if needed
25
  @st.cache_resource
26
  def load_model():
27
+ model_path = "model_ocr.h5" # Update with the correct model file path
28
+ model = tf.keras.models.load_model(model_path, custom_objects={"CTCLayer": CTCLayer})
29
+ return model
30
 
 
31
  model = load_model()
32
 
33
+ # Function to preprocess the image
34
  def prepare_image(img):
35
+ img = img.resize((img_width, img_height)) # Resize to the expected input size for the model
36
+ img_array = img_to_array(img)
37
+ img_array = np.expand_dims(img_array, axis=0) # Add batch dimension
 
 
 
 
38
 
39
+ # The input_length and label_length need to be set according to your data
40
+ input_length = np.ones((img_array.shape[0], 1)) * (img_width // 4) # Example input length
41
+ label_length = np.ones((img_array.shape[0], 1)) * max_length # Example label length
42
 
43
+ # Make prediction
44
+ preds = model.predict([img_array, input_length, label_length])
45
+
46
+ # Decode predictions (use your custom decoding function)
47
+ pred_texts = decode_batch_predictions(preds)
48
 
49
+ return pred_texts
50
 
51
+ # Define a simple batch decoder (adjust as needed)
52
+ def decode_batch_predictions(pred):
53
+ # This function should convert the predictions (logits) to text
54
+ # Modify this function based on your specific character map
55
+ pred_texts = []
56
+ for i in range(pred.shape[0]):
57
+ pred_text = ''.join([characters[int(c)] for c in pred[i] if c != -1]) # Map to characters
58
+ pred_texts.append(pred_text)
59
+ return pred_texts
60
 
 
61
  def run():
62
+ st.title("OCR Model Deployment")
63
+
64
+ # Upload image
65
+ img_file = st.file_uploader("Choose an Image", type=["jpg", "png"])
66
 
67
  if img_file is not None:
68
+ img = Image.open(img_file).convert('L') # Convert to grayscale if needed
69
+ st.image(img, use_column_width=True)
70
 
71
+ # Save the uploaded image
72
+ upload_dir = './upload_images/'
73
+ os.makedirs(upload_dir, exist_ok=True)
74
+ save_image_path = os.path.join(upload_dir, img_file.name)
75
+ with open(save_image_path, "wb") as f:
76
+ f.write(img_file.getbuffer())
77
+
78
+ # Process the image and make prediction
79
+ pred_texts = prepare_image(img)
80
+
81
+ # Show predicted text
82
+ st.success(f"**Predicted Text: {pred_texts[0]}**")
83
 
 
84
  if __name__ == "__main__":
85
  run()