import pandas as pd
import numpy as np
import streamlit as st
import easyocr
from PIL import Image, ImageDraw
from streamlit_drawable_canvas import st_canvas

def rectangle(image, result):
    """Draw rectangles on the image based on predicted coordinates."""
    draw = ImageDraw.Draw(image)
    for res in result:
        top_left = tuple(res[0][0])  # top left coordinates as a tuple
        bottom_right = tuple(res[0][2])  # bottom right coordinates as a tuple
        draw.rectangle((top_left, bottom_right), outline="blue", width=2)
    # Display the image on Streamlit
    st.image(image)

# Main title
st.title("Get text from an image with EasyOCR")

# Subtitle
st.markdown("## EasyOCR with Streamlit")

# Upload image file or draw
st.markdown("## Upload an Image or Draw")
col1, col2 = st.columns(2)

with col1:
    file = st.file_uploader("Upload Here", type=['png', 'jpg', 'jpeg'])

with col2:
    # Drawable canvas
    canvas_result = st_canvas(
        fill_color="rgba(255, 165, 0, 0.3)",
        stroke_width=3,
        stroke_color="#ffffff",
        background_color="#000000",
        background_image=None if file else st.session_state.get("background", None),
        update_streamlit=True,
        width=400,
        height=400,
        drawing_mode="freedraw",
        key="canvas",
    )

image = None  # Initialize image variable

# Process uploaded image or drawing
if file is not None:
    image = Image.open(file)  # Read image with PIL library
elif canvas_result.image_data is not None:
    image = Image.fromarray(np.array(canvas_result.image_data, dtype=np.uint8)).convert('RGB')

if image is not None:
    st.image(image)  # Display the image

    # Initialize EasyOCR reader; you can add or remove languages based on your preference
    reader = easyocr.Reader(['en'], gpu=False)
    result = reader.readtext(np.array(image))  # Convert image to numpy array and process with EasyOCR

    # Print all predicted text
    for idx, res in enumerate(result):
        pred_text = res[1]
        st.write(pred_text)
    
    # Collect the results in a dictionary
    textdic_easyocr = {res[1]: {'pred_confidence': res[2]} for res in result}

    # Create a DataFrame to show the predicted text and prediction confidence
    df = pd.DataFrame.from_dict(textdic_easyocr).T
    st.table(df)

    # Draw rectangles around the detected text in the image
    rectangle(image, result)
else:
    st.write("Please upload an image or use the canvas to draw.")