import easyocr as ocr
import streamlit as st
from PIL import Image
import numpy as np
import base64
from st_clickable_images import clickable_images

if "language" not in st.session_state:
    st.session_state.language=['ch_tra', 'en']

if "output_text" not in st.session_state:
    st.session_state.output_text=""

if "image" not in st.session_state:
    st.session_state.image=""

def extract_text():
    language=st.session_state.language
    st.session_state.image=st.session_state.uploaded_file
    if st.session_state.image:
        st.session_state.output_text=rjxai_extract(
            language,
            st.session_state.image
        )
    else:
        st.session_state.image=""
        st.session_state.output_text=""

def rjxai_extract(
    language,
    uploaded_file,
):  
    input_image=Image.open(uploaded_file)
    reader=ocr.Reader(language, model_storage_directory='.')
    result=reader.readtext(np.array(input_image))
    result_text = [] 
    for text in result:
        result_text.append(text[1])
    return ", ".join(result_text)

# side
with st.sidebar:
    st.header("示例")
    
    images = []
    with open("assets/example/1.jpg", "rb") as image:
        encoded=base64.b64encode(image.read()).decode()
        images.append(f"data:image/jpeg;base64,{encoded}")
    clicked=clickable_images(
        images,
        titles=[f"Image #{str(i)}" for i in range(len(images))],
        div_style={"display": "flex", "justify-content": "center", "flex-wrap": "wrap"},
        img_style={"margin": "5px", "height": "200px"},
        key="clicked"
    )
    if st.session_state.clicked==0:
        st.session_state.image="assets/example/1.jpg"
        st.session_state.output_text=rjxai_extract(
            st.session_state.language,
            "assets/example/1.jpg"
        )

# content
with st.container():
    st.header("Rjxai v4")
    st.subheader("Extract Text from Images")

    # 参考：https://github.com/JaidedAI/EasyOCR/issues/989
    # st.multiselect(
    #     'Select extract language',
    #     ['en', 'ch_sim', 'ch_tra', 'ja', 'es', 'fr', 'fa'],
    #     key="language"
    # )

    uploaded_file=st.file_uploader(
        label = "Upload your image here: png, jpg, jpeg", 
        type=['png','jpg','jpeg'],
        on_change=extract_text,
        key="uploaded_file",
    )

    if st.session_state.image:
        input_image=Image.open(st.session_state.image)
        st.image(input_image)
    
    if st.session_state.output_text:
        st.text_area(
            "extract text:",
            key="output_text",
            disabled=True
        )