Spaces:
Sleeping
Sleeping
File size: 6,257 Bytes
63e71c8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
import streamlit as st
from transformers import AutoModel, AutoTokenizer
import os
import base64
import io
import uuid
import shutil
from pathlib import Path
import time
import tempfile
model_name = "srimanth-d/GOT_CPU"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True, low_cpu_mem_usage=True, use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
model = model.eval()
UPLOAD_FOLDER = "./uploads"
RESULTS_FOLDER = "./results"
for folder in [UPLOAD_FOLDER, RESULTS_FOLDER]:
if not os.path.exists(folder):
os.makedirs(folder)
def image_to_base64(image):
buffered = io.BytesIO()
image.save(buffered, format="PNG")
return base64.b64encode(buffered.getvalue()).decode()
# Cleanup function for removing old files
def cleanup_old_files():
current_time = time.time()
for folder in [UPLOAD_FOLDER, RESULTS_FOLDER]:
for file_path in Path(folder).glob('*'):
if current_time - file_path.stat().st_mtime > 3600: # 1 hour
file_path.unlink()
# Function to search and highlight keywords in text
def search_in_text(text, keywords):
"""Searches for keywords within the text and highlights matches."""
if not keywords:
return text
highlighted_text = text
for keyword in keywords.split():
highlighted_text = highlighted_text.replace(keyword, f"<mark>{keyword}</mark>")
return highlighted_text
# OCR processing function
def run_GOT(image, got_mode, fine_grained_mode="", ocr_color="", ocr_box=""):
unique_id = str(uuid.uuid4())
image_path = os.path.join(UPLOAD_FOLDER, f"{unique_id}.png")
result_path = os.path.join(RESULTS_FOLDER, f"{unique_id}.html")
shutil.copy(image, image_path)
try:
if got_mode == "plain texts OCR":
res = model.chat(tokenizer, image_path, ocr_type='ocr')
return res, None
elif got_mode == "format texts OCR":
res = model.chat(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
elif got_mode == "plain multi-crop OCR":
res = model.chat_crop(tokenizer, image_path, ocr_type='ocr')
return res, None
elif got_mode == "format multi-crop OCR":
res = model.chat_crop(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
elif got_mode == "plain fine-grained OCR":
res = model.chat(tokenizer, image_path, ocr_type='ocr', ocr_box=ocr_box, ocr_color=ocr_color)
return res, None
elif got_mode == "format fine-grained OCR":
res = model.chat(tokenizer, image_path, ocr_type='format', ocr_box=ocr_box, ocr_color=ocr_color, render=True, save_render_file=result_path)
res_markdown = res
if "format" in got_mode and os.path.exists(result_path):
with open(result_path, 'r') as f:
html_content = f.read()
encoded_html = base64.b64encode(html_content.encode('utf-8')).decode('utf-8')
iframe_src = f"data:text/html;base64,{encoded_html}"
iframe = f'<iframe src="{iframe_src}" width="100%" height="600px"></iframe>'
download_link = f'<a href="data:text/html;base64,{encoded_html}" download="result_{unique_id}.html">Download Full Result</a>'
return res_markdown, f"{download_link}<br>{iframe}"
else:
return res_markdown, None
except Exception as e:
return f"Error: {str(e)}", None
finally:
if os.path.exists(image_path):
os.remove(image_path)
# Streamlit interface
st.title("GOT OCR 2.0 Model")
st.markdown("""
Upload your image below and select your preferred mode. Note that more characters may increase wait times.
- **Plain Texts OCR & Format Texts OCR:** Use these modes for basic image-level OCR. Format Text OCR is preferred for better results.
- **Plain Multi-Crop OCR & Format Multi-Crop OCR:** Ideal for images with complex content, offering higher-quality results.
- **Plain Fine-Grained OCR & Format Fine-Grained OCR:** These modes allow you to specify fine-grained regions on the image for more flexible OCR. Regions can be defined by coordinates or colors (red, blue, green, black or white).
""")
uploaded_image = st.file_uploader("Upload your image", type=["png", "jpg", "jpeg"])
got_mode = st.selectbox("Choose OCR mode", [
"plain texts OCR",
"format texts OCR",
"plain multi-crop OCR",
"format multi-crop OCR",
"plain fine-grained OCR",
"format fine-grained OCR"
])
if "fine-grained" in got_mode:
ocr_box = st.text_input("Input OCR box [x1,y1,x2,y2]")
ocr_color = st.selectbox("Choose OCR color", ["red", "green", "blue", "black", "white"])
else:
ocr_box = ""
ocr_color = ""
# Maintain state for OCR result
if 'ocr_result' not in st.session_state:
st.session_state.ocr_result = None
if 'html_result' not in st.session_state:
st.session_state.html_result = None
if st.button("Run OCR"):
if uploaded_image:
with tempfile.NamedTemporaryFile(delete=False) as temp:
temp.write(uploaded_image.read())
ocr_result, html_result = run_GOT(temp.name, got_mode, ocr_box=ocr_box, ocr_color=ocr_color)
st.session_state.ocr_result = ocr_result
st.session_state.html_result = html_result
st.text_area("OCR Result", ocr_result)
else:
st.warning("Please upload an image.")
# Display the OCR result if it has been set
if st.session_state.ocr_result:
st.text_area("OCR Result", st.session_state.ocr_result,key="display_area")
# Keyword search functionality
keywords = st.text_input("Enter keywords for highlighting",key="keyword_input")
if keywords:
highlighted_text = search_in_text(st.session_state.ocr_result, keywords)
st.markdown(highlighted_text, unsafe_allow_html=True)
if st.session_state.html_result:
st.markdown(st.session_state.html_result, unsafe_allow_html=True)
if __name__ == "__main__":
cleanup_old_files()
|