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']


# Preprocessing gambar untuk memastikan bahwa gambar sesuai dengan input yang diinginkan
def prepare_image(img):
    # Resize gambar ke ukuran yang diinginkan
    img = img.resize((img_width, img_height))
    # Konversi ke array dan normalisasi gambar
    img_array = img_to_array(img) / 255.0  # Normalisasi
    # Tambahkan dimensi batch dan sesuaikan dengan dimensi yang diinginkan model
    img_array = np.expand_dims(img_array, axis=0)  # Batch size 1
    img_array = np.transpose(img_array, (0, 2, 1, 3))  # Untuk model dengan dimensi (batch, width, height, channels)
    return img_array


def decode_batch_predictions(pred):
    pred_texts = []
    
    # Loop untuk setiap prediksi dalam batch
    for i in range(pred.shape[0]):
        # Mengambil argmax untuk mendapatkan indeks dengan probabilitas tertinggi
        pred_indices = np.argmax(pred[i], axis=-1)  # Ambil argmax untuk setiap karakter
        
        # Memetakan indeks ke karakter (mengecualikan padding dan placeholder)
        pred_text = ''.join([characters[int(c)] for c in pred_indices if c not in [-1, 0]])
        
        # Menambahkan hasil teks untuk batch ke pred_texts
        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()