| 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 |
|
|
| |
| 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, y_true, y_pred, input_length, label_length): |
| |
| |
| loss = self.loss_fn(y_true, y_pred, input_length, label_length) |
| self.add_loss(loss) |
|
|
| |
| return loss |
|
|
| |
| @st.cache_resource |
| def load_model(): |
| model_path = "model_ocr.h5" |
| model = tf.keras.models.load_model(model_path, custom_objects={"CTCLayer": CTCLayer}) |
| return model |
|
|
| model = load_model() |
|
|
|
|
| |
| img_width, img_height = 200, 50 |
|
|
| |
| max_length = 50 |
|
|
| |
| def prepare_image(img): |
| |
| img = img.resize((img_width, img_height)) |
| |
| |
| img_array = img_to_array(img) |
| |
| |
| img_array = np.expand_dims(img_array, axis=0) |
| img_array = np.transpose(img_array, (0, 2, 1, 3)) |
| |
| |
| input_length = np.ones((img_array.shape[0], 1)) * (img_width // 4) |
| label_length = np.ones((img_array.shape[0], 1)) * max_length |
|
|
| |
| dummy_label = np.zeros((img_array.shape[0], max_length)) |
| |
| |
| preds = model.predict([img_array, input_length, label_length, dummy_label]) |
| pred_texts = decode_batch_predictions(preds) |
|
|
| return pred_texts, preds |
|
|
| def decode_batch_predictions(pred): |
| |
| characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" |
| |
| pred_texts = [] |
|
|
| for i in range(len(pred)): |
| |
| pred_single = pred[i] |
|
|
| |
| pred_indices = np.argmax(pred_single, axis=-1) |
|
|
| |
| pred_indices = pred_indices[:5] |
|
|
| |
| pred_text = ''.join([characters[int(c)] for c in pred_indices if c != -1]) |
|
|
| |
| while len(pred_text) < 5: |
| pred_text += " " |
|
|
| |
| pred_texts.append(pred_text) |
|
|
| return pred_texts |
| |
| def run(): |
| st.title("OCR Model Deployment") |
| |
| |
| img_file = st.file_uploader("Choose an Image", type=["jpg", "png"]) |
|
|
| if img_file is not None: |
| img = Image.open(img_file).convert('L') |
| st.image(img, use_column_width=True) |
|
|
| |
| 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()) |
|
|
| |
| pred_texts = prepare_image(img) |
| |
| |
| st.success(f"**Predicted Text: {pred_texts[0]}**") |
|
|
| if __name__ == "__main__": |
| run() |
|
|