Future-House / paper-qa

High accuracy RAG for answering questions from scientific documents with citations
Apache License 2.0
6.33k stars 600 forks source link

2 Column Layout #115

Closed whitead closed 1 year ago

whitead commented 1 year ago

The current parser has trouble with 2 column PDFs.

suny-sht commented 1 year ago

Hi, I have a hacky fix where an object detector is used to detect regions in the paper, and text regions are separately parsed. I use the layout parsers package to detect regions in images, and tessaract to OCR the images. For that, the pdf is first converted to a list of images. Some libraries needed are

pip install layoutparser
pip install pytesseract
pip install pdf2image

Here is a working example of such a pipeline that prints all text regions

import layoutparser as lp
import pytesseract
from pdf2image import convert_from_path

model = lp.Detectron2LayoutModel(
        config_path="lp://PubLayNet/mask_rcnn_X_101_32x8d_FPN_3x/config", 
        label_map={
            0: "Text",
            1: "Title",
            2: "List",
            3: "Table",
            4: "Figure",
        }, 
        extra_config=["MODEL.ROI_HEADS.SCORE_THRESH_TEST", 0.8],  
    )

def detect_regions(image):
        blocks = model.detect(image)._blocks

        detected_regions = {
            "Text": [],
            "Title": [],
            "List": [],
            "Table": [],
            "Figure": [],
        }

        for block in blocks:
            detected_regions[block.type].append(block.coordinates)

        return detected_regions

filepath = 'path_to_your_pdf'
images = convert_from_path(filepath)

for image in images:
    regions = detect_regions(image)

    for region_type, region_list in regions.items():
        if region_type == "Text":
            for region in region_list:
                # add very small padding to the region to make sure 
                # all the text is included, e.g. 0.5%
                padding = int(max(image.size) * 0.005)
                region = (
                    region[0] - padding,
                    region[1] - padding,
                    region[2] + padding,
                    region[3] + padding,
                )
                cropped_image = image.crop(region)
                text = pytesseract.image_to_string(cropped_image)
                print(text)

I found this to be quite robust for 2 column documents. An added benefit is it only parses text regions, skipping figures/tables. However, there is quite a bit of computational overhead

whitead commented 1 year ago

Fixed in v3