import streamlit as st
from streamlit_drawable_canvas import st_canvas
import cv2
import numpy as np
from tensorflow.keras.models import load_model
from PIL import Image
import easyocr
import pandas as pd

# Load the model for Myanmar character recognition
model = load_model('mm.h5')

# Initialize EasyOCR reader for English
reader = easyocr.Reader(['en'], gpu=False)

class_lists = [
   "0",
   "1",
   "2",
   "3",
   "4",
   "5",
   "6",
   "7",
   "8",
   "9",
    "Ah",
    "Aha",
    "au2",
    "au3",
    "ay2",
    "ba_htoat_chite",
    "ba_kone",
    "da_htway",
    "da_out_chite",
    "da_yay_hmote",
    "da_yin_kout",
    "e1",
    "e2",
    "eeare",
    "ga_khi",
    "ga_nge",
    "ha",
    "hsa_lain",
    "hta_hsin_htu",
    "hta_wun_beare",
    "ka_kji",
    "kha_khway",
    "la",
    "la_kji",
    "ma",
    "na_kji",
    "na_ngear",
    "nga",
    "nga_kyi",
    "O",
    "pa_sout",
    "pfa_u_htoat",
    "sah_lone",
    "ta_thun_lyin_chate",
    "ta_wun_pu",
    "tha",
    "u1",
    "u2",
    "un",
    "wa",
    "yah_kout",
    "yah_pet_let",
    "za_kwear",
    "za_myin_hsware"
]

# Streamlit UI
st.title('Text and Character Recognizer')
st.markdown('''
Select the mode for recognition:
''')

# Choose mode
mode = st.radio("Mode", ('English Text Recognition', 'Myanmar Character Recognition'))

if mode == 'English Text Recognition':
    uploaded_file = st.file_uploader("Upload your file here...", key="uploader_english")
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image', use_column_width=True)

        # EasyOCR to recognize text
        result = reader.readtext(np.array(image))
        for detection in result:
            st.write(f'Detected text: {detection[1]}, Confidence: {detection[2]}')

elif mode == 'Myanmar Character Recognition':
    col1, col2 = st.columns(2)

    with col1:
        uploaded_file = st.file_uploader("Upload your file here...", key="uploader_myanmar")

    with col2:
        # Initialize canvas
        canvas_result = st_canvas(
            fill_color="rgba(255, 165, 0, 0.3)",
            stroke_width=3,
            stroke_color="#ffffff",
            background_color="#000000",
            update_streamlit=True,
            width=200,
            height=200,
            drawing_mode="freedraw",
            key="canvas",
        )

    # Process the image for prediction
    image_data = None
    if uploaded_file is not None:
        image_data = Image.open(uploaded_file).convert('RGB')
    elif canvas_result.image_data is not None:
        image_data = Image.fromarray(np.uint8(canvas_result.image_data)).convert('RGB')

    if image_data is not None:
        # Convert PIL image to OpenCV format
        image_cv = np.array(image_data)
        image_cv = cv2.cvtColor(image_cv, cv2.COLOR_RGB2BGR)
        resized_image = cv2.resize(image_cv, (200, 200))
        # Prepare image for model input
        model_input = resized_image[np.newaxis, :, :, :3]

        st.write('Model Input')
        st.image(model_input, width=200)  # Display the input image to model

        if st.button('Predict Myanmar Character'):
            # Predict the class
            val = model.predict(model_input)
            predicted_class_index = np.argmax(val)
            mm_text = class_lists[predicted_class_index]
            st.write(f'Result: {mm_text}, Index: {predicted_class_index}')
            st.bar_chart(val[0])
    else:
        if mode == 'Myanmar Character Recognition':
            st.write("Please upload an image or draw in the canvas above.")