import os from doctr.io import DocumentFile from doctr.models import ocr_predictor import gradio as gr from PIL import Image from happytransformer import HappyTextToText, TTSettings from transformers import AutoTokenizer, AutoModelForSeq2SeqLM import re from lang_list import ( LANGUAGE_NAME_TO_CODE, T2TT_TARGET_LANGUAGE_NAMES, TEXT_SOURCE_LANGUAGE_NAMES, ) DEFAULT_TARGET_LANGUAGE = "English" from transformers import SeamlessM4TForTextToText from transformers import AutoProcessor model = SeamlessM4TForTextToText.from_pretrained("facebook/hf-seamless-m4t-medium") processor = AutoProcessor.from_pretrained("facebook/hf-seamless-m4t-medium") # 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") 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): 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) return res, _output_name # 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") ], outputs=["text", "file"], title="DocTR 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"], ["Examples/News.png"], ["Examples/Manuscript.jpg"], ["Examples/Files.jpg"]] ) # demo_ocr.launch(debug=True) 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] text_inputs = processor(text = input_text, src_lang=source_language_code , return_tensors="pt") output_tokens = model.generate(**text_inputs, tgt_lang=target_language_code) output = processor.decode(output_tokens[0].tolist(), skip_special_tokens=True) _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 sinister destruction of the holy Akal Takht and the ruthless massacre of thousands of innocent pilgrims had unmasked the deep-seated hatred and animosity that the Indian Government had been nurturing against Sikhs ever since independence", "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()