Armandoliv/document_parser

1import os2import os3os.system('pip install "detectron2@git+https://github.com/facebookresearch/detectron2.git@v0.5#egg=detectron2"')4 5import io6import pandas as pd7import numpy as np8import gradio as gr9 10## for plotting11import matplotlib.pyplot as plt12 13## for ocr14import pdf2image15import cv216import layoutparser as lp17 18from docx import Document19from docx.shared import Inches20 21 22def parse_doc(dic):23    for k,v in dic.items():24        if "Title" in k:25            print('\x1b[1;31m'+ v +'\x1b[0m')26        elif "Figure" in k:27            plt.figure(figsize=(10,5))28            plt.imshow(v)29            plt.show()30        else:31            print(v)32        print(" ")33 34 35def to_image(filename):36  doc = pdf2image.convert_from_path(filename, dpi=350, last_page=1)37  # Save imgs38  folder = "doc"39  if folder not in os.listdir():40      os.makedirs(folder)41 42  p = 143  for page in doc:44      image_name = "page_"+str(p)+".jpg"  45      page.save(os.path.join(folder, image_name), "JPEG")46      p = p+147 48  return doc49 50 51 52def detect(doc):53  # General54  model = lp.Detectron2LayoutModel("lp://PubLayNet/mask_rcnn_X_101_32x8d_FPN_3x/config",55                                 extra_config=["MODEL.ROI_HEADS.SCORE_THRESH_TEST", 0.8],56                                 label_map={0:"Text", 1:"Title", 2:"List", 3:"Table", 4:"Figure"})57  ## turn img into array58  img = np.asarray(doc[0])59 60  ## predict61  detected = model.detect(img)62 63 64  return img, detected 65 66 67# sort detected68def split_page(img, n, axis):69    new_detected, start = [], 070    for s in range(n):71        end = len(img[0])/3 * s if axis == "x" else len(img[1])/372        section = lp.Interval(start=start, end=end, axis=axis).put_on_canvas(img)73        filter_detected = detected.filter_by(section, center=True)._blocks74        new_detected = new_detected + filter_detected75        start = end76    return lp.Layout([block.set(id=idx) for idx,block in enumerate(new_detected)])77 78 79 80def get_detected(img, detected):81  n_cols,n_rows = 1,182 83  ## if single page just sort based on y84  if (n_cols == 1) and (n_rows == 1):85      new_detected = detected.sort(key=lambda x: x.coordinates[1])86      detected = lp.Layout([block.set(id=idx) for idx,block in enumerate(new_detected)])87      88  ## if multi columns sort by x,y89  elif (n_cols > 1) and (n_rows == 1):90      detected = split_page(img, n_cols, axis="x")91 92  ## if multi rows sort by y,x93  elif (n_cols > 1) and (n_rows == 1):94      detected = split_page(img, n_rows, axis="y")95      96  ## if multi columns-rows97  else:98      pass99  100  return detected101 102 103def predict_elements(img, detected)->dict:104  model = lp.TesseractAgent(languages='eng')105  dic_predicted = {}106 107  for block in [block for block in detected if block.type in ["Title","Text", "List"]]:108    ## segmentation109    segmented = block.pad(left=15, right=15, top=5, bottom=5).crop_image(img)110    ## extraction111    extracted = model.detect(segmented)112    ## save113    dic_predicted[str(block.id)+"-"+block.type] = extracted.replace('\n',' ').strip()114 115  for block in [block for block in detected if block.type == "Figure"]:116      ## segmentation117      segmented = block.pad(left=15, right=15, top=5, bottom=5).crop_image(img)118      ## save119      dic_predicted[str(block.id)+"-"+block.type] = segmented120 121 122  for block in [block for block in detected if block.type == "Table"]:123      ## segmentation124      segmented = block.pad(left=15, right=15, top=5, bottom=5).crop_image(img)125      ## extraction126      extracted = model.detect(segmented)127      ## save128      dic_predicted[str(block.id)+"-"+block.type] = pd.read_csv( io.StringIO(extracted) )129 130  131  return dic_predicted132 133def gen_doc(dic_predicted:dict):134  document = Document()135 136  for k,v in dic_predicted.items():137 138    if "Figure" in k:139      cv2.imwrite(f'{k}.jpg', dic_predicted[k])140      document.add_picture(f'{k}.jpg', width=Inches(3))141 142    elif "Table" in k:143      table = document.add_table(rows=v.shape[0], cols=v.shape[1])144      hdr_cells = table.rows[0].cells145      for idx, col in enumerate(v.columns):146        hdr_cells[idx].text = col147      for c in v.iterrows():148        149        for idx, col in enumerate(v.columns):150          try:151            if len(c[1][col].strip())>0:152              row_cells = table.add_row().cells153              row_cells[idx].text = str(c[1][col]) 154          except:155            continue156    157    else:158      document.add_paragraph(str(v))159 160  document.save('demo.docx')161 162 163def main_convert(filename):164  print(filename.name)165  doc = to_image(filename.name)166 167  img, detected = detect(doc)168 169  n_detected = get_detected(img, detected)170 171  dic_predicted = predict_elements(img, n_detected)172 173  gen_doc(dic_predicted)174 175  im_out = lp.draw_box(img, detected, box_width=5, box_alpha=0.2, show_element_type=True)176  dict_out = {}177  for k,v in dic_predicted.items():178    if "figure" not in k.lower():179      dict_out[k] = dic_predicted[k]180 181  return  'demo.docx', im_out, dict_out182  183  184inputs = [gr.File(type='file', label="Original PDF File")]185outputs = [gr.File(label="Converted DOC File"),gr.Image(type="PIL.Image", label="Detected Image"),  gr.JSON()]186 187title = "A Document AI parser"188description = "This demo uses AI Models to detect text, titles, tables, figures and lists as well as table cells from an Scanned document.\nBased on the layout it determines reading order and generates an MS-DOC file to Download."189 190 191io = gr.Interface(fn=main_convert, inputs=inputs, outputs=outputs, title=title, description=description, 192                  css= """.gr-button-primary { background: -webkit-linear-gradient( 193                    90deg, #355764 0%, #55a8a1 100% ) !important;     background: #355764;194                        background: linear-gradient( 195                    90deg, #355764 0%, #55a8a1 100% ) !important;196                        background: -moz-linear-gradient( 90deg, #355764 0%, #55a8a1 100% ) !important;197                        background: -webkit-linear-gradient( 198                    90deg, #355764 0%, #55a8a1 100% ) !important;199                    color:white !important}"""200                  )201                  202io.launch()