app.py

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

# 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, 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)

        # On test time, just return the computed loss
        return loss

# Load the trained model with a custom CTC layer if needed
@st.cache_resource
def load_model():
    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

model = load_model()


# Menambahkan definisi img_width dan img_height
img_width, img_height = 200, 50  # Ganti sesuai dimensi input gambar yang digunakan oleh model Anda

# Definisikan max_length (misalnya panjang label maksimal)
max_length = 50  # Ganti sesuai dengan panjang label teks maksimal yang diinginkan

# Pemetaan karakter yang mencakup huruf (kapital dan kecil) serta angka
characters = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 
              'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 
              'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 
              'w', 'x', 'y', 'z', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 
              'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 
              'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z']


# Function to preprocess the image
def prepare_image(img):
    # Resize gambar sesuai dengan ukuran yang diharapkan oleh model
    img = img.resize((img_width, img_height))  # Resize to (200, 50)
    
    # Konversi gambar ke array
    img_array = img_to_array(img)
    
    # Tambahkan dimensi untuk batch (menjadi 1, 50, 200) dan reshape ke bentuk (1, 50, 200, 1)
    img_array = np.expand_dims(img_array, axis=0)  # Tambahkan dimensi untuk batch
    img_array = np.transpose(img_array, (0, 2, 1, 3))  # Mengubah urutan dimensi menjadi (1, 200, 50, 1)
    
    # Menyusun input_length dan label_length untuk model OCR
    input_length = np.ones((img_array.shape[0], 1)) * (img_width // 4)  # Sesuaikan dengan input panjang
    label_length = np.ones((img_array.shape[0], 1)) * max_length  # Example label length

    # Menambahkan input dummy untuk label (jika perlu untuk prediksi)
    dummy_label = np.zeros((img_array.shape[0], max_length))  # Input dummy jika model mengharapkan label input
    
    # Melakukan prediksi
    preds = model.predict([img_array, input_length, label_length, dummy_label])  # Berikan 4 input
    pred_texts = decode_batch_predictions(preds)

    return pred_texts, preds

def decode_batch_predictions(pred):
    pred_texts = []
    
    # Loop untuk setiap batch (jika ada lebih dari satu batch)
    for i in range(pred.shape[0]):
        # Ambil argmax untuk mendapatkan indeks karakter yang diprediksi (yaitu karakter yang paling mungkin)
        pred_indices = np.argmax(pred[i], axis=-1)  # Cari indeks dengan probabilitas tertinggi
        
        # Sekarang kita memetakan indeks ini ke karakter, mengecualikan nilai yang tidak valid seperti -1 atau 0 (padding)
        pred_text = ''.join([characters[int(c)] for c in pred_indices if c != -1 and c != 0])
        
        pred_texts.append(pred_text)  # Simpan teks hasil decoding untuk batch yang sedang diproses
    
    return pred_texts

    
def run():
    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).convert('L')  # Convert to grayscale if needed
        st.image(img, use_column_width=True)

        # 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]}**")

if __name__ == "__main__":
    run()