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

# 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):
    # Daftar karakter yang digunakan dalam model OCR Anda
    characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"  # Sesuaikan dengan karakter model Anda
    
    pred_texts = []

    for i in range(len(pred)):
        # Ambil hasil prediksi untuk sampel i
        pred_single = pred[i]

        # Ambil argmax untuk setiap langkah (timesteps) untuk mendapatkan indeks karakter
        pred_indices = np.argmax(pred_single, axis=-1)

        # Ambil hasil prediksi dengan membatasi pada 5 karakter (untuk CAPTCHA dengan panjang 5 karakter)
        pred_indices = pred_indices[:5]  # Pastikan kita hanya mengambil 5 karakter pertama

        # Gabungkan prediksi menjadi string, menghindari padding (-1) dan pastikan panjangnya 5
        pred_text = ''.join([characters[int(c)] for c in pred_indices if c != -1])

        # Pastikan panjang teks hasil prediksi adalah 5 karakter
        while len(pred_text) < 5:
            pred_text += " "  # Bisa sesuaikan dengan apa yang ingin dilakukan jika panjangnya kurang dari 5

        # Append prediksi teks untuk batch
        pred_texts.append(pred_text)

    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()