pranavdaware commited on
Commit
026ea7b
·
verified ·
1 Parent(s): 391041b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +74 -0
app.py ADDED
@@ -0,0 +1,74 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModel, AutoTokenizer
3
+ import os
4
+ import re
5
+
6
+ # Load the model and tokenizer from local path
7
+ # Assuming your model and tokenizer are stored in '/content/model' directory in Colab
8
+ model_path = 'pranavdaware/web_ocr'
9
+
10
+ # Load the model and tokenizer from the local directory
11
+ tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
12
+ model = AutoModel.from_pretrained(model_path, trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True)
13
+ model = model.eval().cuda()
14
+
15
+ # Function to extract text using OCR
16
+ def ocr_processing(image_file):
17
+ try:
18
+ # Perform OCR on the uploaded image
19
+ result = model.chat(tokenizer, image_file, ocr_type='ocr')
20
+ return result
21
+ except Exception as e:
22
+ return str(e)
23
+
24
+ # Function to search for keywords in extracted text
25
+ def search_keyword(ocr_text, keyword):
26
+ try:
27
+ # Use regex to search for the keyword and highlight matches
28
+ matches = re.findall(rf"({keyword})", ocr_text, re.IGNORECASE)
29
+ if matches:
30
+ highlighted_text = re.sub(rf"({keyword})", r'<mark>\1</mark>', ocr_text, flags=re.IGNORECASE)
31
+ return highlighted_text
32
+ else:
33
+ return f"No matches found for '{keyword}' in the extracted text."
34
+ except Exception as e:
35
+ return str(e)
36
+
37
+ # Gradio interface
38
+ def main():
39
+ # Gradio app layout
40
+ with gr.Blocks() as demo:
41
+ gr.Markdown("# OCR and Keyword Search Application")
42
+
43
+ with gr.Row():
44
+ with gr.Column():
45
+ image_input = gr.Image(type="filepath", label="Upload your image")
46
+ keyword_input = gr.Textbox(label="Enter keyword to search")
47
+ ocr_output = gr.Textbox(label="OCR Output")
48
+ search_output = gr.HTML(label="Search Results")
49
+
50
+ # Button for OCR processing
51
+ process_button = gr.Button("Process Image for OCR")
52
+
53
+ # Connect the OCR processing function to the button
54
+ process_button.click(
55
+ fn=ocr_processing,
56
+ inputs=image_input,
57
+ outputs=ocr_output
58
+ )
59
+
60
+ # Button for keyword search
61
+ search_button = gr.Button("Search Keyword in OCR Text")
62
+
63
+ # Connect the search function to the button
64
+ search_button.click(
65
+ fn=search_keyword,
66
+ inputs=[ocr_output, keyword_input],
67
+ outputs=search_output
68
+ )
69
+
70
+ # Launch the Gradio demo
71
+ demo.launch()
72
+
73
+ if __name__ == "__main__":
74
+ main()