DiT_Extractor/base_utils.py

# Copyright (c) 2022, Lawrence Livermore National Security, LLC. 
# All rights reserved.
# See the top-level LICENSE and NOTICE files for details.
# LLNL-CODE-838964

# SPDX-License-Identifier: Apache-2.0-with-LLVM-exception

from pdfminer.pdfpage import PDFParser
from pdfminer.pdfpage import PDFDocument
from pdfminer.pdfpage import PDFPage
from pdfminer.layout import LTTextBoxHorizontal
from pdfminer.layout import LTTextLineHorizontal
from pdfminer.layout import LTChar
from pdfminer.layout import LAParams
from pdfminer.layout import LTRect
from pdfminer.layout import LTFigure

from pdfminer.converter import PDFPageAggregator
from pdfminer.pdfinterp import PDFResourceManager
from pdfminer.pdfinterp import PDFPageInterpreter
from pdfminer import pdfinterp

from collections.abc import Iterable 
from collections import Counter
from collections import OrderedDict

import os

# This is use for highlighting in PDFs 
from PyPDF2.generic import (
    DictionaryObject,
    NumberObject,
    FloatObject,
    NameObject,
    TextStringObject,
    ArrayObject
)

# Used to extract pages
from PyPDF2 import PdfFileReader, PdfFileWriter

def get_page_sizes(document):
    parser = PDFParser(open(document, 'rb'))
    doc = PDFDocument(parser)
    pageSizesList = []
    for page in PDFPage.create_pages(doc):
        # the media box that is the page size as list of 4 integers x0 y0 x1 y1
        pageSizesList.append(page.mediabox) # <- appending
    return pageSizesList
    
def get_page_count(document):
    # Is there a better way of getting the page count than doing this?   
    parser = PDFParser(document)
    tmpdoc = PDFDocument(parser)
    page_count = pdfinterp.resolve1(tmpdoc.catalog['Pages'])['Count']
    return page_count
    
def get_pdf_page_count(filename):
    with open(filename, 'rb') as document:    
        return get_page_count(document)
        
def get_pages(document, page_numbers = None):
    #Create resource manager
    rsrcmgr = PDFResourceManager()
    # Set parameters for analysis.
    laparams = LAParams()
    # Create a PDF page aggregator object.
    device = PDFPageAggregator(rsrcmgr, laparams=laparams)
    interpreter = PDFPageInterpreter(rsrcmgr, device)
    
    page_count = get_page_count(document)
    
    if page_numbers is None:
        page_numbers = range(page_count)
        
    for page, page_number in zip(PDFPage.get_pages(document, page_numbers), page_numbers):
        interpreter.process_page(page)
        # receive the LTPage object for the page.
        layout = device.get_result()
        #print("Yield page:", page_number)
        yield layout, page_number
                
def partial_overlaps(box, other):
    """
    Determine if the two bounding boxes overlap eachother.
    TODO: Really should just use a standard Python library for this.
    
    box -- 2 coordinate bounding box (x1,y1,x2,y2)
    other -- 2 coordinate bounding box (x1,y1,x2,y2)
    """
    # a1 x1  a2 x2      
    # <------------------>
    x_intersects = (other[0] < box[0] and other[2] > box[0]) or (
                    other[0] < box[2] and other[2] > box[2])
    y_intersects = (other[1] < box[1] and other[3] > box[1]) or (
                    other[1] < box[3] and other[3] > box[3]) 
 
    intersects = x_intersects or y_intersects
    # TODO: Simplify?
    return intersects and overlaps(box, other)
    #return intersects

def overlaps(box, other):
    """
    Determine if the two bounding boxes overlap eachother.
    TODO: Really should just use a standard Python library for this.
    
    box -- 2 coordinate bounding box (x1,y1,x2,y2)
    other -- 2 coordinate bounding box (x1,y1,x2,y2)
    """
    x_intersects = box[0] > other[2] or box[2] < other[0]
    y_intersects = box[1] > other[3] or box[3] < other[1] 
 
    intersects = not (x_intersects or y_intersects)
    return intersects

def union(src, other):
    """
    Expand src by union of other bbox
    
    src -- 2 coordinate bounding box (x1,y1,x2,y2)
    other -- 2 coordinate bounding box (x1,y1,x2,y2)
    
    returns union of src and other
    """
    xmin = min(src[0], other[0])
    ymin = min(src[1], other[1])
    xmax = max(src[2], other[2])
    ymax = max(src[3], other[3])
    
    return [xmin, ymin, xmax, ymax]



# See: https://gist.github.com/agentcooper/4c55133f5d95866acdee5017cd318558#file-pypdf2highlight-py
# x1, y1 starts in bottom left corner
def createHighlight(x1, y1, x2, y2, meta, color = [1, 0, 0]):
    newHighlight = DictionaryObject()

    newHighlight.update({
        NameObject("/F"): NumberObject(4),
        NameObject("/Type"): NameObject("/Annot"),
        NameObject("/Subtype"): NameObject("/Highlight"),

        NameObject("/T"): TextStringObject(meta["author"]),
        NameObject("/Contents"): TextStringObject(meta["contents"]),

        NameObject("/C"): ArrayObject([FloatObject(c) for c in color]),
        NameObject("/Rect"): ArrayObject([
            FloatObject(x1),
            FloatObject(y1),
            FloatObject(x2),
            FloatObject(y2)
        ]),
        NameObject("/QuadPoints"): ArrayObject([
            FloatObject(x1),
            FloatObject(y2),
            FloatObject(x2),
            FloatObject(y2),
            FloatObject(x1),
            FloatObject(y1),
            FloatObject(x2),
            FloatObject(y1)
        ]),
    })

    return newHighlight

def addHighlightToPage(highlight, page, output):
    highlight_ref = output._addObject(highlight);

    if "/Annots" in page:
        page[NameObject("/Annots")].append(highlight_ref)
    else:
        page[NameObject("/Annots")] = ArrayObject([highlight_ref])
    
def get_pdf_words(document, page_numbers=None):
    """
    Get all words from LTChar or LTTextLineHorizontal objects from the document.
    
    :param document: string path of the PDF file to process
    :returns: A map of page #'s containing lists of coordinates and PDFMiner
    objects. Ex.: {page_number: [[x1, y1, x2, y2, <LTTextLineHorizontal>],]}
    """
    pdf_doc = open(document, 'rb')
        
    bboxes = {}
    for layout, page in get_pages(pdf_doc, page_numbers):
        #print(element.get_text())
        bboxes[page] = []
        for element in layout:
            if not isinstance(element, Iterable):
                continue # not iterable
            for subElement in element:
                #print('Subelement type:', type(subElement))
                if isinstance(subElement, LTChar):                
                    if (subElement.get_text() == ' '):
                        pass  # TODO: Handle word deliminator                                    
                    # Print the character in this class
                    # print(subElement.get_text(), end='')      
                    item = list(subElement.bbox)
                    item.append(subElement)
                    bboxes[page].append(item)
                elif isinstance(subElement, LTTextLineHorizontal):
                    #print(subElement.bbox)
                    item = list(subElement.bbox)
                    item.append(subElement)
                    bboxes[page].append(item)
                else:
                    pass
    return bboxes

def get_paragraphs(words):
    paragraph_tolerance = 0.1
    max_height_diff = 1
    paragraphs = []
    
    for page, elements in words.items():
        # Find nominal font size
        # Round to int
        freq = Counter()
        for element in elements:
            height = int(element[3] - element[1])
            #print(height,end=' ')
            freq[height] += 1
            
        nominal_font = freq.most_common(1)[0][0]
        print("Nominal font is:", nominal_font)

        print("Page:", page)
        x_offset_prev_line = None
        prev_x_offset = None
        prev_y_offset = None
        paragraph_content = ""
        #print("Element count:", len(elements))
        first_line = False
        processed_first_line = False
        
        for element in elements:
            x_offset = element[0]
            y_offset = element[1]
            height = int(element[3] - element[1])
            text = element[4].get_text()
            
            if x_offset_prev_line != None:
                large_x_offset = (abs(x_offset_prev_line - x_offset) > paragraph_tolerance)
        
            # Font size mismatch?
            if abs(height - nominal_font) > max_height_diff:
                if len(paragraph_content) > 0:
                    print("Content append:", len(paragraph_content))
                    paragraphs.append(paragraph_content)
                    paragraph_content = ""
                print("Continue due to height != nominal_font")
                continue
                
            print("ELEMENT:", element[0:4], text[0:15])
            if prev_y_offset is not None and len(paragraph_content) > 0:
                if y_offset < prev_y_offset - height * 1.5:
                    print("Content append:", len(paragraph_content))
                    if len(paragraph_content) > 0:
                        paragraphs.append(paragraph_content)
                        paragraph_content = text
                    prev_y_offset = None
                    continue    
                
            prev_y_offset = y_offset
                
            prev_y_offset = y_offset
            #print("element:", element)
            if not isinstance(element[4], LTTextLineHorizontal):
                continue
            
            #print("Running text:", text)
            #print(f"x_offset_prev_line , x_offset]: {x_offset_prev_line, x_offset}")

            
            # Find first paragraph
            if x_offset_prev_line is None:
                #print("x_offset_prev is none")
                x_offset_prev_line = x_offset
                if not processed_first_line:
                    first_line = True
                    processed_first_line = True
                if height == nominal_font:
                    paragraph_content += text
                #print("Continue due to x_offset_prev_line is none")
                continue
                

            
            # Check case if first line was indented
            if x_offset_prev_line > x_offset and first_line:
                #print("x_offset < element[0]")
                first_line = False
                paragraph_content += text
                x_offset_prev_line = x_offset
                #print("Continue due to  x_offset_prev_line > x_offset and first_line")
                continue
                
            # is this indented?
            # and ignore small changes
            if x_offset_prev_line < x_offset and large_x_offset:
                #print(f"x_offset_prev_line > x_offset: {x_offset_prev_line, x_offset}")
                if height == nominal_font and len(paragraph_content) > 0:
                    paragraphs.append(paragraph_content)
                
                paragraph_content = text
                # Reset at next line read
                # What if next paragraph is also indented???
                x_offset_prev_line = None
                #print("Continue due to  x_offset_prev_line < x_offset and large_x_offset")
                continue                

            #print(element[0:4])
            if height == nominal_font:
                paragraph_content += text
            #print("End of loop")
                
            # TODO: Remove redundant space
        if paragraph_content != "":
            paragraphs.append(paragraph_content)
    
    # Find paragraph indexes
    c = 0
    indexes = []
    for p in paragraphs:
        c += len(p)
        indexes.append(c)
        
    return paragraphs, indexes
    
def get_pdf_elements(document, element_type, page_numbers=None):
    pdf_doc = open(document, 'rb')
        
    items = {}
    for layout, page in get_pages(pdf_doc, page_numbers):
        #print(element.get_text())
        items[page] = []
        for element in layout:
            if isinstance(element, element_type):
                item = list(element.bbox)
                if hasattr(element, 'non_stroking_color'):
                    item.append(element.non_stroking_color)
                items[page].append(item)
    print(items)
    return items
                
def get_large_colored_background_rectangles(document, page_numbers=None):
    # Only include rectangles that are at least 4" x 1" in size
    min_size = (288.0, 72.0)
    
    elements = get_pdf_elements(document, LTRect, page_numbers)
    rects_out = {}
    for page, rects in elements.items():
        print("Rects:", rects)
        for rect in rects:
            width = rect[2] - rect[0] 
            height = rect[3] - rect[1] 
            print("Dimensions:", width, height)
            if (width > min_size[0] and
                height > min_size[1]):
                if not page in rects_out:
                    rects_out[page] = []
                rects_out[page].append(rect)
    return rects_out

def extract_pages(document, output, page_numbers=None):
    pdf = PdfFileReader(document)
    
    pdf_writer = PdfFileWriter()
    for page in page_numbers:
        current_page = pdf.getPage(page)
        pdf_writer.addPage(current_page)
    
    with open(output, "wb") as out:
        pdf_writer.write(out)