Kethan20's picture
Update app.py
48af6c5 verified
raw
history blame
1.37 kB
from transformers import AutoModel, AutoTokenizer
import streamlit as st
from PIL import Image
import tempfile
def perform_ocr(image):
tokenizer = AutoTokenizer.from_pretrained('srimanth-d/GOT_CPU', trust_remote_code=True)
model = AutoModel.from_pretrained('srimanth-d/GOT_CPU', trust_remote_code=True, low_cpu_mem_usage=True, use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
model = model.eval()
res = model.chat(tokenizer, image, ocr_type='ocr')
return res
# Title and instructions
st.title(' OCR and Document Search Web Application Prototype')
st.write('Upload an image and extract text in Hindi and English. You can also search for keywords within the extracted text.')
# Upload the image
uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "png", "jpeg"])
# If an image is uploaded
if uploaded_file is not None:
image = Image.open(uploaded_file)
st.image(image, caption='Uploaded Image.', use_column_width=True)
with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
temp_file.write(uploaded_file.getbuffer())
temp_file_path = temp_file.name
# Perform OCR on the uploaded image
st.write("Extracting text...")
extracted_text = perform_ocr(temp_file_path)
st.write("Extracted Text:")
st.text_area("OCR Output", extracted_text, height=200)