Spaces:
Sleeping
Sleeping
| from transformers import AutoModel, AutoTokenizer | |
| import streamlit as st | |
| from PIL import Image | |
| import re | |
| import os | |
| import uuid | |
| # Set the page layout to wide | |
| st.set_page_config(layout="wide") | |
| # Load the model and tokenizer only once | |
| 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).eval() | |
| else: | |
| 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).eval().to('cuda') | |
| return model, tokenizer | |
| if "model" not in st.session_state or "tokenizer" not in st.session_state: | |
| model, tokenizer = load_model("OCR for English or Hindi (CPU)") | |
| st.session_state.update({"model": model, "tokenizer": tokenizer}) | |
| # Function to run the GOT model for multilingual OCR | |
| def run_ocr(image, model, tokenizer): | |
| image_path = f"{uuid.uuid4()}.png" | |
| image.save(image_path) | |
| try: | |
| res = model.chat(tokenizer, image_path, ocr_type='ocr') | |
| return res if isinstance(res, str) else str(res) | |
| except Exception as e: | |
| return f"Error: {str(e)}" | |
| finally: | |
| os.remove(image_path) | |
| # Function to highlight keyword in text | |
| def highlight_text(text, search_term): | |
| return re.sub(re.escape(search_term), lambda m: f'<span style="background-color: red;">{m.group()}</span>', text, flags=re.IGNORECASE) if search_term else text | |
| # Streamlit App | |
| st.title(":blue[Optical Character Recognition Application]") | |
| st.write("upload 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"): | |
| if uploaded_image: | |
| with st.spinner("Processing..."): | |
| model, tokenizer = load_model(model_option) | |
| result_text = run_ocr(image, model, tokenizer) | |
| if "Error" not in result_text: | |
| st.session_state["extracted_text"] = result_text | |
| 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: | |
| search_term = st.text_input("Enter a word or phrase to highlight:") | |
| st.subheader("Extracted Text:") | |
| st.markdown(f'<div style="white-space: pre-wrap;">{highlight_text(st.session_state["extracted_text"], search_term)}</div>', unsafe_allow_html=True) | |