pdf2data-tools 0.0.2__py3-none-any.whl

pdf2data/__init__.py

@@ -0,0 +1,5 @@
+"""Top-level package for pdf2data."""
+
+__author__ = """Daniel Pereira Costa"""
+__email__ = "daniel.pereira.costa@tecnico.ulisboa.pt"
+__version__ = "0.0.2"

pdf2data/cli/block_extractor.py

@@ -0,0 +1,157 @@
+import os
+from typing import List, Optional
+
+import click
+
+
+@click.command()
+@click.argument("input_folder", type=str)
+@click.argument("output_folder", type=str)
+@click.option(
+    "--pipeline",
+    default="MinerU", #NotDefined, MinerU
+    help="Define the pipeline to be used",
+)
+@click.option(
+    "--layout_model",
+    default="DocLayout-YOLO-DocStructBench",
+    help="model used to determine the overall layout",
+)
+@click.option(
+    "--table_model",
+    default=None,
+    help="model used to identify the tables",
+)
+@click.option(
+    "--layout_model_threshold",
+    default=0.6,
+    help="threshold of the layout model",
+)
+@click.option(
+    "--table_model_threshold",
+    default=0.6,
+    help="threshold of the table detection model",
+)
+@click.option(
+    "--struct_model",
+    default="microsoft/table-structure-recognition-v1.1-all",
+    help="table structure detection model",
+)
+@click.option(
+    "--device",
+    default="cpu",
+    help="device to run the mask models",
+)
+def block_extractor(
+    input_folder: str,
+    output_folder: str,
+    pipeline: Optional[str],
+    layout_model: str,
+    table_model: str,
+    layout_model_threshold: float,
+    table_model_threshold: float,
+    struct_model: str,
+    device: str,
+) -> None:
+    if pipeline == "NotDefined":
+        from pdf2data.pdf2data_pipeline import PDF2Data
+        pdf2data_pipeline: PDF2Data = PDF2Data(layout_model=layout_model, 
+                                                layout_model_threshold=layout_model_threshold,
+                                                table_model=table_model,
+                                                table_model_threshold=table_model_threshold,
+                                                table_structure_model=struct_model,
+                                                device=device, 
+                                                input_folder=input_folder, 
+                                                output_folder=output_folder,
+                                                extract_text=False,
+                                                extract_equations=False)
+        pdf2data_pipeline.pdf_transform()
+        """file_list: List[str] = get_doc_list(input_folder, "pdf")
+        extractor: BlockExtractor = BlockExtractor(
+            ocr_model=ocr_model,
+            word_detection_model=word_detection_model,
+            word_detection_threshold=word_detection_model_threshold,
+            extract_tables=extract_tables,
+            extract_figures=extract_figures,
+            correct_struct=correct_struct,
+            table_zoom=table_zoom,
+            cell_zoom=cell_zoom,
+            figure_zoom=figure_zoom,
+            x_table_corr=x_table_corr,
+            y_table_corr=y_table_corr,
+            iou_lines=iou_lines,
+            iou_struct=iou_struct,
+            word_factor=word_factor,
+            word_iou=word_iou,
+            structure_model=struct_model,
+            struct_model_threshold=struct_model_threshold,
+            reconstructor_type=reconstructor_type,
+            brightness=brightness,
+            contrast=contrast,
+            letter_ratio=letter_ratio
+        )
+        mask: LayoutParser = LayoutParser(
+            model=layout_model,
+            model_threshold=layout_model_threshold,
+            table_model=table_model,
+            table_model_threshold=table_model_threshold,
+            device_type=device
+        )
+        mask.model_post_init(None)
+        total_docs: int = len(file_list)
+        doc_number: int = 1
+        for file in file_list:
+            print(f'{doc_number}//{total_docs} processed')
+            print(file)
+            file_path: str = input_folder + "/" + file
+            layout = mask.get_layout(file_path)
+            extractor.get_blocks(file_path, layout, output_folder)
+            doc_number += 1"""
+    elif pipeline == "MinerU":
+        from pdf2data.mineru import MinerU
+        miner_pipeline: MinerU = MinerU(
+            input_folder=input_folder,
+            output_folder=output_folder,
+            extract_equations=False,
+            extract_text=False)
+        miner_pipeline.pdf_transform()
+    elif pipeline == "Docling":
+        from pdf2data.docling import Docling
+        docling_pipeline: Docling = Docling(
+            input_folder=input_folder,
+            output_folder=output_folder,
+            extract_equations=False,
+            extract_text=False)
+        docling_pipeline.pdf_transform()
+    elif pipeline in ["PaddlePPStructure", "PaddleVL"]:
+        from pdf2data.padle_pipeline import PaddlePPStructure
+        paddle_pipeline: PaddlePPStructure = PaddlePPStructure(
+            extractor_name=pipeline,
+            input_folder=input_folder,
+            output_folder=output_folder,
+            extract_equations=False,
+            extract_text=False)
+        paddle_pipeline.pdf_transform()
+    elif pipeline == "MinerUVL":
+        from pdf2data.mineru_vlm import MinerUVLM
+        mineru_vlm_pipeline: MinerUVLM = MinerUVLM(
+            input_folder=input_folder,
+            output_folder=output_folder,
+            extract_equations=False,
+            extract_text=False)
+        mineru_vlm_pipeline.pdf_transform()
+    
+    
+            
+
+def main():
+    block_extractor()
+
+"""@click.option(
+    "--table_model",
+    default="TableBank_faster_rcnn_R_101_FPN_3x",
+    help="model used to identify the tables",
+)"""
+
+if __name__ == "__main__":
+    main()

pdf2data/cli/block_finder.py

@@ -0,0 +1,50 @@
+import os
+from typing import Any, Dict, List
+import json
+from pdf2data.support import get_doc_list
+
+import click
+
+from pdf2data.keywords import BlockFinder
+
+@click.command()
+@click.argument("input_folder", type=str)
+@click.argument("output_folder", type=str)
+@click.argument("keywords_file", type=str)
+@click.option(
+    "--generic_file",
+    default=None,
+    help="file containing generic keywords",
+)
+@click.option(
+    "--find_tables",
+    default=True,
+    help="True to look for tables, False otherwise.",
+)
+@click.option(
+    "--find_figures",
+    default=False,
+    help="True to look for figures, False otherwise.",
+)
+def block_finder(input_folder: str, output_folder: str, keywords_file:str, generic_file: str, find_tables: bool, find_figures: bool):
+    if os.path.isdir(output_folder) is False:
+        os.mkdir(output_folder)
+    finder: BlockFinder = BlockFinder(keywords_file_path=keywords_file, generic_keywords_file_path=generic_file)
+    doc_list: List[str] = get_doc_list(input_folder, "")
+    final_results_dict: Dict[str, Any] = {}
+    results_path = f"{output_folder}/found_blocks.json"
+    for doc in doc_list:
+        print(doc)
+        blocks_path: str = f"{input_folder}/{doc}/{doc}_blocks.json"
+        blocks = finder.find(blocks_path, tables=find_tables, figures=find_figures)
+        final_results_dict[doc] = blocks
+    result_json = json.dumps(final_results_dict, indent=4)
+    with open(results_path, "w") as f:
+        f.write(result_json)
+
+def main():
+    block_finder()
+
+
+if __name__ == "__main__":
+    main() 

pdf2data/cli/evaluator.py

@@ -0,0 +1,45 @@
+import json
+import os
+
+import click
+
+from pdf2data.evaluator import Evaluator
+from pdf2data.support import get_doc_list
+
+@click.command()
+@click.argument("ref_folder", type=str)
+@click.argument("results_folder", type=str)
+@click.argument("eval_file_path", type=str)
+@click.argument("type", type=str)
+@click.option(
+    "--string_threshold",
+    default=0.6,
+    help="similarity threshold between strings",
+)
+@click.option(
+    "--box_threshold",
+    default=0.5,
+    help="similarity threshold between boxes",
+)
+
+def evaluator(ref_folder: str, results_folder: str, eval_file_path: str, type: str, string_threshold: float, box_threshold: float):
+    available_types: set = set(["metadata", "text", "blocks", "table_detection"])
+    if type not in available_types:
+        raise AttributeError(f"The specified type is not availabe, the available types are {available_types}")
+    evaluator = Evaluator(ref_folder=ref_folder, result_folder=results_folder, eval_file_path=eval_file_path, string_similarity=string_threshold, iou_threshold=box_threshold)
+    if type == "metadata":
+        evaluator.eval_metadata()
+    elif type == "text":
+        evaluator.eval_text()
+    elif type == "blocks":
+        evaluator.eval_blocks()
+    elif type == "table_detection":
+        evaluator.eval_table_detector()
+    
+
+def main():
+    evaluator()
+
+
+if __name__ == "__main__":
+    main()

pdf2data/cli/metadata_finder.py

@@ -0,0 +1,67 @@
+import json
+import os
+from typing import Any, Dict, List, Optional
+
+import click
+
+from pdf2data.metadata import Metadata
+from pdf2data.support import get_doc_list
+
+
+@click.command()
+@click.argument("path", type=str)
+@click.option(
+    "--output_folder",
+    default=None,
+    help="Output folder path",
+)
+@click.option(
+    "--file_type",
+    default="all",
+    help="file type to be analysed, options are all, pdf or cermxml",
+)
+def metadata_finder(path: str, output_folder: Optional[str], file_type: str) -> None:
+    """Create a json file containing the metadata for each file in path
+
+    Parameters
+    ----------
+    path : str
+        path to the file or folder to be processed
+    output_folder : Optional[str]
+        Output folder path to save the genereate json file containing the metdata
+    """
+    if os.path.isfile(path):
+        file_list: List[str] = [path]
+        input_folder: str = ""
+    elif file_type == "all":
+        file_list = get_doc_list(path, "pdf") + get_doc_list(path, "cermxml")
+        input_folder = path + "/"
+    elif file_type == "pdf":
+        file_list = get_doc_list(path, "pdf")
+        input_folder = path + "/"
+    elif file_type == "cermxml":
+        file_list = get_doc_list(path, "cermxml")
+        input_folder = path + "/"
+    if output_folder is None:
+        output_path = ""
+    else:
+        output_path = output_folder + "/"
+    for file in file_list:
+        file_path: str = f"{input_folder}{file}"
+        file_name: str = os.path.splitext(file)[0]
+        metadata: Metadata = Metadata(file_path=file_path)
+        metadata.update()
+        metadata_dict: Dict[str, Any] = metadata.__dict__
+        del metadata_dict["file_path"]
+        json_metadata = json.dumps(metadata_dict, indent=4)
+        with open(f"{output_path}{file_name}_metadata.json", "w") as j:
+            # convert the dictionary into a json variable
+            j.write(json_metadata)
+
+
+def main():
+    metadata_finder()
+
+
+if __name__ == "__main__":
+    main()

pdf2data/cli/pdf2data.py

@@ -0,0 +1,354 @@
+import os
+from typing import Any, Dict, List
+import json
+import time
+
+import click
+
+from pdf2data.references import References
+from pdf2data.metadata import Metadata
+
+@click.command()
+@click.argument("input_folder", type=str)
+@click.argument("output_folder", type=str)
+@click.option(
+    "--pipeline",
+    default="MinerU", #NotDefined, MinerU
+    help="Define the pipeline to be used",
+)
+@click.option(
+    "--layout_model",
+    default="PP-DocLayout-L",
+    help="model used to determine the overall layout",
+)
+@click.option(
+    "--layout_model_threshold",
+    default=0.7,
+    help="threshold of the layout model",
+)
+@click.option(
+    "--table_model",
+    default=None,
+    help="model used to identify the tables",
+)
+@click.option(
+    "--table_model_threshold",
+    default=0.5,
+    help="threshold of the table detection model",
+)
+@click.option(
+    "--text_extractor_type",
+    default="layoutparser",
+    help="type of the text extractor, available: ['layoutparser', 'cermine', 'minersix']",
+)
+@click.option(
+    "--reference_extractor_type",
+    default="anystyle",
+    help="type of the reference extractor, available: ['anystyle', 'cermine']",
+)
+@click.option(
+    "--metadata_extractor_type",
+    default="pdf2doi",
+    help="type of the metadata extractor, available: ['pdf2doi', 'cermine']",
+)
+@click.option(
+    "--extract_tables",
+    default=True,
+    help="True to extract tables, False otherwise",
+)
+@click.option(
+    "--extract_figures",
+    default=True,
+    help="True to extract figures, False otherwise",
+)
+@click.option(
+    "--correct_struct",
+    default=False,
+    help="True to to correct the table structure using the words position, False otherwise",
+)
+@click.option(
+    "--table_zoom",
+    default=1.5,
+    help="zoom of the image containing the table",
+)
+@click.option(
+    "--figure_zoom",
+    default=3,
+    help="zoom of the figures extracted",
+)
+@click.option(
+    "--x_table_corr",
+    default=0.01,
+    help="factor correct the table coordinates in the x axis",
+)
+@click.option(
+    "--y_table_corr",
+    default=0.01,
+    help="factor correct the table coordinates in the y axis",
+)
+@click.option(
+    "--iou_lines",
+    default=0.5,
+    help="iou value to supress collumns and rows",
+)
+@click.option(
+    "--iou_struct",
+    default=0.02,
+    help="minimum iou between table and row/collumn to consider the row/collumn as correct",
+)
+@click.option(
+    "--word_factor",
+    default=1.0,
+    help="factor used to determine the maximum distance to consider two different words as a single table entry",
+)
+@click.option(
+    "--word_iou",
+    default=0.00001,
+    help="iou value to consider that a word is inside a specific table entry",
+)
+@click.option(
+    "--struct_model_threshold",
+    default=0.3,
+    help="table structure detection model threshold",
+)
+@click.option(
+    "--reconstructor_type",
+    default="entry_by_entry",
+    help="type of reconstructor used, options: 'entry_by_entry' or 'word_by_word'",
+)
+@click.option(
+    "--brightness",
+    default=1.0,
+    help="brightness factor of the Table image",
+)
+@click.option(
+    "--contrast",
+    default=1.1,
+    help="contrast factor of the Table image",
+)
+@click.option(
+    "--device",
+    default="cpu",
+    help="device to run the mask models",
+)
+@click.option(
+    "--letter_ratio",
+    default=4.0,
+    help="minimum ratio between letter and ratio to consider a column as a row index or a row as a collumn header",
+)
+def pdf2data(input_folder: str,
+    output_folder: str,
+    pipeline: str,
+    layout_model: str,
+    table_model: str,
+    layout_model_threshold: float,
+    table_model_threshold: float,
+    text_extractor_type: float,
+    reference_extractor_type: float,
+    metadata_extractor_type: float,
+    extract_tables: bool,
+    extract_figures: bool,
+    correct_struct: bool,
+    table_zoom: float,
+    figure_zoom: float,
+    x_table_corr: float,
+    y_table_corr: float,
+    iou_lines: float,
+    iou_struct: float,
+    word_factor: float,
+    word_iou: float,
+    struct_model_threshold: float,
+    reconstructor_type: str,
+    brightness: float,
+    contrast: float,
+    device: str,
+    letter_ratio: float,
+) -> None:
+    start_time = time.time()
+    if pipeline == "NotDefined":
+        from pdf2data.pdf2data_pipeline import PDF2Data
+        pdf2data_pipeline: PDF2Data = PDF2Data(layout_model=layout_model, 
+                                               layout_model_threshold=layout_model_threshold, 
+                                               table_model=table_model, 
+                                               table_model_threshold=table_model_threshold, device=device, 
+                                               input_folder=input_folder, 
+                                               output_folder=output_folder,
+                                               extract_references=True)
+        pdf2data_pipeline.pdf_transform()
+        """file_list: List[str] = get_doc_list(input_folder, "pdf")
+        possible_types = set(["layoutparser", "cermine", "minersix"])
+        if text_extractor_type not in possible_types:
+            raise AttributeError(
+                f"{type} is not a available type, try one of the following: {possible_types}"
+            )
+        if text_extractor_type == "layoutparser":
+            text_extension: str = ".pdf"
+        else:
+            generator: TextFileGenerator = TextFileGenerator(
+                input_folder=input_folder, output_folder=output_folder
+            )
+            generator.model_post_init(None)
+            if text_extractor_type == "cermine":
+                generator.pdf_to_cermxml()
+                text_extension = ".cermxml"
+            elif text_extractor_type == "minersix":
+                generator.pdf_to_miner("txt")
+                text_extension = ".txt"
+        if reference_extractor_type == "anystyle":
+            reference_extension: str = ".pdf"
+        elif reference_extractor_type == "cermine":
+            if text_extractor_type != "cermine":
+                raise ValueError("Can only use reference_extractor as 'cermine' if the text_extrator is 'cermine'")
+            else:
+                reference_extension: str = ".cermxml"
+        if metadata_extractor_type == "pdf2doi":
+            metadata_extension: str = ".pdf"
+        elif metadata_extractor_type == "cermine":
+            if text_extractor_type != "cermine":
+                raise ValueError("Can only use metadata_extractor as 'cermine' if the text_extrator is 'cermine'")
+            else:
+                metadata_extension: str = ".cermxml"
+        extractor: BlockExtractor = BlockExtractor(
+            extract_tables=extract_tables,
+            extract_figures=extract_figures,
+            correct_struct=correct_struct,
+            table_zoom=table_zoom,
+            figure_zoom=figure_zoom,
+            x_table_corr=x_table_corr,
+            y_table_corr=y_table_corr,
+            iou_lines=iou_lines,
+            iou_struct=iou_struct,
+            word_factor=word_factor,
+            word_iou=word_iou,
+            struct_model_threshold=struct_model_threshold,
+            reconstructor_type=reconstructor_type,
+            brightness=brightness,
+            contrast=contrast,
+            letter_ratio=letter_ratio
+        )
+        extractor.model_post_init(None)
+        mask: LayoutParser = LayoutParser(
+            model=layout_model,
+            model_threshold=layout_model_threshold,
+            table_model=table_model,
+            table_model_threshold=table_model_threshold,
+            device_type=device
+        )
+        mask.model_post_init(None)
+        total_docs: int = len(file_list)
+        doc_number: int = 1
+        for file in file_list:
+            print(f'{doc_number}//{total_docs} processed')
+            print(file)
+            doc_number += 1
+            file_name = os.path.splitext(file)[0]
+            file_folder = output_folder + "/" + file_name
+
+            if os.path.isdir(file_folder) is False:
+                os.mkdir(file_folder)
+            file_path = input_folder + "/" + file
+            try:
+                pdf = PyPDF2.PdfReader(file_path)
+            except PyPDF2.errors.PdfReadError:
+                continue
+            layout: Dict[str, Any] = mask.get_layout(file_path)
+            text_doc = file_name + text_extension
+            text_path: str = f"{input_folder}/{text_doc}"
+            text_extractor: TextExtractor = TextExtractor(
+                    input_file=text_path, output_folder=file_folder
+                )
+            if text_extractor_type == "layoutparser":
+                text_extractor.extract_layoutparser(f"{file_name}_text", layout)
+            else:
+                if text_extractor_type == "cermine":
+                    text_extractor.extract_cermine(f"{file_name}_text")
+                elif text_extractor_type == "minersix":
+                    text_extractor.extract_txt(f"{file_name}_text")
+            reference_path = input_folder + "/" + file_name + reference_extension
+            references: References = References(
+                file_path=reference_path, output_folder=file_folder
+            )
+            references.generate_reference_file()
+            metadata_path = input_folder + "/" + file_name + metadata_extension
+            metadata: Metadata = Metadata(file_path=metadata_path)
+            metadata.update()
+            doi: str = metadata.doi
+            metadata_dict: Dict[str, Any] = metadata.__dict__
+            del metadata_dict["file_path"]
+            json_metadata = json.dumps(metadata_dict, indent=4)
+            with open(f"{file_folder}/{file_name}_metadata.json", "w") as j:
+                # convert the dictionary into a json variable
+                j.write(json_metadata)
+            extractor.get_blocks(file_path, layout, file_folder, doi=doi)
+            if text_extractor_type in set(["cermine", "minersix"]):
+                shutil.move(
+                    text_path, file_folder + "/" + file_name + text_extension
+                )
+            shutil.move(
+                    file_path, file_folder + "/" + file
+                )"""
+    elif pipeline == "MinerU":
+        from pdf2data.mineru import MinerU
+        miner_pipeline: MinerU = MinerU(
+            input_folder=input_folder,
+            output_folder=output_folder,
+            extract_references=True)
+        miner_pipeline.pdf_transform()
+    elif pipeline == "Docling":
+        from pdf2data.docling import Docling
+        docling_pipeline: Docling = Docling(
+            input_folder=input_folder,
+            output_folder=output_folder,
+            extract_references=True)
+        docling_pipeline.pdf_transform()
+    elif pipeline in ["PaddlePPStructure", "PaddleVL"]:
+        from pdf2data.padle_pipeline import PaddlePPStructure
+        paddle_pipeline: PaddlePPStructure = PaddlePPStructure(
+            extractor_name=pipeline,
+            input_folder=input_folder,
+            output_folder=output_folder,
+            extract_references=True)
+        paddle_pipeline.pdf_transform()
+    elif pipeline == "MinerUVL":
+        from pdf2data.mineru_vlm import MinerUVLM
+        mineru_vlm_pipeline: MinerUVLM = MinerUVLM(
+            input_folder=input_folder,
+            output_folder=output_folder,
+            extract_references=True)
+        mineru_vlm_pipeline.pdf_transform()
+    document_extraction_time = time.time() - start_time / 60
+    folders_list = os.listdir(output_folder)
+    number = 1
+    for folder in folders_list:
+        print(f"{number} // {len(folders_list)} processed")
+        print(folder)
+        pdf_file_path = input_folder + "/" + folder + ".pdf"
+        file_folder = output_folder + "/" + folder
+        txt_file_path = file_folder + "/" + folder + "_references.txt"
+        with open(file_folder + "/" + folder + "_content.json", "r") as f:
+            content_dict = json.load(f)
+        reference_parser =References(
+            file_path=txt_file_path, output_folder=file_folder
+        )
+        content_dict["references"] = reference_parser.generate_reference_list()
+        os.remove(txt_file_path)
+        metadata: Metadata = Metadata(file_path=pdf_file_path)
+        metadata.update()
+        metadata_dict: Dict[str, Any] = metadata.__dict__
+        content_dict["metadata"] = metadata_dict
+        new_content_dict = {"metadata": content_dict["metadata"]}
+        new_content_dict.update(content_dict)
+        content_json = json.dumps(new_content_dict, indent=4)
+        with open(file_folder + "/" + folder + "_content.json", "w") as f:
+            f.write(content_json)
+    full_time = time.time() - start_time / 60
+    extraction_metadata_dict = {"pipeline": pipeline, "document_extraction_time": document_extraction_time, "full_time": full_time}
+    with open(output_folder + "/extraction_metadata.json", "w") as f:
+        f.write(json.dumps(extraction_metadata_dict, indent=4))
+
+def main():
+    pdf2data()
+
+
+if __name__ == "__main__":
+    main()