app.py

import os
import traceback
import argparse
from typing import List, Tuple, Set, Dict

import time 
from PIL import Image
import numpy as np
from doctr.models import ocr_predictor
import logging
import pandas as pd
from bs4 import BeautifulSoup
import gradio

from utils import cropImages
from utils import draw_only_box,draw_box_with_text,getlogger,Annotation
from ocr_component1 import OCRComponent1
from detectionAndOcrTable1 import DetectionAndOcrTable1
from detectionAndOcrTable2 import DetectionAndOcrTable2
from detectionAndOcrTable3 import DetectionAndOcrTable3
from detectionAndOcrTable4 import DetectionAndOcrTable4
from ocrTable1 import OcrTable1
from ocrTable2 import OcrTable2 
from pdf2image import convert_from_path


def convertHTMLToCSV(html:str,output_path:str)->str:
    
    # empty list
    data = []
    
    # for getting the header from
    # the HTML file
    list_header = []
    soup = BeautifulSoup(html,'html.parser')
    header = soup.find_all("table")[0].find("tr")
    
    for items in header:
        try:
            list_header.append(items.get_text())
        except:
            continue
    
    # for getting the data 
    HTML_data = soup.find_all("table")[0].find_all("tr")[1:]
    
    for element in HTML_data:
        sub_data = []
        for sub_element in element:
            try:
                sub_data.append(sub_element.get_text())
            except:
                continue
        data.append(sub_data)
    
    # Storing the data into Pandas
    # DataFrame 
    dataFrame = pd.DataFrame(data = data, columns = list_header)
    
    # Converting Pandas DataFrame
    # into CSV file
    dataFrame.to_csv(output_path)

def saveResults(image_list, results, labels, output_dir='output/', threshold=0.5):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    for idx, im in enumerate(image_list):
        im = draw_only_box(im, results[idx], labels, threshold=threshold)

        out_path = os.path.join(output_dir, f"{idx}.jpg")
        im.save(out_path, quality=95)
        print("save result to: " + out_path)

def InputToImages(input_path:str,resolution=300)-> List[Image.Image]:
    """
    input is file location to image 
    return : List of Pillow image objects
    """
    images=[]
    try:
        img =Image.open(input_path)
        if img.mode == 'RGBA':
            img = img.convert('RGB')
        images.append(img)
    except Exception as e:
        traceback.print_exc()
    return images

def drawTextDetRes(bxs :List[List[float]],img:Image.Image,output_path:str):
    """
    draw layout analysis results 
    """
    """bxs_draw is xmin, ymin, xmax, ymax"""
    bxs_draw = [[b[0][0], b[0][1], b[1][0], b[-1][1]] for b in bxs if b[0][0] <= b[1][0] and b[0][1] <= b[-1][1]]
    
    #images_to_recognizer = cropImage(bxs, img)
    img_to_save = draw_only_box(img, bxs_draw)
    img_to_save.save(output_path, quality=95)

def test_ocr_component1(test_file="TestingFiles/OCRTest1German.pdf", debug_folder = './res/table1/',englishFlag = False):
    #Takes as input image of a single page and returns the detected lines and words

    images = convert_from_path(test_file)
    ocr = OCRComponent1(englishFlag)
    ocr_results = {}

    all_text_in_pages = {}
    for page_number,img in enumerate(images):
        text_in_page = ""
        
        line_annotations= ocr.predict(img = np.array(img))
        ocr_results[page_number] = line_annotations

        """
        boxes_to_draw =[]
        for list_of_ann in word_annotations:
            for ann in list_of_ann:
                logger.info(ann.text)
                b = ann.box
                boxes_to_draw.append(b)

        img_to_save = draw_only_box(img,boxes_to_draw)
        img_to_save.save("res/12June_2_lines.png", quality=95)
        """

        line_boxes_to_draw =[]
        #print("Detected lines are ")
        #print(len(line_annotations.items()))
        for index,ann in line_annotations.items():

            b = ann.box
            line_boxes_to_draw.append(b)
            line_words = ""
            #print("detected words per line")
            #print(len(ann.words))
            for wordann in ann.words:
                line_words += wordann.text +" "
            print(line_words)
            text_in_page += line_words +"\n"
                
        img_to_save1 = draw_only_box(img,line_boxes_to_draw)
        imgname = test_file.split("/")[-1][:-4]
        img_to_save1.save(debug_folder+imgname+"_"+str(page_number)+"_bbox_detection.png", quality=95)

        all_text_in_pages[page_number] = text_in_page

    return ocr_results, all_text_in_pages


def test_tableOcrOnly1(test_file :Image.Image , debug_folder = './res/table1/',denoise = False,englishFlag = False):
    #Hybrid Unitable +DocTR
    #Good at these kind of tables - with a lot of texts 
    table = OcrTable1(englishFlag)
    image = test_file.convert("RGB")
    """
    parts = test_file.split("/")
    filename = parts[-1][:-4]
    debugfolder_filename_page_name= debug_folder+filename+"_"
    
    table_code = table.predict([image],debugfolder_filename_page_name,denoise = denoise)
    with open(debugfolder_filename_page_name+'output.txt', 'w') as file:
        file.write(table_code)
    """

    table_code = table.predict([image],denoise = denoise)
    return table_code


def test_tableOcrOnly2(test_file:Image.Image, debug_folder = './res/table2/'):
    table = OcrTable2()
    #FullUnitable
    #Good at these kind of tables - with not much text 
 
    image = test_file.convert("RGB")
    table.predict([image],debug_folder)

def test_table_component1(test_file = 'TestingFiles/TableOCRTestEnglish.pdf', debug_folder ='./res/table_debug2/',denoise = False,englishFlag = True):
    table_predictor = DetectionAndOcrTable1(englishFlag)

    images = convert_from_path(test_file)
    for page_number,img in enumerate(images):

        #print(img.mode)
        print("Looking at page:")
        print(page_number)
        parts = test_file.split("/")
        filename = parts[-1][:-4]
        debugfolder_filename_page_name= debug_folder+filename+"_"+ str(page_number)+'_'
        table_codes = table_predictor.predict(img,debugfolder_filename_page_name=debugfolder_filename_page_name,denoise = denoise)
        for index, table_code in enumerate(table_codes):
            with open(debugfolder_filename_page_name+str(index)+'output.xls', 'w') as file:
                file.write(table_code)
        return table_codes

def test_table_component2(test_file = 'TestingFiles/TableOCRTestEnglish.pdf', debug_folder ='./res/table_debug2/'):
    #This components can take in entire pdf page as input , scan for tables and return the table in html format
    #Uses the full unitable model

    table_predictor = DetectionAndOcrTable2()

    images = convert_from_path(test_file)
    for page_number,img in enumerate(images):
        print("Looking at page:")
        print(page_number)
        parts = test_file.split("/")
        filename = parts[-1][:-4]
        debugfolder_filename_page_name= debug_folder+filename+"_"+ str(page_number)+'_'
        table_codes = table_predictor.predict(img,debugfolder_filename_page_name=debugfolder_filename_page_name)
        for index, table_code in enumerate(table_codes):
            with open(debugfolder_filename_page_name+str(index)+'output.xls', 'w') as file:
                file.write(table_code)
        return table_codes

def test_table_component3(test_file = 'TestingFiles/TableOCRTestEnglish.pdf',debug_folder ='./res/table_debug3/',denoise = False,englishFlag = True):
    table_predictor = DetectionAndOcrTable3(englishFlag)

    images = convert_from_path(test_file)
    for page_number,img in enumerate(images):
        #print(img.mode)
        print("Looking at page:")
        print(page_number)
        parts = test_file.split("/")
        filename = parts[-1][:-4]
        debugfolder_filename_page_name= debug_folder+filename+"_"+ str(page_number)+'_'
        table_codes = table_predictor.predict(img,debugfolder_filename_page_name=debugfolder_filename_page_name)
        for index, table_code in enumerate(table_codes):
            with open(debugfolder_filename_page_name+str(index)+'output.xls', 'w') as file:
                file.write(table_code)
        return table_codes



def test_table_component4(test_file = 'TestingFiles/TableOCRTestEnglish.pdf',debug_folder ='./res/table_debug3/'):
    table_predictor = DetectionAndOcrTable4()

    images = convert_from_path(test_file)
    for page_number,img in enumerate(images):
        #print(img.mode)
        print("Looking at page:")
        print(page_number)
        parts = test_file.split("/")
        filename = parts[-1][:-4]
        debugfolder_filename_page_name= debug_folder+filename+"_"+ str(page_number)+'_'
        table_codes = table_predictor.predict(img,debugfolder_filename_page_name=debugfolder_filename_page_name)
        for index, table_code in enumerate(table_codes):
            with open(debugfolder_filename_page_name+str(index)+'output.xls', 'w') as file:
                file.write(table_code)
        return table_codes


"""
    parser = argparse.ArgumentParser(description='Process some strings.')
    parser.add_argument('ocr', type=str, help='type in id of the component to test')
    parser.add_argument('--test_file',type=str, help='path to the testing file')
    parser.add_argument('--debug_folder',type=str, help='path to the folder you want to save your results in')
    parser.add_argument('--englishFlag',type=bool, help='Whether your pdf is in english => could lead to better results ')
    parser.add_argument('--denoise',type=bool, help='preprocessing for not clean scans  ')

    args = parser.parse_args()
    start = time.time()
    if args.ocr == "ocr1":
        test_ocr_component1(args.test_file,args.debug_folder, args.englishFlag)
    elif args.ocr == "table1":
        test_tableOcrOnly1(args.test_file,args.debug_folder,args.englishFlag,args.denoise)
    elif args.ocr == "table2":
        test_tableOcrOnly2(args.test_file,args.debug_folder)
    elif args.ocr =="pdftable1":
        test_table_component1(args.test_file,args.debug_folder,args.englishFlag,args.denoise)
    elif args.ocr =="pdftable2":
        test_table_component2(args.test_file,args.debug_folder)
    elif args.ocr =="pdftable3":
        test_table_component3(args.test_file,args.debug_folder,args.englishFlag,args.denoise)
    elif args.ocr =="pdftable4":
        test_table_component4(args.test_file,args.debug_folder)

"""
import gradio as gr
from gradio_pdf import PDF

with gr.Blocks() as demo:
    gr.Markdown("# OCR component")
    inputs_for_ocr = [PDF(label="Document"), gr.Textbox(label="internal debug folder",placeholder = "./res/table1/"),gr.Checkbox(label ="English Document?",value =False)]
    ocr_btn = gr.Button("Run ocr")
   
    gr.Examples(
        examples=[["TestingFiles/OCRTest1German.pdf",'./res/table1/',False]],
        inputs=inputs_for_ocr
                )
    
    outputs_for_ocr = [gr.Textbox(label="List of annotation objects"), gr.Textbox("Text in page")]

    ocr_btn.click(fn=test_ocr_component1, 
                  inputs = inputs_for_ocr,
                  outputs = outputs_for_ocr,
                  api_name="OCR"
                  )
    
    gr.Markdown("# Table OCR components that takes a pdf, extract table and return their html code ")
    gr.Markdown("## Component 1 uses table transformer and doctr +Unitable")
    inputs_for_pdftable1 = [PDF(label="Document"), gr.Textbox(label="internal debug folder",placeholder = "./res/table1/"),gr.Checkbox(label ="Denoise?",value =False),gr.Checkbox(label ="English Document?",value =False)]
    table1_btn = gr.Button("Run pdftable1")

    gr.Examples(
        examples=[["TestingFiles/OCRTest5English.pdf",'./res/table1/',False]],
        inputs=inputs_for_pdftable1
                )
    outputs_for_pdftable1 = [gr.Textbox(label="Table code")]
    
    table1_btn.click(fn=test_table_component1,
                    inputs = inputs_for_pdftable1,
                    outputs = outputs_for_pdftable1,
                    api_name="pdfTable1"
                    )
    
    gr.Markdown("## Component 2 uses table transformer and Unitable")
    inputs_for_pdftable2 = [PDF(label="Document"), gr.Textbox(label="internal debug folder",placeholder = "./res/table1/")]
    table2_btn = gr.Button("Run pdftable2")

    gr.Examples(
        examples=[["TestingFiles/OCRTest5English.pdf",'./res/table1/',False]],
        inputs=inputs_for_pdftable1
                )
    outputs_for_pdftable2 = [gr.Textbox(label="Table code")]

    table2_btn.click(fn=test_table_component2,
                    inputs = inputs_for_pdftable2,
                    outputs = outputs_for_pdftable2,
                    api_name="pdfTable2"
                    )
    
    gr.Markdown("## Component 3 uses Yolo and Unitable+doctr")
    inputs_for_pdftable3 = [PDF(label="Document"), gr.Textbox(label="internal debug folder",placeholder = "./res/table1/"),gr.Checkbox(label ="Denoise?",value =False),gr.Checkbox(label ="English Document?",value =False)]
    table3_btn = gr.Button("Run pdftable3")
    

    gr.Examples(
        examples=[["TestingFiles/TableOCRTestEnglish.pdf",'./res/table1/',False]],
        inputs=inputs_for_pdftable1
                )
    outputs_for_pdftable3 = [gr.Textbox(label="Table code")]

    table3_btn.click(fn=test_table_component3,
                    inputs = inputs_for_pdftable3,
                    outputs = outputs_for_pdftable3,
                    api_name="pdfTable3"
                    )
    
    gr.Markdown("## Component 4 uses Yolo and Unitable")
    inputs_for_pdftable4 = [PDF(label="Document"), gr.Textbox(label="internal debug folder",placeholder = "./res/table1/")]
    table4_btn = gr.Button("Run pdftable4")
    
    gr.Examples(
        examples=[["TestingFiles/TableOCRTestEasier.pdf",'./res/table1/',False]],
        inputs=inputs_for_pdftable1
                )
    outputs_for_pdftable4 = [gr.Textbox(label="Table code")]


    table4_btn.click(fn=test_table_component4,
                    inputs = inputs_for_pdftable4,
                    outputs = outputs_for_pdftable4,
                    api_name="pdfTable4"
                    )
    

    gr.Markdown("# Table OCR component that takes image of an cropped tavle, extract table and return their html code ")

    inputs_for_table1 = [gr.Image(label="Image of cropped table",type='pil'), gr.Textbox(label="internal debug folder",placeholder = "./res/table1/"),gr.Checkbox(label ="Denoise?",value =False),gr.Checkbox(label ="English Document?",value =False)]
    onlytable1_btn = gr.Button("Run table1")
    
    gr.Examples(
        examples=[[Image.open("cropped_table.png"),'./res/table1/',False]],
        inputs=inputs_for_table1
                )
    outputs_for_table1 = [gr.HTML(label="Table code")]


    onlytable1_btn.click(fn=test_tableOcrOnly1,
                    inputs = inputs_for_table1,
                    outputs = outputs_for_table1,
                    api_name="table1"
                    )
    
    gr.Markdown("## Another Table OCR component that takes image of an cropped table, extract table and return their html code ")

    inputs_for_table2 = [gr.Image(label="Image of cropped table",type='pil'), gr.Textbox(label="internal debug folder",placeholder = "./res/table1/")]
    onlytable2_btn = gr.Button("Run table2")
    
    
    gr.Examples(
        examples=[[Image.open("cropped_table.png"),'./res/table1/',False]],
        inputs=inputs_for_table2
                )
    outputs_for_table2 = [gr.HTML(label="Table code")]
    
    onlytable2_btn.click(fn=test_tableOcrOnly2,
                    inputs = inputs_for_table2,
                    outputs = outputs_for_table2,
                    api_name="table2"
                    )


    

demo.launch(share=True)