app.py

ffrom transformers import AutoModel, AutoTokenizer
import streamlit as st
from PIL import Image
import re
import os
import uuid

# Load the model and tokenizer only once
if "model" not in st.session_state or "tokenizer" not in st.session_state:
    @st.cache_resource
    def load_model(model_name):
        if model_name == "OCR for English or Hindi (CPU)":
            tokenizer = AutoTokenizer.from_pretrained('srimanth-d/GOT_CPU', trust_remote_code=True)
            model = AutoModel.from_pretrained('srimanth-d/GOT_CPU', trust_remote_code=True, use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
            model = model.eval()
        elif model_name == "OCR for English (GPU)":
            tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
            model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
            model = model.eval().to('cuda')
        return model, tokenizer

    # Load and store in session state
    model_option = "OCR for English or Hindi (CPU)"  # Default value for loading purposes
    model, tokenizer = load_model(model_option)
    st.session_state["model"] = model
    st.session_state["tokenizer"] = tokenizer
else:
    model = st.session_state["model"]
    tokenizer = st.session_state["tokenizer"]

# Function to run the GOT model for multilingual OCR
def run_ocr(image, model, tokenizer):
    unique_id = str(uuid.uuid4())
    image_path = f"{unique_id}.png"
    
    # Save image to disk
    image.save(image_path)

    try:
        # Use the model to extract text from the image
        res = model.chat(tokenizer, image_path, ocr_type='ocr')
        if isinstance(res, str):
            return res
        else:
            return str(res)
    except Exception as e:
        return f"Error: {str(e)}"
    finally:
        # Clean up the saved image
        if os.path.exists(image_path):
            os.remove(image_path)

# Function to highlight keyword in text
def highlight_text(text, search_term):
    if not search_term:
        return text
    pattern = re.compile(re.escape(search_term), re.IGNORECASE)
    return pattern.sub(lambda m: f'<span style="background-color: yellow;">{m.group()}</span>', text)

# Streamlit App
st.title("GOT-OCR Multilingual Demo")
st.write("Upload an image for OCR")

# Create two columns
col1, col2 = st.columns(2)

# Left column - Display the uploaded image
with col1:
    uploaded_image = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg"])
    if uploaded_image:
        image = Image.open(uploaded_image)
        st.image(image, caption='Uploaded Image', use_column_width=True)

# Right column - Model selection, options, and displaying extracted text
with col2:
    model_option = st.selectbox("Select Model", ["OCR for English or Hindi (CPU)", "OCR for English (GPU)"])
    
    if st.button("Run OCR"):
        with st.spinner("Processing..."):
            # Run OCR and store the result in session state
            if uploaded_image:
                result_text = run_ocr(image, model, tokenizer)
                if "Error" not in result_text:
                    st.session_state["extracted_text"] = result_text  # Store the result in session state
                else:
                    st.error(result_text)
            else:
                st.error("Please upload an image before running OCR.")

    # Display the extracted text if it exists in session state
    if "extracted_text" in st.session_state:
        extracted_text = st.session_state["extracted_text"]

        # Keyword input for search
        search_term = st.text_input("Enter a word or phrase to highlight:")
        
        # Highlight keyword in the extracted text
        highlighted_text = highlight_text(extracted_text, search_term)
        
        # Display the highlighted text using markdown
        st.subheader("Extracted Text:")
        st.markdown(f'<div style="white-space: pre-wrap;">{highlighted_text}</div>', unsafe_allow_html=True)