Spaces:
Sleeping
Sleeping
| import os | |
| import spaces | |
| import nltk | |
| nltk.download('punkt',quiet=True) | |
| from doctr.io import DocumentFile | |
| from doctr.models import ocr_predictor | |
| import gradio as gr | |
| from PIL import Image | |
| import base64 | |
| from utils import HocrParser | |
| from happytransformer import HappyTextToText, TTSettings | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,logging | |
| from transformers.integrations import deepspeed | |
| import re | |
| from lang_list import ( | |
| LANGUAGE_NAME_TO_CODE, | |
| T2TT_TARGET_LANGUAGE_NAMES, | |
| TEXT_SOURCE_LANGUAGE_NAMES, | |
| ) | |
| logging.set_verbosity_error() | |
| DEFAULT_TARGET_LANGUAGE = "English" | |
| from transformers import SeamlessM4TForTextToText | |
| from transformers import AutoProcessor | |
| model = SeamlessM4TForTextToText.from_pretrained("facebook/hf-seamless-m4t-large") | |
| processor = AutoProcessor.from_pretrained("facebook/hf-seamless-m4t-large") | |
| import pytesseract as pt | |
| import cv2 | |
| # OCR Predictor initialization | |
| OCRpredictor = ocr_predictor(det_arch='db_mobilenet_v3_large', reco_arch='crnn_vgg16_bn', pretrained=True) | |
| # Grammar Correction Model initialization | |
| happy_tt = HappyTextToText("T5", "vennify/t5-base-grammar-correction") | |
| grammar_args = TTSettings(num_beams=5, min_length=1) | |
| # Spell Check Model initialization | |
| OCRtokenizer = AutoTokenizer.from_pretrained("Bhuvana/t5-base-spellchecker", use_fast=False) | |
| OCRmodel = AutoModelForSeq2SeqLM.from_pretrained("Bhuvana/t5-base-spellchecker") | |
| # zero = torch.Tensor([0]).cuda() | |
| # print(zero.device) | |
| def correct_spell(inputs): | |
| input_ids = OCRtokenizer.encode(inputs, return_tensors='pt') | |
| sample_output = OCRmodel.generate( | |
| input_ids, | |
| do_sample=True, | |
| max_length=512, | |
| top_p=0.99, | |
| num_return_sequences=1 | |
| ) | |
| res = OCRtokenizer.decode(sample_output[0], skip_special_tokens=True) | |
| return res | |
| def process_text_in_chunks(text, process_function, max_chunk_size=256): | |
| # Split text into sentences | |
| sentences = re.split(r'(?<=[.!?])\s+', text) | |
| processed_text = "" | |
| for sentence in sentences: | |
| # Further split long sentences into smaller chunks | |
| chunks = [sentence[i:i + max_chunk_size] for i in range(0, len(sentence), max_chunk_size)] | |
| for chunk in chunks: | |
| processed_text += process_function(chunk) | |
| processed_text += " " # Add space after each processed sentence | |
| return processed_text.strip() | |
| def greet(img, apply_grammar_correction, apply_spell_check,lang_of_input): | |
| if (lang_of_input=="Hindi"): | |
| res = pt.image_to_string(img,lang='hin') | |
| _output_name = "RESULT_OCR.txt" | |
| open(_output_name, 'w').write(res) | |
| return res, _output_name, None | |
| if (lang_of_input=="Punjabi"): | |
| res = pt.image_to_string(img,lang='pan') | |
| _output_name = "RESULT_OCR.txt" | |
| open(_output_name, 'w').write(res) | |
| return res, _output_name, None | |
| img.save("out.jpg") | |
| doc = DocumentFile.from_images("out.jpg") | |
| output = OCRpredictor(doc) | |
| res = "" | |
| for obj in output.pages: | |
| for obj1 in obj.blocks: | |
| for obj2 in obj1.lines: | |
| for obj3 in obj2.words: | |
| res += " " + obj3.value | |
| res += "\n" | |
| res += "\n" | |
| # Process in chunks for grammar correction | |
| if apply_grammar_correction: | |
| res = process_text_in_chunks(res, lambda x: happy_tt.generate_text("grammar: " + x, args=grammar_args).text) | |
| # Process in chunks for spell check | |
| if apply_spell_check: | |
| res = process_text_in_chunks(res, correct_spell) | |
| _output_name = "RESULT_OCR.txt" | |
| open(_output_name, 'w').write(res) | |
| # Convert OCR output to searchable PDF | |
| _output_name_pdf="RESULT_OCR.pdf" | |
| xml_outputs = output.export_as_xml() | |
| parser = HocrParser() | |
| base64_encoded_pdfs = list() | |
| for i, (xml, img) in enumerate(zip(xml_outputs, doc)): | |
| xml_element_tree = xml[1] | |
| parser.export_pdfa(_output_name_pdf, | |
| hocr=xml_element_tree, image=img) | |
| with open(_output_name_pdf, 'rb') as f: | |
| base64_encoded_pdfs.append(base64.b64encode(f.read())) | |
| return res, _output_name, _output_name_pdf | |
| # Gradio Interface for OCR | |
| demo_ocr = gr.Interface( | |
| fn=greet, | |
| inputs=[ | |
| gr.Image(type="pil"), | |
| gr.Checkbox(label="Apply Grammar Correction"), | |
| gr.Checkbox(label="Apply Spell Check"), | |
| gr.Dropdown(["English","Hindi","Punjabi"], label="Select Language", value="English") | |
| ], | |
| outputs=[ | |
| gr.Textbox(label="OCR Text"), | |
| gr.File(label="Text file"), | |
| gr.File(label="Searchable PDF File(English only)") | |
| ], | |
| title="OCR with Grammar and Spell Check", | |
| description="Upload an image to get the OCR results. Optionally, apply grammar and spell check.", | |
| examples=[ | |
| ["Examples/Book.png",False, False, "English"], | |
| ["Examples/News.png",False, False, "English"], | |
| ["Examples/Manuscript.jpg",False, False, "English"], | |
| ["Examples/Files.jpg",False, False, "English"], | |
| ["Examples/Hindi.jpg",False, False, "Hindi"], | |
| ["Examples/Hindi-manu.jpg",False, False, "Hindi"], | |
| ["Examples/Punjabi_machine.png",False, False, "Punjabi"]], | |
| cache_examples=False | |
| ) | |
| # demo_ocr.launch(debug=True) | |
| def split_text_into_batches(text, max_tokens_per_batch): | |
| sentences = nltk.sent_tokenize(text) # Tokenize text into sentences | |
| batches = [] | |
| current_batch = "" | |
| for sentence in sentences: | |
| if len(current_batch) + len(sentence) + 1 <= max_tokens_per_batch: # Add 1 for space | |
| current_batch += sentence + " " # Add sentence to current batch | |
| else: | |
| batches.append(current_batch.strip()) # Add current batch to batches list | |
| current_batch = sentence + " " # Start a new batch with the current sentence | |
| if current_batch: | |
| batches.append(current_batch.strip()) # Add the last batch | |
| return batches | |
| def run_t2tt(file_uploader , input_text: str, source_language: str, target_language: str) -> (str, bytes): | |
| if file_uploader is not None: | |
| with open(file_uploader, 'r') as file: | |
| input_text=file.read() | |
| source_language_code = LANGUAGE_NAME_TO_CODE[source_language] | |
| target_language_code = LANGUAGE_NAME_TO_CODE[target_language] | |
| max_tokens_per_batch= 256 | |
| batches = split_text_into_batches(input_text, max_tokens_per_batch) | |
| translated_text = "" | |
| for batch in batches: | |
| text_inputs = processor(text=batch, src_lang=source_language_code, return_tensors="pt") | |
| output_tokens = model.generate(**text_inputs, tgt_lang=target_language_code) | |
| translated_batch = processor.decode(output_tokens[0].tolist(), skip_special_tokens=True) | |
| translated_text += translated_batch + " " | |
| output=translated_text.strip() | |
| _output_name = "result.txt" | |
| open(_output_name, 'w').write(output) | |
| return str(output), _output_name | |
| with gr.Blocks() as demo_t2tt: | |
| with gr.Row(): | |
| with gr.Column(): | |
| with gr.Group(): | |
| file_uploader = gr.File(label="Upload a text file (Optional)") | |
| input_text = gr.Textbox(label="Input text") | |
| with gr.Row(): | |
| source_language = gr.Dropdown( | |
| label="Source language", | |
| choices=TEXT_SOURCE_LANGUAGE_NAMES, | |
| value="Punjabi", | |
| ) | |
| target_language = gr.Dropdown( | |
| label="Target language", | |
| choices=T2TT_TARGET_LANGUAGE_NAMES, | |
| value=DEFAULT_TARGET_LANGUAGE, | |
| ) | |
| btn = gr.Button("Translate") | |
| with gr.Column(): | |
| output_text = gr.Textbox(label="Translated text") | |
| output_file = gr.File(label="Translated text file") | |
| gr.Examples( | |
| examples=[ | |
| [ | |
| None, | |
| "The annual harvest festival of Baisakhi in Punjab showcases the region's agricultural prosperity and cultural vibrancy. This joyful occasion brings together people of all ages to celebrate with traditional music, dance, and feasts, reflecting the enduring spirit and community bond of Punjab's people", | |
| "English", | |
| "Punjabi", | |
| ], | |
| [ | |
| None, | |
| "It contains. much useful information about administrative, revenue, judicial and ecclesiastical activities in various areas which, it is hoped, would supplement the information available in official records.", | |
| "English", | |
| "Hindi", | |
| ], | |
| [ | |
| None, | |
| "दुनिया में बहुत सी अलग-अलग भाषाएं हैं और उनमें अपने वर्ण और शब्दों का भंडार होता है. इसमें में कुछ उनके अपने शब्द होते हैं तो कुछ ऐसे भी हैं, जो दूसरी भाषाओं से लिए जाते हैं.", | |
| "Hindi", | |
| "Punjabi", | |
| ], | |
| [ | |
| None, | |
| "ਸੂੂਬੇ ਦੇ ਕਈ ਜ਼ਿਲ੍ਹਿਆਂ ’ਚ ਬੁੱਧਵਾਰ ਸਵੇਰੇ ਸੰਘਣੀ ਧੁੰਦ ਛਾਈ ਰਹੀ ਤੇ ਤੇਜ਼ ਹਵਾਵਾਂ ਨੇ ਕਾਂਬਾ ਹੋਰ ਵਧਾ ਦਿੱਤਾ। ਸੱਤ ਸ਼ਹਿਰਾਂ ’ਚ ਦਿਨ ਦਾ ਤਾਪਮਾਨ ਦਸ ਡਿਗਰੀ ਸੈਲਸੀਅਸ ਦੇ ਆਸਪਾਸ ਰਿਹਾ। ਸੂਬੇ ’ਚ ਵੱਧ ਤੋਂ ਵੱਧ ਤਾਪਮਾਨ ’ਚ ਵੀ ਦਸ ਡਿਗਰੀ ਸੈਲਸੀਅਸ ਦੀ ਗਿਰਾਵਟ ਦਰਜ ਕੀਤੀ ਗਈ", | |
| "Punjabi", | |
| "English", | |
| ], | |
| ], | |
| inputs=[file_uploader ,input_text, source_language, target_language], | |
| outputs=[output_text, output_file], | |
| fn=run_t2tt, | |
| cache_examples=False, | |
| api_name=False, | |
| ) | |
| gr.on( | |
| triggers=[input_text.submit, btn.click], | |
| fn=run_t2tt, | |
| inputs=[file_uploader, input_text, source_language, target_language], | |
| outputs=[output_text, output_file], | |
| api_name="t2tt", | |
| ) | |
| with gr.Blocks() as demo: | |
| with gr.Tabs(): | |
| with gr.Tab(label="OCR"): | |
| demo_ocr.render() | |
| with gr.Tab(label="Translate"): | |
| demo_t2tt.render() | |
| if __name__ == "__main__": | |
| demo.launch() |