Spaces:
Sleeping
Sleeping
File size: 3,958 Bytes
c8267f0 3cefb7d c8267f0 1739588 c8267f0 cadd139 c8267f0 1739588 1ec038e c8267f0 1739588 c8267f0 a6afa92 c8267f0 de484e9 57613ba 20edee5 57613ba c8267f0 a54ad4b c8267f0 1739588 c8267f0 1739588 c8267f0 1739588 |
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 |
import gradio as gr
import spaces
from transformers import AutoModel, AutoTokenizer
import os
import base64
import io
import uuid
import time
import shutil
from pathlib import Path
import re
import easyocr
tokenizer = AutoTokenizer.from_pretrained('RufusRubin777/GOT-OCR2_0_CPU', trust_remote_code=True, device_map='cpu')
model = AutoModel.from_pretrained('RufusRubin777/GOT-OCR2_0_CPU', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cpu', use_safetensors=True)
model = model.eval().cpu()
reader = easyocr.Reader(['hi'])
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()
def run_GOT(image, language):
unique_id = str(uuid.uuid4())
image_path = os.path.join(UPLOAD_FOLDER, f"{unique_id}.png")
shutil.copy(image, image_path)
try:
english_extraction = model.chat(tokenizer, image_path, ocr_type='ocr')
hindi_extraction = reader.readtext(image)
hindi_extract = ''
for x in hindi_extraction:
hindi_extract += x[1] + '\n'
return english_extraction + '\n' + hindi_extract
except Exception as e:
return f"Error: {str(e)}", None
finally:
if os.path.exists(image_path):
os.remove(image_path)
def search_keyword(text, keyword):
if not keyword:
return "<p>Please enter a keyword to search.</p>"
# Convert text and keyword to lowercase for case-insensitive search
text_lower = text.lower()
keyword_lower = keyword.lower()
# Find all occurrences of the keyword
matches = list(re.finditer(re.escape(keyword_lower), text_lower))
if not matches:
return f"<p>Keyword '{keyword}' not found in the text.</p>"
# Highlight all occurrences of the keyword
result = []
last_end = 0
for match in matches:
start, end = match.span()
result.append(text[last_end:start])
result.append(f'<mark>{text[start:end]}</mark>')
last_end = end
result.append(text[last_end:])
highlighted_text = ''.join(result)
# Wrap the result in a scrollable div with some basic styling
return f"""
<div style="max-height: 300px; overflow-y: auto; border: 1px solid #ccc; padding: 10px; background-color: #f9f9f9;">
<p>{highlighted_text}</p>
</div>
<p>Found {len(matches)} occurrence(s) of '{keyword}'</p>
"""
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()
title_html = """
<h1><span class="gradient-text" id="text">IIT Roorkee - OCR and Document Search Web Application Prototype (General OCR Theory (GOT), a 580M end-to-end OCR 2.0 model.)</span></h1>
"""
with gr.Blocks() as scan_master_web_app:
gr.HTML(title_html)
with gr.Row():
with gr.Column():
image_input = gr.Image(type="filepath", label="Upload your image")
submit_button = gr.Button("Extract Text")
ocr_result = gr.Textbox(label="Extracted Text")
with gr.Column():
keyword = gr.Textbox(label="Search a keyword in the extracted text")
search_button = gr.Button("Search Keyword")
search_result = gr.HTML(label="Search result")
submit_button.click(
run_GOT,
inputs=[image_input],
outputs=[ocr_result]
)
search_button.click(
search_keyword,
inputs=[ocr_result, keyword],
outputs=[search_result]
)
if __name__ == "__main__":
cleanup_old_files()
scan_master_web_app.launch() |