Spaces:
Sleeping
Sleeping
| from 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: | |
| 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 = st.selectbox("Select Model", ["OCR for English or Hindi (CPU)", "OCR for English (GPU)"]) | |
| 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") | |
| # Upload image | |
| 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) | |
| if st.button("Run OCR"): | |
| with st.spinner("Processing..."): | |
| # Run OCR and store the result in session state | |
| 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) | |
| # Display the extracted text if it exists in session state | |
| if "extracted_text" in st.session_state: | |
| extracted_text = st.session_state["extracted_text"] | |
| st.subheader("Extracted Text:") | |
| st.text(extracted_text) # Display the raw extracted text | |
| # Keyword input for search | |
| search_term = st.text_input("Enter a word or phrase to highlight:") | |
| # Highlight keyword in the extracted text | |
| if search_term: | |
| highlighted_text = highlight_text(extracted_text, search_term) | |
| # Display the highlighted text using markdown | |
| st.markdown(highlighted_text, unsafe_allow_html=True) | |