app.py

import os
import cv2
import numpy as np
from PIL import Image
from path import Path
import streamlit as st
from typing import Tuple
import easyocr  # Import EasyOCR

from pathlib import Path
import sys

# Add the 'app' directory to the sys.path
# Assuming 'app' is in the current working directory
sys.path.append(str(Path(__file__).parent / 'app'))
from app.dataloader_iam import Batch
from app.model import Model, DecoderType
from app.preprocessor import Preprocessor
from streamlit_drawable_canvas import st_canvas

# Set page config at the very beginning (only executed once)
st.set_page_config(
    page_title="HTR App",
    page_icon=":pencil:",
    layout="centered",
    initial_sidebar_state="auto",
)

ms = st.session_state
if "themes" not in ms: 
  ms.themes = {"current_theme": "light",
                    "refreshed": True,
                    
                    "light": {"theme.base": "dark",
                              "theme.backgroundColor": "black",
                              "theme.primaryColor": "#c98bdb",
                              "theme.secondaryBackgroundColor": "#5591f5",
                              "theme.textColor": "white",
                              "theme.textColor": "white",
                              "button_face": "🌜"},

                    "dark":  {"theme.base": "light",
                              "theme.backgroundColor": "white",
                              "theme.primaryColor": "#5591f5",
                              "theme.secondaryBackgroundColor": "#82E1D7",
                              "theme.textColor": "#0a1464",
                              "button_face": "🌞"},
                    }
  

def ChangeTheme():
  previous_theme = ms.themes["current_theme"]
  tdict = ms.themes["light"] if ms.themes["current_theme"] == "light" else ms.themes["dark"]
  for vkey, vval in tdict.items(): 
    if vkey.startswith("theme"): st._config.set_option(vkey, vval)

  ms.themes["refreshed"] = False
  if previous_theme == "dark": ms.themes["current_theme"] = "light"
  elif previous_theme == "light": ms.themes["current_theme"] = "dark"


btn_face = ms.themes["light"]["button_face"] if ms.themes["current_theme"] == "light" else ms.themes["dark"]["button_face"]
st.button(btn_face, on_click=ChangeTheme)

if ms.themes["refreshed"] == False:
  ms.themes["refreshed"] = True
  st.rerun()


def get_img_size(line_mode: bool = False) -> Tuple[int, int]:
    """
    Auxiliary method that sets the height and width
    Height is fixed while width is set according to the Model used.
    """
    if line_mode:
        return 256, get_img_height()
    return 128, get_img_height()

def get_img_height() -> int:
    """
    Auxiliary method that sets the height, which is fixed for the Neural Network.
    """
    return 32

def infer(line_mode: bool, model: Model, fn_img: Path) -> None:
    """
    Auxiliary method that does inference using the pretrained models:
    Recognizes text in an image given its path.
    """
    img = cv2.imread(fn_img, cv2.IMREAD_GRAYSCALE)
    assert img is not None

    preprocessor = Preprocessor(get_img_size(line_mode), dynamic_width=True, padding=16)
    img = preprocessor.process_img(img)

    batch = Batch([img], None, 1)
    recognized, probability = model.infer_batch(batch, True)
    return [recognized, probability]

def infer_super_model(image_path) -> None:
    reader = easyocr.Reader(['en'])  # Initialize EasyOCR reader
    result = reader.readtext(image_path)
    recognized_texts = [text[1] for text in result]  # Extract recognized texts
    probabilities = [text[2] for text in result]  # Extract probabilities
    return recognized_texts, probabilities



def main():

    st.title('Extract text from Image Demo')
    
    st.markdown("""
    Streamlit Web Interface for Handwritten Text Recognition (HTR), Optical Character Recognition (OCR) 
                implemented with TensorFlow and trained on the IAM off-line HTR dataset. 
                The model takes images of single words or text lines (multiple words) as input and outputs the recognized text. 
    """, unsafe_allow_html=True)
    
    st.markdown("""
    Predictions can be made using one of two models:
    - Single_Model (Trained on Single Word Images) 
    - Line_Model (Trained on Text Line Images)    
    - Super_Model ( Most Robust Option for English )
    - Burmese (Link)
    """, unsafe_allow_html=True)

    st.subheader('Select a Model, Choose the Arguments and Draw in the box below or Upload an Image to obtain a prediction.')

    #Selectors for the model and decoder
    modelSelect = st.selectbox("Select a Model", ['Single_Model', 'Line_Model', 'Super_Model'])
    

    if modelSelect != 'Super_Model':
        decoderSelect = st.selectbox("Select a Decoder", ['Bestpath', 'Beamsearch', 'Wordbeamsearch'])


    #Mappings (dictionaries) for the model and decoder. Asigns the directory or the DecoderType of the selected option.
    modelMapping = {
        "Single_Model": '../model/word-model',
        "Line_Model": '../model/line-model'
    }

    decoderMapping = {
        'Bestpath': DecoderType.BestPath,
        'Beamsearch': DecoderType.BeamSearch,
        'Wordbeamsearch': DecoderType.WordBeamSearch
    }

    #Slider for pencil width
    strokeWidth = st.slider("Stroke Width: ", 1, 25, 6)

    #Canvas/Text Box for user input. BackGround Color must be white (#FFFFFF) or else text will not be properly recognised.
    inputDrawn = st_canvas(
        fill_color="rgba(255, 165, 0, 0.3)", 
        stroke_width=strokeWidth,
        update_streamlit=True,
        background_image=None,
        height = 200,
        width = 400,
        drawing_mode='freedraw',
        key="canvas",
        background_color = '#FFFFFF'
    )

    #Buffer for user input (images uploaded from the user's device)
    inputBuffer = st.file_uploader("Upload an Image", type=["png"])

    #Inference Button
    inferBool = st.button("Recognize Text")

    # After clicking the "Recognize Text" button, check if the model selected is Super_Model
    if inferBool:
        if modelSelect == 'Super_Model':
            inputArray = None  # Initialize inputArray to None

            # Handling uploaded file
            if inputBuffer is not None:
                with Image.open(inputBuffer).convert('RGB') as img:
                    inputArray = np.array(img)

            # Handling canvas data
            elif inputDrawn.image_data is not None:
                # Convert RGBA to RGB
                inputArray = cv2.cvtColor(np.array(inputDrawn.image_data, dtype=np.uint8), cv2.COLOR_RGBA2RGB)

            # Now check if inputArray has been set
            if inputArray is not None:
                # Initialize EasyOCR Reader
                reader = easyocr.Reader(['en'])  # Assuming English language; adjust as necessary
                # Perform OCR
                results = reader.readtext(inputArray)

                # Display results
                all_text = ''
                for (bbox, text, prob) in results:
                    all_text += f'{text} (confidence: {prob:.2f})\n'

                st.write("**Recognized Texts and their Confidence Scores:**")
                st.text(all_text)
            else:
                st.write("No image data found. Please upload an image or draw on the canvas.")


        else:
            # Handle other model selections as before
            if ((inputDrawn.image_data is not None or inputBuffer is not None) and inferBool == True):
                #We turn the input into a numpy array
                if inputDrawn.image_data is not None:
                    inputArray = np.array(inputDrawn.image_data)
                
                if inputBuffer is not None:
                    inputBufferImage = Image.open(inputBuffer)
                    inputArray = np.array(inputBufferImage)

                #We turn this array into a .png format and save it. 
                inputImage = Image.fromarray(inputArray.astype('uint8'), 'RGBA')
                inputImage.save('userInput.png')
                #We obtain the model directory and the decoder type from their mapping
                modelDir = modelMapping[modelSelect]
                decoderType = decoderMapping[decoderSelect]

                #Finally, we call the model with this image as attribute and display the Best Candidate and its probability on the Interface
                model = Model(list(open(modelDir + "/charList.txt").read()), modelDir, decoderType, must_restore=True)
                inferedText = infer(modelDir == '../model/line-model', model, 'userInput.png')

                st.write("**Best Candidate: **", inferedText[0][0])
                st.write("**Probability: **", str(inferedText[1][0]*100) + "%")

if __name__ == "__main__":
    main()