File size: 1,930 Bytes
e535f9e |
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 |
import numpy as np
import easyocr
from transformers import AutoModel, AutoTokenizer
from PIL import Image
import warnings
from transformers import logging
import re
#To Surpaas warnings
warnings.filterwarnings("ignore", message="The attention mask and the pad token id were not set.")
warnings.filterwarnings("ignore", message="Setting `pad_token_id` to `eos_token_id`")
warnings.filterwarnings("ignore", message="The `seen_tokens` attribute is deprecated")
logging.set_verbosity_error()
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()
easyocr_reader = easyocr.Reader(['hi'], gpu=False)
# Function to perform OCR based on selected language
def perform_ocr(image, language):
if language == "Hindi":
image_np = np.array(image)
result = easyocr_reader.readtext(image_np, detail=0)
return ' '.join(result)
elif language == "English":
image_path = 'temp_image.png'
image.save(image_path)
result = model.chat(tokenizer, image_path, ocr_type='ocr')
return result
else:
return "Invalid language selection. Please choose Hindi or English."
def process_keyword(image, language, keyword):
extracted_text = perform_ocr(image, language)
if keyword:
keyword_regex = re.escape(keyword)
highlighted_text = re.sub(
f'({keyword_regex})', r'<mark style="background-color: yellow">\1</mark>', extracted_text, flags=re.IGNORECASE
)
if highlighted_text != extracted_text:
return highlighted_text
else:
return f"No keyword '{keyword}' found in the text."
else:
return extracted_text
|