visualizingjp/pdf-ocr
3
1# -*- coding: utf-8 -*-2"""3 4easyocr.py - A wrapper for easyocr to convert pdf to images to text5"""6 7import logging8from pathlib import Path9 10logging.basicConfig(11 level=logging.INFO,12 format="%(asctime)s %(levelname)s %(message)s",13 datefmt="%m/%d/%Y %I:%M:%S",14)15 16 17import os18import pprint as pp19import re20import shutil21import time22from datetime import date, datetime23from os.path import basename, dirname, join24from pathlib import Path25 26from cleantext import clean27from doctr.io import DocumentFile28from doctr.models import ocr_predictor29from libretranslatepy import LibreTranslateAPI30from natsort import natsorted31from spellchecker import SpellChecker32from tqdm.auto import tqdm33 34 35def simple_rename(filepath, target_ext=".txt"):36 _fp = Path(filepath)37 basename = _fp.stem38 return f"OCR_{basename}_{target_ext}"39 40 41def rm_local_text_files(name_contains="RESULT_"):42 """43 rm_local_text_files - remove local text files44 45 Args:46 name_contains (str, optional): [description]. Defaults to "OCR_".47 """48 files = [49 f50 for f in Path.cwd().iterdir()51 if f.is_file() and f.suffix == ".txt" and name_contains in f.name52 ]53 logging.info(f"removing {len(files)} text files")54 for f in files:55 os.remove(f)56 logging.info("done")57 58 59def corr(60 s: str,61 add_space_when_numerics=False,62 exceptions=["e.g.", "i.e.", "etc.", "cf.", "vs.", "p."],63) -> str:64 """corrects spacing in a string65 66 Args:67 s (str): the string to correct68 add_space_when_numerics (bool, optional): [add a space when a period is between two numbers, example 5.73]. Defaults to False.69 exceptions (list, optional): [do not change these substrings]. Defaults to ['e.g.', 'i.e.', 'etc.', 'cf.', 'vs.', 'p.'].70 71 Returns:72 str: the corrected string73 """74 if add_space_when_numerics:75 s = re.sub(r"(\d)\.(\d)", r"\1. \2", s)76 77 s = re.sub(r"\s+", " ", s)78 s = re.sub(r'\s([?.!"](?:\s|$))', r"\1", s)79 80 # fix space before apostrophe81 s = re.sub(r"\s\'", r"'", s)82 # fix space after apostrophe83 s = re.sub(r"'\s", r"'", s)84 # fix space before comma85 s = re.sub(r"\s,", r",", s)86 87 for e in exceptions:88 expected_sub = re.sub(r"\s", "", e)89 s = s.replace(expected_sub, e)90 91 return s92 93 94def fix_punct_spaces(string):95 """96 fix_punct_spaces - replace spaces around punctuation with punctuation. For example, "hello , there" -> "hello, there"97 98 Parameters99 ----------100 string : str, required, input string to be corrected101 102 Returns103 -------104 str, corrected string105 """106 107 fix_spaces = re.compile(r"\s*([?!.,]+(?:\s+[?!.,]+)*)\s*")108 string = fix_spaces.sub(lambda x: "{} ".format(x.group(1).replace(" ", "")), string)109 string = string.replace(" ' ", "'")110 string = string.replace(' " ', '"')111 return string.strip()112 113 114def clean_OCR(ugly_text: str):115 """116 clean_OCR - clean the OCR text files.117 118 Parameters119 ----------120 ugly_text : str, required, input string to be cleaned121 122 Returns123 -------124 str, cleaned string125 """126 # Remove all the newlines.127 cleaned_text = ugly_text.replace("\n", " ")128 # Remove all the tabs.129 cleaned_text = cleaned_text.replace("\t", " ")130 # Remove all the double spaces.131 cleaned_text = cleaned_text.replace(" ", " ")132 # Remove all the spaces at the beginning of the text.133 cleaned_text = cleaned_text.lstrip()134 # remove all instances of "- " and " - "135 cleaned_text = cleaned_text.replace("- ", "")136 cleaned_text = cleaned_text.replace(" -", "")137 return fix_punct_spaces(cleaned_text)138 139 140def move2completed(from_dir, filename, new_folder="completed", verbose=False):141 142 # this is the better version143 old_filepath = join(from_dir, filename)144 145 new_filedirectory = join(from_dir, new_folder)146 147 if not os.path.isdir(new_filedirectory):148 os.mkdir(new_filedirectory)149 if verbose:150 print("created new directory for files at: \n", new_filedirectory)151 new_filepath = join(new_filedirectory, filename)152 153 try:154 shutil.move(old_filepath, new_filepath)155 logging.info("successfully moved the file {} to */completed.".format(filename))156 except:157 logging.info(158 "ERROR! unable to move file to \n{}. Please investigate".format(159 new_filepath160 )161 )162 163 164"""## pdf2text functions165 166"""167 168 169custom_replace_list = {170 "t0": "to",171 "'$": "'s",172 ",,": ", ",173 "_ ": " ",174 " '": "'",175}176 177replace_corr_exceptions = {178 "i. e.": "i.e.",179 "e. g.": "e.g.",180 "e. g": "e.g.",181 " ,": ",",182}183 184 185spell = SpellChecker()186 187 188def check_word_spelling(word: str) -> bool:189 """190 check_word_spelling - check the spelling of a word191 192 Args:193 word (str): word to check194 195 Returns:196 bool: True if word is spelled correctly, False if not197 """198 199 misspelled = spell.unknown([word])200 201 return len(misspelled) == 0202 203 204def eval_and_replace(text: str, match_token: str = "- ") -> str:205 """206 eval_and_replace - conditionally replace all instances of a substring in a string based on whether the eliminated substring results in a valid word207 208 Args:209 text (str): text to evaluate210 match_token (str, optional): token to replace. Defaults to "- ".211 212 Returns:213 str: text with replaced tokens214 """215 216 if match_token not in text:217 return text218 else:219 while True:220 full_before_text = text.split(match_token, maxsplit=1)[0]221 before_text = [222 char for char in full_before_text.split()[-1] if char.isalpha()223 ]224 before_text = "".join(before_text)225 full_after_text = text.split(match_token, maxsplit=1)[-1]226 after_text = [char for char in full_after_text.split()[0] if char.isalpha()]227 after_text = "".join(after_text)228 full_text = before_text + after_text229 if check_word_spelling(full_text):230 text = full_before_text + full_after_text231 else:232 text = full_before_text + " " + full_after_text233 if match_token not in text:234 break235 return text236 237 238def cleantxt_ocr(ugly_text, lower=False, lang: str = "en") -> str:239 """240 cleantxt_ocr - clean text from OCR241 242 Args:243 ugly_text (str): text to clean244 lower (bool, optional): _description_. Defaults to False.245 lang (str, optional): _description_. Defaults to "en".246 247 Returns:248 str: cleaned text249 """250 # a wrapper for clean text with options different than default251 252 # https://pypi.org/project/clean-text/253 cleaned_text = clean(254 ugly_text,255 fix_unicode=True, # fix various unicode errors256 to_ascii=True, # transliterate to closest ASCII representation257 lower=lower, # lowercase text258 no_line_breaks=True, # fully strip line breaks as opposed to only normalizing them259 no_urls=True, # replace all URLs with a special token260 no_emails=True, # replace all email addresses with a special token261 no_phone_numbers=False, # replace all phone numbers with a special token262 no_numbers=False, # replace all numbers with a special token263 no_digits=False, # replace all digits with a special token264 no_currency_symbols=False, # replace all currency symbols with a special token265 no_punct=False, # remove punctuations266 replace_with_punct="", # instead of removing punctuations you may replace them267 replace_with_url="<URL>",268 replace_with_email="<EMAIL>",269 replace_with_phone_number="<PHONE>",270 replace_with_number="<NUM>",271 replace_with_digit="0",272 replace_with_currency_symbol="<CUR>",273 lang=lang, # set to 'de' for German special handling274 )275 276 return cleaned_text277 278 279def format_ocr_out(OCR_data):280 281 if isinstance(OCR_data, list):282 text = " ".join(OCR_data)283 else:284 text = str(OCR_data)285 _clean = cleantxt_ocr(text)286 return corr(_clean)287 288 289def postprocess(text: str) -> str:290 """to be used after recombining the lines"""291 292 proc = corr(cleantxt_ocr(text))293 294 for k, v in custom_replace_list.items():295 proc = proc.replace(str(k), str(v))296 297 proc = corr(proc)298 299 for k, v in replace_corr_exceptions.items():300 proc = proc.replace(str(k), str(v))301 302 return eval_and_replace(proc)303 304 305def result2text(result, as_text=False) -> str or list:306 """Convert OCR result to text"""307 308 full_doc = []309 for i, page in enumerate(result.pages, start=1):310 text = ""311 for block in page.blocks:312 text += "\n\t"313 for line in block.lines:314 for word in line.words:315 # print(dir(word))316 text += word.value + " "317 full_doc.append(text)318 319 return "\n".join(full_doc) if as_text else full_doc320 321 322def convert_PDF_to_Text(323 PDF_file,324 ocr_model=None,325 max_pages: int = 20,326):327 328 st = time.perf_counter()329 PDF_file = Path(PDF_file)330 ocr_model = ocr_predictor(pretrained=True) if ocr_model is None else ocr_model331 logging.info(f"starting OCR on {PDF_file.name}")332 doc = DocumentFile.from_pdf(PDF_file)333 truncated = False334 if len(doc) > max_pages:335 logging.warning(336 f"PDF has {len(doc)} pages, which is more than {max_pages}.. truncating"337 )338 doc = doc[:max_pages]339 truncated = True340 341 # Analyze342 logging.info(f"running OCR on {len(doc)} pages")343 result = ocr_model(doc)344 raw_text = result2text(result)345 proc_text = [format_ocr_out(r) for r in raw_text]346 fin_text = [postprocess(t) for t in proc_text]347 348 ocr_results = "\n\n".join(fin_text)349 350 fn_rt = time.perf_counter() - st351 352 logging.info("OCR complete")353 354 results_dict = {355 "num_pages": len(doc),356 "runtime": round(fn_rt, 2),357 "date": str(date.today()),358 "converted_text": ocr_results,359 "truncated": truncated,360 "length": len(ocr_results),361 }362 363 return results_dict364 365 366# @title translation functions367 368lt = LibreTranslateAPI("https://translate.astian.org/")369 370 371def translate_text(text, source_l, target_l="en"):372 373 return str(lt.translate(text, source_l, target_l))374 375 376def translate_doc(filepath, lang_start, lang_end="en", verbose=False):377 """translate a document from lang_start to lang_end378 379 {'code': 'en', 'name': 'English'},380 {'code': 'fr', 'name': 'French'},381 {'code': 'de', 'name': 'German'},382 {'code': 'it', 'name': 'Italian'},"""383 384 src_folder = dirname(filepath)385 src_folder = Path(src_folder)386 trgt_folder = src_folder / f"translated_{lang_end}"387 trgt_folder.mkdir(exist_ok=True)388 with open(filepath, "r", encoding="utf-8", errors="ignore") as f:389 foreign_t = f.readlines()390 in_name = basename(filepath)391 translated_doc = []392 for line in tqdm(393 foreign_t, total=len(foreign_t), desc="translating {}...".format(in_name[:10])394 ):395 translated_line = translate_text(line, lang_start, lang_end)396 translated_doc.append(translated_line)397 t_out_name = "[To {}]".format(lang_end) + simple_rename(in_name) + ".txt"398 out_path = join(trgt_folder, t_out_name)399 with open(out_path, "w", encoding="utf-8", errors="ignore") as f_o:400 f_o.writelines(translated_doc)401 if verbose:402 print("finished translating the document! - ", datetime.now())403 return out_path404