main.py

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

import time 
from PIL import Image
import numpy as np
import logging
import pandas as pd
from bs4 import BeautifulSoup

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
    #
    if not os.path.exists(debug_folder):
        os.makedirs(debug_folder)
    images = convert_from_path(test_file)
    ocr = OCRComponent1(englishFlag)
    ocr_results = {}
    for page_number,img in enumerate(images):

        
        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)
                
        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)

    return ocr_results


def test_tableOcrOnly1(test_file = './cropped_table_0.png' , debug_folder = './res/table1/',denoise = False,englishFlag = False):
    if not os.path.exists(debug_folder):
        os.makedirs(debug_folder)
    #Hybrid Unitable +DocTR
    #Good at these kind of tables - with a lot of texts 
    table = OcrTable1(englishFlag)
    image = Image.open(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)
    return table_code

def test_tableOcrOnly2(test_file = './cropped_table_1.png' , debug_folder = './res/table2/'):
    if not os.path.exists(debug_folder):
        os.makedirs(debug_folder)
    table = OcrTable2()

    #FullUnitable
    #Good at these kind of tables - with not much text 
 
    image = Image.open(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)
    if not os.path.exists(debug_folder):
        os.makedirs(debug_folder)

    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)

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
    if not os.path.exists(debug_folder):
        os.makedirs(debug_folder)
    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',encoding='utf-8') as file:
                file.write(table_code)

def test_table_component3(test_file = 'TestingFiles/TableOCRTestEnglish.pdf',debug_folder ='./res/table_debug3/',denoise = False,englishFlag = True):
    if not os.path.exists(debug_folder):
        os.makedirs(debug_folder)
    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',encoding='utf-8') as file:
                file.write(table_code)



def test_table_component4(test_file = 'TestingFiles/TableOCRTestEnglish.pdf',debug_folder ='./res/table_debug3/'):
    if not os.path.exists(debug_folder):
        os.makedirs(debug_folder)
    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',encoding='utf-8') as file:
                file.write(table_code)


if __name__ == '__main__':
    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)

    #test_table_component1()
    #test_ocr_component1("TestingFilesImages/OCRTest3English_0.jpg")
    #test_tableOcrOnly2()
    #test_tableOcrOnly2()
    #test_tableOcrOnly2_singleImage()

    # Example run 
    # python main.py pdftable2 --test_file TestingFiles/TableOCRTestEnglish.pdf --debug_foler ./res/table_debug2/ 
    # python main.py ocr1 --test_file TestingFiles/OCRTest1German.pdf --debug_repo ./res/ocrdebug1/ 
    # python main.py table1 --test_file ./cropped_table_0.png  --debug_repo ./res/table1/ 
    # python main.py table2 --test_file ./cropped_table_1.png  --debug_repo ./res/table2/ 
    # python main.py pdftable1 --test_file TestingFiles/TableOCRTestEnglish.pdf --debug_foler ./res/table_debug2/ 


    
    
    
    
    end = time.time()
    print("The entire pipeline took " , end-start)