Update app.py
Browse files
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 |
-
#
|
8 |
-
class
|
9 |
-
def __init__(self,
|
10 |
-
super(
|
11 |
-
self.
|
12 |
-
self.activation = activation
|
13 |
-
self.dense = tf.keras.layers.Dense(units)
|
14 |
|
15 |
-
def call(self,
|
16 |
-
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
|
21 |
-
|
22 |
-
|
23 |
-
config.update({
|
24 |
-
'units': self.units,
|
25 |
-
'activation': self.activation,
|
26 |
-
})
|
27 |
-
return config
|
28 |
|
29 |
-
#
|
30 |
@st.cache_resource
|
31 |
def load_model():
|
32 |
-
model_path = "
|
33 |
-
|
|
|
34 |
|
35 |
-
# Load the model
|
36 |
model = load_model()
|
37 |
|
38 |
-
# Function to
|
39 |
def prepare_image(img):
|
40 |
-
|
41 |
-
|
42 |
-
|
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 |
-
|
49 |
-
|
|
|
50 |
|
51 |
-
|
52 |
-
|
|
|
|
|
|
|
53 |
|
54 |
-
|
55 |
|
56 |
-
|
57 |
-
|
58 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
59 |
|
60 |
-
# Main function to run the Streamlit app
|
61 |
def run():
|
62 |
-
st.title("
|
63 |
-
|
|
|
|
|
64 |
|
65 |
if img_file is not None:
|
66 |
-
img = Image.open(img_file)
|
67 |
-
st.image(img,
|
68 |
|
69 |
-
#
|
70 |
-
|
71 |
-
|
72 |
-
|
73 |
-
|
74 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
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()
|