First commit

.gitattributes

@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.traineddata filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,239 @@
+import streamlit as st
+
+import PIL
+import numpy as np
+import torch
+from collections import defaultdict
+
+import cv2
+from doctr.io import DocumentFile
+from doctr.models import ocr_predictor
+from doctr.utils.visualization import visualize_page
+
+import pytesseract
+from pytesseract import Output
+
+from bs4 import BeautifulSoup as bs
+
+import sys, json
+
+import postprocess
+
+
+ocr_predictor = ocr_predictor('db_resnet50', 'crnn_vgg16_bn', pretrained=True)
+structure_model = torch.hub.load('ultralytics/yolov5', 'custom', 'weights/structure_wts.pt', force_reload=True)
+imgsz = 640
+
+structure_class_names = [
+    'table', 'table column', 'table row', 'table column header',
+    'table projected row header', 'table spanning cell', 'no object'
+]
+structure_class_map = {k: v for v, k in enumerate(structure_class_names)}
+structure_class_thresholds = {
+    "table": 0.5,
+    "table column": 0.5,
+    "table row": 0.5,
+    "table column header": 0.25,
+    "table projected row header": 0.25,
+    "table spanning cell": 0.25,
+    "no object": 10
+}
+
+
+def table_structure(filename):
+    image = cv2.imread(filename)
+    pred = structure_model(image, size=imgsz)
+    pred = pred.xywhn[0]    
+    result = pred.cpu().numpy()
+    return result
+
+
+def ocr(filename):
+    doc = DocumentFile.from_images(filename)
+    result = ocr_predictor(doc).export()
+    result = result['pages'][0]
+    H, W = result['dimensions']
+    ocr_res = []
+    for block in result['blocks']:
+        for line in block['lines']:
+            for word in line['words']:
+                bbox = word['geometry']
+                word_info = {
+                    'bbox': [int(bbox[0][0] * W), int(bbox[0][1] * H), int(bbox[1][0] * W), int(bbox[1][1] * H)],
+                    'text': word['value']
+                }
+                ocr_res.append(word_info)
+    return ocr_res
+
+
+def convert_stucture(page_tokens, filename, structure_result):
+    image = cv2.imread(filename)
+    width = image.shape[1]
+    height = image.shape[0]
+    # print(width, height)
+    
+    bboxes = []
+    scores = []
+    labels = []
+    for i, result in enumerate(structure_result):
+        class_id = int(result[5])
+        score = float(result[4])
+        min_x = result[0]
+        min_y = result[1]
+        w = result[2]
+        h = result[3]
+        
+        x1 = int((min_x-w/2)*width)
+        y1 = int((min_y-h/2)*height)
+        x2 = int((min_x+w/2)*width)
+        y2 = int((min_y+h/2)*height)
+        # print(x1, y1, x2, y2)
+
+        bboxes.append([x1, y1, x2, y2])
+        scores.append(score)
+        labels.append(class_id)
+
+    table_objects = []
+    for bbox, score, label in zip(bboxes, scores, labels):
+        table_objects.append({'bbox': bbox, 'score': score, 'label': label})
+    # print('table_objects:', table_objects)
+        
+    table = {'objects': table_objects, 'page_num': 0}
+    
+    table_class_objects = [obj for obj in table_objects if obj['label'] == structure_class_map['table']]
+    if len(table_class_objects) > 1:
+        table_class_objects = sorted(table_class_objects, key=lambda x: x['score'], reverse=True)
+    try:
+        table_bbox = list(table_class_objects[0]['bbox'])
+    except:
+        table_bbox = (0,0,1000,1000)
+    # print('table_class_objects:', table_class_objects)
+    # print('table_bbox:', table_bbox)
+    
+    tokens_in_table = [token for token in page_tokens if postprocess.iob(token['bbox'], table_bbox) >= 0.5]
+    # print('tokens_in_table:', tokens_in_table)
+    
+    table_structures, cells, confidence_score = postprocess.objects_to_cells(table, table_objects, tokens_in_table, structure_class_names, structure_class_thresholds)
+    
+    return table_structures, cells, confidence_score
+
+
+def visualize_cells(filename, cells, ax):    
+    image = cv2.imread(filename)
+    for i, cell in enumerate(cells):
+        bbox = cell['bbox']
+        x1 = int(bbox[0])
+        y1 = int(bbox[1])
+        x2 = int(bbox[2])
+        y2 = int(bbox[3])
+        cv2.rectangle(image, (x1, y1), (x2, y2), color=(0, 255, 0))
+    ax.image(image)
+
+
+def pytess(cell_pil_img):
+    return ' '.join(pytesseract.image_to_data(cell_pil_img, output_type=Output.DICT, config='-c tessedit_char_blacklist=œ˜â€œï¬â™Ã©œ¢!|”?«“¥ --tessdata-dir tessdata --oem 3 --psm 6')['text']).strip()
+
+
+def resize(pil_img, size=1800):
+    length_x, width_y = pil_img.size
+    factor = max(1, size / length_x)
+    size = int(factor * length_x), int(factor * width_y)
+    pil_img = pil_img.resize(size, PIL.Image.ANTIALIAS)
+    return pil_img, factor
+
+
+def image_smoothening(img):
+    ret1, th1 = cv2.threshold(img, 180, 255, cv2.THRESH_BINARY)
+    ret2, th2 = cv2.threshold(th1, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+    blur = cv2.GaussianBlur(th2, (1, 1), 0)
+    ret3, th3 = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+    return th3
+
+
+def remove_noise_and_smooth(pil_img):
+    img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2GRAY)
+    filtered = cv2.adaptiveThreshold(img.astype(np.uint8), 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 41, 3)
+    kernel = np.ones((1, 1), np.uint8)
+    opening = cv2.morphologyEx(filtered, cv2.MORPH_OPEN, kernel)
+    closing = cv2.morphologyEx(opening, cv2.MORPH_CLOSE, kernel)
+    img = image_smoothening(img)
+    or_image = cv2.bitwise_or(img, closing)
+    pil_img = PIL.Image.fromarray(or_image)
+    return pil_img
+
+
+def extract_text_from_cells(filename, cells):
+    pil_img = PIL.Image.open(filename)
+    pil_img, factor = resize(pil_img)
+    #pil_img = remove_noise_and_smooth(pil_img)
+    #display(pil_img)
+    for cell in cells:
+        bbox = [x * factor for x in cell['bbox']]
+        cell_pil_img = pil_img.crop(bbox)
+        #cell_pil_img = remove_noise_and_smooth(cell_pil_img)
+        #cell_pil_img = tess_prep(cell_pil_img)        
+        cell['text'] = pytess(cell_pil_img)
+    return cells
+
+
+def cells_to_html(cells):
+    n_cols = max(cell['column_nums'][-1] for cell in cells) + 1
+    n_rows = max(cell['row_nums'][-1] for cell in cells) + 1
+    html_code = ''
+    for r in range(n_rows):
+        r_cells = [cell for cell in cells if cell['row_nums'][0] == r]
+        r_cells.sort(key=lambda x: x['column_nums'][0])
+        r_html = ''
+        for cell in r_cells:
+            rowspan = cell['row_nums'][-1] - cell['row_nums'][0] + 1
+            colspan = cell['column_nums'][-1] - cell['column_nums'][0] + 1
+            r_html += f'<td rowspan="{rowspan}" colspan="{colspan}">{cell["text"]}</td>'
+        html_code += f'<tr>{r_html}</tr>'
+    html_code = '''<html>
+                   <head>
+                   <meta charset="UTF-8">
+                   <style>
+                   table, th, td {
+                     border: 1px solid black;
+                     font-size: 10px;
+                   }
+                   </style>
+                   </head>
+                   <body>
+                   <table frame="hsides" rules="groups" width="100%%">
+                     %s
+                   </table>
+                   </body>
+                   </html>''' % html_code
+    soup = bs(html_code)
+    html_code = soup.prettify()
+    return html_code
+
+
+def main():
+
+    st.set_page_config(layout="wide")
+    st.title("Table Structure Recognition Demo")
+    st.write('\n')
+
+    cols = st.beta_columns((1, 1, 1))
+    cols[0].subheader("Input page")
+    cols[1].subheader("Structure output")
+    cols[2].subheader("HTML output")
+
+    st.sidebar.title("Image upload")
+    st.set_option('deprecation.showfileUploaderEncoding', False)
+    filename = st.sidebar.file_uploader("Upload files", type=['png', 'jpeg', 'jpg'])
+
+    cols[0].image(cv2.imread(filename))
+
+    ocr_res = ocr(filename)
+    structure_result = table_structure(filename)
+    table_structures, cells, confidence_score = convert_stucture(ocr_res, filename, structure_result)
+    visualize_cells(filename, cells, cols[1])
+
+    cells = extract_text_from_cells(filename, cells)
+    html_code = cells_to_html(cells)
+
+    cols[2].html(html_code)
+    

packages.txt

@@ -0,0 +1,6 @@
+ffmpeg 
+libsm6 
+libxext6 
+libgl1
+tesseract-ocr-eng
+python3-opencv

postprocess.py

@@ -0,0 +1,895 @@
+"""
+Copyright (C) 2021 Microsoft Corporation
+"""
+from collections import defaultdict
+
+from fitz import Rect
+
+
+def apply_threshold(objects, threshold):
+    """
+    Filter out objects below a certain score.
+    """
+    return [obj for obj in objects if obj['score'] >= threshold]
+
+
+def apply_class_thresholds(bboxes, labels, scores, class_names, class_thresholds):
+    """
+    Filter out bounding boxes whose confidence is below the confidence threshold for
+    its associated class label. 
+    """
+    # Apply class-specific thresholds
+    indices_above_threshold = [idx for idx, (score, label) in enumerate(zip(scores, labels))
+                                    if score >= class_thresholds[
+                                        class_names[label]
+                                    ]
+                                ]
+    bboxes = [bboxes[idx] for idx in indices_above_threshold]
+    scores = [scores[idx] for idx in indices_above_threshold]
+    labels = [labels[idx] for idx in indices_above_threshold]
+
+    return bboxes, scores, labels
+
+
+def iou(bbox1, bbox2):
+    """
+    Compute the intersection-over-union of two bounding boxes.
+    """
+    intersection = Rect(bbox1).intersect(bbox2)
+    union = Rect(bbox1).include_rect(bbox2)
+    
+    union_area = union.get_area()  # getArea()
+    if union_area > 0:
+        return intersection.get_area() / union.get_area()  # .getArea()
+    
+    return 0
+
+
+def iob(bbox1, bbox2):
+    """
+    Compute the intersection area over box area, for bbox1.
+    """
+    intersection = Rect(bbox1).intersect(bbox2)
+    
+    bbox1_area = Rect(bbox1).get_area()  # .getArea()
+    if bbox1_area > 0:
+        return intersection.get_area() / bbox1_area  # getArea()
+    
+    return 0
+
+
+def objects_to_cells(table, objects_in_table, tokens_in_table, class_map, class_thresholds):
+    """
+    Process the bounding boxes produced by the table structure recognition model
+    and the token/word/span bounding boxes into table cells.
+
+    Also return a confidence score based on how well the text was able to be
+    uniquely slotted into the cells detected by the table model.
+    """
+
+    table_structures = objects_to_table_structures(table, objects_in_table, tokens_in_table, class_map,
+                                                   class_thresholds)
+
+    # Check for a valid table
+    if len(table_structures['columns']) < 1 or len(table_structures['rows']) < 1:
+        cells = []#None
+        confidence_score = 0
+    else:
+        cells, confidence_score = table_structure_to_cells(table_structures, tokens_in_table, table['bbox'])
+
+    return table_structures, cells, confidence_score
+
+
+def objects_to_table_structures(table_object, objects_in_table, tokens_in_table, class_names, class_thresholds):
+    """
+    Process the bounding boxes produced by the table structure recognition model into
+    a *consistent* set of table structures (rows, columns, supercells, headers).
+    This entails resolving conflicts/overlaps, and ensuring the boxes meet certain alignment
+    conditions (for example: rows should all have the same width, etc.).
+    """
+
+    page_num = table_object['page_num']
+
+    table_structures = {}
+
+    columns = [obj for obj in objects_in_table if class_names[obj['label']] == 'table column']
+    rows = [obj for obj in objects_in_table if class_names[obj['label']] == 'table row']
+    headers = [obj for obj in objects_in_table if class_names[obj['label']] == 'table column header']
+    supercells = [obj for obj in objects_in_table if class_names[obj['label']] == 'table spanning cell']
+    for obj in supercells:
+        obj['subheader'] = False
+    subheaders = [obj for obj in objects_in_table if class_names[obj['label']] == 'table projected row header']
+    for obj in subheaders:
+        obj['subheader'] = True
+    supercells += subheaders
+    for obj in rows:
+        obj['header'] = False
+        for header_obj in headers:
+            if iob(obj['bbox'], header_obj['bbox']) >= 0.5:
+                obj['header'] = True
+
+    for row in rows:
+        row['page'] = page_num
+
+    for column in columns:
+        column['page'] = page_num
+
+    #Refine table structures
+    rows = refine_rows(rows, tokens_in_table, class_thresholds['table row'])
+    columns = refine_columns(columns, tokens_in_table, class_thresholds['table column'])
+
+    # Shrink table bbox to just the total height of the rows
+    # and the total width of the columns
+    row_rect = Rect()
+    for obj in rows:
+        row_rect.include_rect(obj['bbox'])
+    column_rect = Rect() 
+    for obj in columns:
+        column_rect.include_rect(obj['bbox'])
+    table_object['row_column_bbox'] = [column_rect[0], row_rect[1], column_rect[2], row_rect[3]]
+    table_object['bbox'] = table_object['row_column_bbox']
+
+    # Process the rows and columns into a complete segmented table
+    columns = align_columns(columns, table_object['row_column_bbox'])
+    rows = align_rows(rows, table_object['row_column_bbox'])
+
+    table_structures['rows'] = rows
+    table_structures['columns'] = columns
+    table_structures['headers'] = headers
+    table_structures['supercells'] = supercells
+
+    if len(rows) > 0 and len(columns) > 1:
+        table_structures = refine_table_structures(table_object['bbox'], table_structures, tokens_in_table, class_thresholds)
+
+    return table_structures
+
+
+def refine_rows(rows, page_spans, score_threshold):
+    """
+    Apply operations to the detected rows, such as
+    thresholding, NMS, and alignment.
+    """
+    
+    #MODIFY
+    rows = [obj for obj in rows if obj['score'] >= score_threshold or obj['header']]
+    ###
+    
+    rows = nms_by_containment(rows, page_spans, overlap_threshold=0.5)
+    # remove_objects_without_content(page_spans, rows)  # TODO
+    if len(rows) > 1:
+        rows = sort_objects_top_to_bottom(rows)
+
+    return rows
+
+
+def refine_columns(columns, page_spans, score_threshold):
+    """
+    Apply operations to the detected columns, such as
+    thresholding, NMS, and alignment.
+    """
+
+    #MODIFY
+    columns = [obj for obj in columns if obj['score'] >= score_threshold]
+    ###
+    
+    columns = nms_by_containment(columns, page_spans, overlap_threshold=0.5)
+    # remove_objects_without_content(page_spans, columns)  # TODO
+    if len(columns) > 1:
+        columns = sort_objects_left_to_right(columns)
+
+    return columns
+
+
+def nms_by_containment(container_objects, package_objects, overlap_threshold=0.5):
+    """
+    Non-maxima suppression (NMS) of objects based on shared containment of other objects.
+    """
+    container_objects = sort_objects_by_score(container_objects)
+    num_objects = len(container_objects)
+    suppression = [False for obj in container_objects]
+
+    packages_by_container, _, _ = slot_into_containers(container_objects, package_objects, overlap_threshold=overlap_threshold,
+                                                 unique_assignment=True, forced_assignment=False)
+
+    for object2_num in range(1, num_objects):
+        object2_packages = set(packages_by_container[object2_num])
+        if len(object2_packages) == 0:
+            suppression[object2_num] = True
+        for object1_num in range(object2_num):
+            if not suppression[object1_num]:
+                object1_packages = set(packages_by_container[object1_num])
+                if len(object2_packages.intersection(object1_packages)) > 0:
+                    suppression[object2_num] = True
+
+    final_objects = [obj for idx, obj in enumerate(container_objects) if not suppression[idx]]
+    return final_objects
+
+
+def slot_into_containers(container_objects, package_objects, overlap_threshold=0.5,
+                         unique_assignment=True, forced_assignment=False):
+    """
+    Slot a collection of objects into the container they occupy most (the container which holds the largest fraction of the object).
+    """
+    best_match_scores = []
+
+    container_assignments = [[] for container in container_objects]
+    package_assignments = [[] for package in package_objects]
+
+    if len(container_objects) == 0 or len(package_objects) == 0:
+        return container_assignments, package_assignments, best_match_scores
+
+    match_scores = defaultdict(dict)
+    for package_num, package in enumerate(package_objects):
+        match_scores = []
+        package_rect = Rect(package['bbox'])
+        package_area = package_rect.get_area()  # getArea()        
+        for container_num, container in enumerate(container_objects):
+            container_rect = Rect(container['bbox'])
+            intersect_area = container_rect.intersect(package['bbox']).get_area()  # getArea()
+            overlap_fraction = intersect_area / package_area
+            match_scores.append({'container': container, 'container_num': container_num, 'score': overlap_fraction})
+
+        sorted_match_scores = sort_objects_by_score(match_scores)
+
+        best_match_score = sorted_match_scores[0]
+        best_match_scores.append(best_match_score['score'])
+        if forced_assignment or best_match_score['score'] >= overlap_threshold:
+            container_assignments[best_match_score['container_num']].append(package_num)
+            package_assignments[package_num].append(best_match_score['container_num'])
+
+        if not unique_assignment: # slot package into all eligible slots
+            for match_score in sorted_match_scores[1:]:
+                if match_score['score'] >= overlap_threshold:
+                    container_assignments[match_score['container_num']].append(package_num)
+                    package_assignments[package_num].append(match_score['container_num'])
+                else:
+                    break
+            
+    return container_assignments, package_assignments, best_match_scores
+
+
+def sort_objects_by_score(objects, reverse=True):
+    """
+    Put any set of objects in order from high score to low score.
+    """
+    if reverse:
+        sign = -1
+    else:
+        sign = 1
+    return sorted(objects, key=lambda k: sign*k['score'])
+
+
+def remove_objects_without_content(page_spans, objects):
+    """
+    Remove any objects (these can be rows, columns, supercells, etc.) that don't
+    have any text associated with them.
+    """
+    for obj in objects[:]:
+        object_text, _ = extract_text_inside_bbox(page_spans, obj['bbox'])
+        if len(object_text.strip()) == 0:
+            objects.remove(obj)
+            
+            
+def extract_text_inside_bbox(spans, bbox):
+    """
+    Extract the text inside a bounding box.
+    """
+    bbox_spans = get_bbox_span_subset(spans, bbox)
+    bbox_text = extract_text_from_spans(bbox_spans, remove_integer_superscripts=True)
+
+    return bbox_text, bbox_spans
+
+
+def get_bbox_span_subset(spans, bbox, threshold=0.5):
+    """
+    Reduce the set of spans to those that fall within a bounding box.
+
+    threshold: the fraction of the span that must overlap with the bbox.
+    """
+    span_subset = []
+    for span in spans:
+        if overlaps(span['bbox'], bbox, threshold):
+            span_subset.append(span)
+    return span_subset
+
+
+def overlaps(bbox1, bbox2, threshold=0.5):
+    """
+    Test if more than "threshold" fraction of bbox1 overlaps with bbox2.
+    """
+    rect1 = Rect(list(bbox1))
+    area1 = rect1.get_area()  # .getArea()
+    if area1 == 0:
+        return False
+    return rect1.intersect(list(bbox2)).get_area()/area1 >= threshold  # getArea()
+
+
+def extract_text_from_spans(spans, join_with_space=True, remove_integer_superscripts=True):
+    """
+    Convert a collection of page tokens/words/spans into a single text string.
+    """
+
+    if join_with_space:
+        join_char = " "
+    else:
+        join_char = ""
+    spans_copy = spans[:]
+    
+    if remove_integer_superscripts:
+        for span in spans:
+            flags = span['flags']
+            if flags & 2**0: # superscript flag
+                if is_int(span['text']):
+                    spans_copy.remove(span)
+                else:
+                    span['superscript'] = True
+
+    if len(spans_copy) == 0:
+        return ""
+    
+    spans_copy.sort(key=lambda span: span['span_num'])
+    spans_copy.sort(key=lambda span: span['line_num'])
+    spans_copy.sort(key=lambda span: span['block_num'])
+    
+    # Force the span at the end of every line within a block to have exactly one space
+    # unless the line ends with a space or ends with a non-space followed by a hyphen
+    line_texts = []
+    line_span_texts = [spans_copy[0]['text']]
+    for span1, span2 in zip(spans_copy[:-1], spans_copy[1:]):
+        if not span1['block_num'] == span2['block_num'] or not span1['line_num'] == span2['line_num']:
+            line_text = join_char.join(line_span_texts).strip()
+            if (len(line_text) > 0
+                    and not line_text[-1] == ' '
+                    and not (len(line_text) > 1 and line_text[-1] == "-" and not line_text[-2] == ' ')):
+                if not join_with_space:
+                    line_text += ' '
+            line_texts.append(line_text)
+            line_span_texts = [span2['text']]
+        else:
+            line_span_texts.append(span2['text'])
+    line_text = join_char.join(line_span_texts)
+    line_texts.append(line_text)
+            
+    return join_char.join(line_texts).strip()
+
+
+def sort_objects_left_to_right(objs):
+    """
+    Put the objects in order from left to right.
+    """
+    return sorted(objs, key=lambda k: k['bbox'][0] + k['bbox'][2])
+
+
+def sort_objects_top_to_bottom(objs):
+    """
+    Put the objects in order from top to bottom.
+    """
+    return sorted(objs, key=lambda k: k['bbox'][1] + k['bbox'][3])
+
+
+def align_columns(columns, bbox):
+    """
+    For every column, align the top and bottom boundaries to the final
+    table bounding box.
+    """
+    try:
+        for column in columns:
+            column['bbox'][1] = bbox[1]
+            column['bbox'][3] = bbox[3]
+    except Exception as err:
+        print("Could not align columns: {}".format(err))
+        pass
+
+    return columns
+
+
+def align_rows(rows, bbox):
+    """
+    For every row, align the left and right boundaries to the final
+    table bounding box.
+    """
+    try:
+        for row in rows:
+            row['bbox'][0] = bbox[0]
+            row['bbox'][2] = bbox[2]
+    except Exception as err:
+        print("Could not align rows: {}".format(err))
+        pass
+
+    return rows
+
+
+def refine_table_structures(table_bbox, table_structures, page_spans, class_thresholds):
+    """
+    Apply operations to the detected table structure objects such as
+    thresholding, NMS, and alignment.
+    """
+    rows = table_structures["rows"]
+    columns = table_structures['columns']
+
+    #columns = fill_column_gaps(columns, table_bbox)
+    #rows = fill_row_gaps(rows, table_bbox)
+
+    # Process the headers
+    headers = table_structures['headers']
+    headers = apply_threshold(headers, class_thresholds["table column header"])
+    headers = nms(headers)
+    headers = align_headers(headers, rows)
+
+    # Process supercells
+    supercells = [elem for elem in table_structures['supercells'] if not elem['subheader']]
+    subheaders = [elem for elem in table_structures['supercells'] if elem['subheader']]
+    supercells = apply_threshold(supercells, class_thresholds["table spanning cell"])
+    subheaders = apply_threshold(subheaders, class_thresholds["table projected row header"])
+    supercells += subheaders
+    # Align before NMS for supercells because alignment brings them into agreement
+    # with rows and columns first; if supercells still overlap after this operation,
+    # the threshold for NMS can basically be lowered to just above 0
+    supercells = align_supercells(supercells, rows, columns)
+    supercells = nms_supercells(supercells)
+
+    header_supercell_tree(supercells)
+
+    table_structures['columns'] = columns
+    table_structures['rows'] = rows
+    table_structures['supercells'] = supercells
+    table_structures['headers'] = headers
+
+    return table_structures
+
+
+def nms(objects, match_criteria="object2_overlap", match_threshold=0.05, keep_metric="score", keep_higher=True):
+    """
+    A customizable version of non-maxima suppression (NMS).
+    
+    Default behavior: If a lower-confidence object overlaps more than 5% of its area
+    with a higher-confidence object, remove the lower-confidence object.
+
+    objects: set of dicts; each object dict must have a 'bbox' and a 'score' field
+    match_criteria: how to measure how much two objects "overlap"
+    match_threshold: the cutoff for determining that overlap requires suppression of one object
+    keep_metric: which metric to use to determine the object to keep
+    keep_higher: if True, keep the object with the higher metric; otherwise, keep the lower
+    """
+    if len(objects) == 0:
+        return []
+
+    if keep_metric=="score":
+        objects = sort_objects_by_score(objects, reverse=keep_higher)
+    elif keep_metric=="area":
+        objects = sort_objects_by_area(objects, reverse=keep_higher)
+
+    num_objects = len(objects)
+    suppression = [False for obj in objects]
+
+    for object2_num in range(1, num_objects):
+        object2_rect = Rect(objects[object2_num]['bbox'])
+        object2_area = object2_rect.get_area()  # .getArea()
+        for object1_num in range(object2_num):
+            if not suppression[object1_num]:
+                object1_rect = Rect(objects[object1_num]['bbox'])
+                object1_area = object1_rect.get_area()  # .getArea()
+                intersect_area = object1_rect.intersect(object2_rect).get_area()  # .getArea()
+                try:
+                    if match_criteria=="object1_overlap":
+                        metric = intersect_area / object1_area
+                    elif match_criteria=="object2_overlap":
+                        metric = intersect_area / object2_area
+                    elif match_criteria=="iou":
+                        metric = intersect_area / (object1_area + object2_area - intersect_area)
+                    if metric >= match_threshold:
+                        suppression[object2_num] = True
+                        break
+                except Exception:
+                    # Intended to recover from divide-by-zero
+                    pass
+
+    return [obj for idx, obj in enumerate(objects) if not suppression[idx]]
+
+
+def align_headers(headers, rows):
+    """
+    Adjust the header boundary to be the convex hull of the rows it intersects
+    at least 50% of the height of.
+
+    For now, we are not supporting tables with multiple headers, so we need to
+    eliminate anything besides the top-most header.
+    """
+    
+    aligned_headers = []
+
+    for row in rows:
+        row['header'] = False
+
+    header_row_nums = []
+    for header in headers:
+        for row_num, row in enumerate(rows):
+            row_height = row['bbox'][3] - row['bbox'][1]
+            min_row_overlap = max(row['bbox'][1], header['bbox'][1])
+            max_row_overlap = min(row['bbox'][3], header['bbox'][3])
+            overlap_height = max_row_overlap - min_row_overlap
+            if overlap_height / row_height >= 0.5:
+                header_row_nums.append(row_num)
+
+    if len(header_row_nums) == 0:
+        return aligned_headers
+
+    header_rect = Rect()
+    if header_row_nums[0] > 0:
+        header_row_nums = list(range(header_row_nums[0]+1)) + header_row_nums
+
+    last_row_num = -1
+    for row_num in header_row_nums:
+        if row_num == last_row_num + 1:
+            row = rows[row_num]
+            row['header'] = True
+            header_rect = header_rect.include_rect(row['bbox'])
+            last_row_num = row_num
+        else:
+            # Break as soon as a non-header row is encountered.
+            # This ignores any subsequent rows in the table labeled as a header.
+            # Having more than 1 header is not supported currently.
+            break
+
+    header = {'bbox': list(header_rect)}
+    aligned_headers.append(header)
+
+    return aligned_headers
+
+
+def align_supercells(supercells, rows, columns):
+    """
+    For each supercell, align it to the rows it intersects 50% of the height of,
+    and the columns it intersects 50% of the width of.
+    Eliminate supercells for which there are no rows and columns it intersects 50% with.
+    """
+    aligned_supercells = []
+
+    for supercell in supercells:
+        supercell['header'] = False
+        row_bbox_rect = None
+        col_bbox_rect = None
+        intersecting_header_rows = set()
+        intersecting_data_rows = set()
+        for row_num, row in enumerate(rows):
+            row_height = row['bbox'][3] - row['bbox'][1]
+            supercell_height = supercell['bbox'][3] - supercell['bbox'][1]
+            min_row_overlap = max(row['bbox'][1], supercell['bbox'][1])
+            max_row_overlap = min(row['bbox'][3], supercell['bbox'][3])
+            overlap_height = max_row_overlap - min_row_overlap
+            if 'span' in supercell:
+                overlap_fraction = max(overlap_height/row_height,
+                                       overlap_height/supercell_height)
+            else:
+                overlap_fraction = overlap_height / row_height
+            if overlap_fraction >= 0.5:
+                if 'header' in row and row['header']:
+                    intersecting_header_rows.add(row_num)
+                else:
+                    intersecting_data_rows.add(row_num)
+
+        # Supercell cannot span across the header boundary; eliminate whichever
+        # group of rows is the smallest
+        supercell['header'] = False
+        if len(intersecting_data_rows) > 0 and len(intersecting_header_rows) > 0:
+            if len(intersecting_data_rows) > len(intersecting_header_rows):
+                intersecting_header_rows = set()
+            else:
+                intersecting_data_rows = set()
+        if len(intersecting_header_rows) > 0:
+            supercell['header'] = True
+        elif 'span' in supercell:
+            continue # Require span supercell to be in the header
+        intersecting_rows = intersecting_data_rows.union(intersecting_header_rows)
+        # Determine vertical span of aligned supercell
+        for row_num in intersecting_rows:
+            if row_bbox_rect is None:
+                row_bbox_rect = Rect(rows[row_num]['bbox'])
+            else:
+                row_bbox_rect = row_bbox_rect.include_rect(rows[row_num]['bbox'])
+        if row_bbox_rect is None:
+            continue
+
+        intersecting_cols = []
+        for col_num, col in enumerate(columns):
+            col_width = col['bbox'][2] - col['bbox'][0]
+            supercell_width = supercell['bbox'][2] - supercell['bbox'][0]
+            min_col_overlap = max(col['bbox'][0], supercell['bbox'][0])
+            max_col_overlap = min(col['bbox'][2], supercell['bbox'][2])
+            overlap_width = max_col_overlap - min_col_overlap
+            if 'span' in supercell:
+                overlap_fraction = max(overlap_width/col_width,
+                                       overlap_width/supercell_width)
+                # Multiply by 2 effectively lowers the threshold to 0.25
+                if supercell['header']:
+                    overlap_fraction = overlap_fraction * 2
+            else:
+                overlap_fraction = overlap_width / col_width
+            if overlap_fraction >= 0.5:
+                intersecting_cols.append(col_num)
+                if col_bbox_rect is None:
+                    col_bbox_rect = Rect(col['bbox'])
+                else:
+                    col_bbox_rect = col_bbox_rect.include_rect(col['bbox'])
+        if col_bbox_rect is None:
+            continue
+
+        supercell_bbox = list(row_bbox_rect.intersect(col_bbox_rect))
+        supercell['bbox'] = supercell_bbox
+
+        # Only a true supercell if it joins across multiple rows or columns
+        if (len(intersecting_rows) > 0 and len(intersecting_cols) > 0
+                and (len(intersecting_rows) > 1 or len(intersecting_cols) > 1)):
+            supercell['row_numbers'] = list(intersecting_rows)
+            supercell['column_numbers'] = intersecting_cols
+            aligned_supercells.append(supercell)
+
+            # A span supercell in the header means there must be supercells above it in the header
+            if 'span' in supercell and supercell['header'] and len(supercell['column_numbers']) > 1:
+                for row_num in range(0, min(supercell['row_numbers'])):
+                    new_supercell = {'row_numbers': [row_num], 'column_numbers': supercell['column_numbers'],
+                                     'score': supercell['score'], 'propagated': True}
+                    new_supercell_columns = [columns[idx] for idx in supercell['column_numbers']]
+                    new_supercell_rows = [rows[idx] for idx in supercell['row_numbers']]
+                    bbox = [min([column['bbox'][0] for column in new_supercell_columns]),
+                            min([row['bbox'][1] for row in new_supercell_rows]),
+                            max([column['bbox'][2] for column in new_supercell_columns]),
+                            max([row['bbox'][3] for row in new_supercell_rows])]
+                    new_supercell['bbox'] = bbox
+                    aligned_supercells.append(new_supercell)
+
+    return aligned_supercells
+
+
+def nms_supercells(supercells):
+    """
+    A NMS scheme for supercells that first attempts to shrink supercells to
+    resolve overlap.
+    If two supercells overlap the same (sub)cell, shrink the lower confidence
+    supercell to resolve the overlap. If shrunk supercell is empty, remove it.
+    """
+
+    supercells = sort_objects_by_score(supercells)
+    num_supercells = len(supercells)
+    suppression = [False for supercell in supercells]
+
+    for supercell2_num in range(1, num_supercells):
+        supercell2 = supercells[supercell2_num]
+        for supercell1_num in range(supercell2_num):
+            supercell1 = supercells[supercell1_num]
+            remove_supercell_overlap(supercell1, supercell2)
+        if ((len(supercell2['row_numbers']) < 2 and len(supercell2['column_numbers']) < 2)
+                or len(supercell2['row_numbers']) == 0 or len(supercell2['column_numbers']) == 0):
+            suppression[supercell2_num] = True
+
+    return [obj for idx, obj in enumerate(supercells) if not suppression[idx]]
+
+
+def header_supercell_tree(supercells):
+    """
+    Make sure no supercell in the header is below more than one supercell in any row above it.
+    The cells in the header form a tree, but a supercell with more than one supercell in a row
+    above it means that some cell has more than one parent, which is not allowed. Eliminate
+    any supercell that would cause this to be violated.
+    """
+    header_supercells = [supercell for supercell in supercells if 'header' in supercell and supercell['header']]
+    header_supercells = sort_objects_by_score(header_supercells)
+    
+    for header_supercell in header_supercells[:]:
+        ancestors_by_row = defaultdict(int)
+        min_row = min(header_supercell['row_numbers'])
+        for header_supercell2 in header_supercells:
+            max_row2 = max(header_supercell2['row_numbers'])
+            if max_row2 < min_row:
+                if (set(header_supercell['column_numbers']).issubset(
+                    set(header_supercell2['column_numbers']))):
+                    for row2 in header_supercell2['row_numbers']:
+                        ancestors_by_row[row2] += 1
+        for row in range(0, min_row):
+            if not ancestors_by_row[row] == 1:
+                supercells.remove(header_supercell)
+                break
+                
+                
+def table_structure_to_cells(table_structures, table_spans, table_bbox):
+    """
+    Assuming the row, column, supercell, and header bounding boxes have
+    been refined into a set of consistent table structures, process these
+    table structures into table cells. This is a universal representation
+    format for the table, which can later be exported to Pandas or CSV formats.
+    Classify the cells as header/access cells or data cells
+    based on if they intersect with the header bounding box.
+    """
+    columns = table_structures['columns']
+    rows = table_structures['rows']
+    supercells = table_structures['supercells']
+    cells = []
+    subcells = []
+
+    # Identify complete cells and subcells
+    for column_num, column in enumerate(columns):
+        for row_num, row in enumerate(rows):
+            column_rect = Rect(list(column['bbox']))
+            row_rect = Rect(list(row['bbox']))
+            cell_rect = row_rect.intersect(column_rect)
+            header = 'header' in row and row['header']
+            cell = {'bbox': list(cell_rect), 'column_nums': [column_num], 'row_nums': [row_num],
+                    'header': header}
+
+            cell['subcell'] = False
+            for supercell in supercells:
+                supercell_rect = Rect(list(supercell['bbox']))
+                if (supercell_rect.intersect(cell_rect).get_area()  # .getArea()
+                        / cell_rect.get_area()) > 0.5:  # getArea()
+                    cell['subcell'] = True
+                    break
+
+            if cell['subcell']:
+                subcells.append(cell)
+            else:
+                #cell_text = extract_text_inside_bbox(table_spans, cell['bbox'])
+                #cell['cell_text'] = cell_text
+                cell['subheader'] = False
+                cells.append(cell)
+
+    for supercell in supercells:
+        supercell_rect = Rect(list(supercell['bbox']))
+        cell_columns = set()
+        cell_rows = set()
+        cell_rect = None
+        header = True
+        for subcell in subcells:
+            subcell_rect = Rect(list(subcell['bbox']))
+            subcell_rect_area = subcell_rect.get_area()  # .getArea()
+            if (subcell_rect.intersect(supercell_rect).get_area()  # .getArea()
+                    / subcell_rect_area) > 0.5:
+                if cell_rect is None:
+                    cell_rect = Rect(list(subcell['bbox']))
+                else:
+                    cell_rect.include_rect(Rect(list(subcell['bbox'])))
+                cell_rows = cell_rows.union(set(subcell['row_nums']))
+                cell_columns = cell_columns.union(set(subcell['column_nums']))
+                # By convention here, all subcells must be classified
+                # as header cells for a supercell to be classified as a header cell;
+                # otherwise, this could lead to a non-rectangular header region
+                header = header and 'header' in subcell and subcell['header']
+        if len(cell_rows) > 0 and len(cell_columns) > 0:
+            cell = {'bbox': list(cell_rect), 'column_nums': list(cell_columns), 'row_nums': list(cell_rows),
+                    'header': header, 'subheader': supercell['subheader']}
+            cells.append(cell)
+
+    # Compute a confidence score based on how well the page tokens
+    # slot into the cells reported by the model
+    _, _, cell_match_scores = slot_into_containers(cells, table_spans)
+    try:
+        mean_match_score = sum(cell_match_scores) / len(cell_match_scores)
+        min_match_score = min(cell_match_scores)
+        confidence_score = (mean_match_score + min_match_score)/2
+    except:
+        confidence_score = 0
+
+    # Dilate rows and columns before final extraction
+    #dilated_columns = fill_column_gaps(columns, table_bbox)
+    dilated_columns = columns
+    #dilated_rows = fill_row_gaps(rows, table_bbox)
+    dilated_rows = rows
+    for cell in cells:
+        column_rect = Rect()
+        for column_num in cell['column_nums']:
+            column_rect.include_rect(list(dilated_columns[column_num]['bbox']))
+        row_rect = Rect()
+        for row_num in cell['row_nums']:
+            row_rect.include_rect(list(dilated_rows[row_num]['bbox']))
+        cell_rect = column_rect.intersect(row_rect)
+        cell['bbox'] = list(cell_rect)
+
+    span_nums_by_cell, _, _ = slot_into_containers(cells, table_spans, overlap_threshold=0.001,
+                                               unique_assignment=True, forced_assignment=False)
+
+    for cell, cell_span_nums in zip(cells, span_nums_by_cell):
+        cell_spans = [table_spans[num] for num in cell_span_nums]
+        # TODO: Refine how text is extracted; should be character-based, not span-based;
+        # but need to associate 
+        # cell['cell_text'] = extract_text_from_spans(cell_spans, remove_integer_superscripts=False)  # TODO
+        cell['spans'] = cell_spans
+        
+    # Adjust the row, column, and cell bounding boxes to reflect the extracted text
+    num_rows = len(rows)
+    rows = sort_objects_top_to_bottom(rows)
+    num_columns = len(columns)
+    columns = sort_objects_left_to_right(columns)
+    min_y_values_by_row = defaultdict(list)
+    max_y_values_by_row = defaultdict(list)
+    min_x_values_by_column = defaultdict(list)
+    max_x_values_by_column = defaultdict(list)
+    for cell in cells:
+        min_row = min(cell["row_nums"])
+        max_row = max(cell["row_nums"])
+        min_column = min(cell["column_nums"])
+        max_column = max(cell["column_nums"])
+        for span in cell['spans']:
+            min_x_values_by_column[min_column].append(span['bbox'][0])
+            min_y_values_by_row[min_row].append(span['bbox'][1])
+            max_x_values_by_column[max_column].append(span['bbox'][2])
+            max_y_values_by_row[max_row].append(span['bbox'][3])
+    for row_num, row in enumerate(rows):
+        if len(min_x_values_by_column[0]) > 0:
+            row['bbox'][0] = min(min_x_values_by_column[0])
+        if len(min_y_values_by_row[row_num]) > 0:
+            row['bbox'][1] = min(min_y_values_by_row[row_num])
+        if len(max_x_values_by_column[num_columns-1]) > 0:
+            row['bbox'][2] = max(max_x_values_by_column[num_columns-1])
+        if len(max_y_values_by_row[row_num]) > 0:
+            row['bbox'][3] = max(max_y_values_by_row[row_num])
+    for column_num, column in enumerate(columns):
+        if len(min_x_values_by_column[column_num]) > 0:
+            column['bbox'][0] = min(min_x_values_by_column[column_num])
+        if len(min_y_values_by_row[0]) > 0:
+            column['bbox'][1] = min(min_y_values_by_row[0])
+        if len(max_x_values_by_column[column_num]) > 0:
+            column['bbox'][2] = max(max_x_values_by_column[column_num])
+        if len(max_y_values_by_row[num_rows-1]) > 0:
+            column['bbox'][3] = max(max_y_values_by_row[num_rows-1])
+    for cell in cells:
+        row_rect = Rect()
+        column_rect = Rect()
+        for row_num in cell['row_nums']:
+            row_rect.include_rect(list(rows[row_num]['bbox']))
+        for column_num in cell['column_nums']:
+            column_rect.include_rect(list(columns[column_num]['bbox']))
+        cell_rect = row_rect.intersect(column_rect)
+        if cell_rect.get_area() > 0:  # getArea()
+            cell['bbox'] = list(cell_rect)
+            pass
+
+    return cells, confidence_score
+
+
+def remove_supercell_overlap(supercell1, supercell2):
+    """
+    This function resolves overlap between supercells (supercells must be
+    disjoint) by iteratively shrinking supercells by the fewest grid cells
+    necessary to resolve the overlap.
+    Example:
+    If two supercells overlap at grid cell (R, C), and supercell #1 is less
+    confident than supercell #2, we eliminate either row R from supercell #1
+    or column C from supercell #1 by comparing the number of columns in row R
+    versus the number of rows in column C. If the number of columns in row R
+    is less than the number of rows in column C, we eliminate row R from
+    supercell #1. This resolves the overlap by removing fewer grid cells from
+    supercell #1 than if we eliminated column C from it.
+    """
+    common_rows = set(supercell1['row_numbers']).intersection(set(supercell2['row_numbers']))
+    common_columns = set(supercell1['column_numbers']).intersection(set(supercell2['column_numbers']))
+
+    # While the supercells have overlapping grid cells, continue shrinking the less-confident
+    # supercell one row or one column at a time
+    while len(common_rows) > 0 and len(common_columns) > 0:
+        # Try to shrink the supercell as little as possible to remove the overlap;
+        # if the supercell has fewer rows than columns, remove an overlapping column,
+        # because this removes fewer grid cells from the supercell;
+        # otherwise remove an overlapping row
+        if len(supercell2['row_numbers']) < len(supercell2['column_numbers']):
+            min_column = min(supercell2['column_numbers'])
+            max_column = max(supercell2['column_numbers'])
+            if max_column in common_columns:
+                common_columns.remove(max_column)
+                supercell2['column_numbers'].remove(max_column)
+            elif min_column in common_columns:
+                common_columns.remove(min_column)
+                supercell2['column_numbers'].remove(min_column)
+            else:
+                supercell2['column_numbers'] = []
+                common_columns = set()
+        else:
+            min_row = min(supercell2['row_numbers'])
+            max_row = max(supercell2['row_numbers'])
+            if max_row in common_rows:
+                common_rows.remove(max_row)
+                supercell2['row_numbers'].remove(max_row)
+            elif min_row in common_rows:
+                common_rows.remove(min_row)
+                supercell2['row_numbers'].remove(min_row)
+            else:
+                supercell2['row_numbers'] = []
+                common_rows = set()

requirements.txt

@@ -0,0 +1,9 @@
+-e git+https://github.com/mindee/doctr.git#egg=python-doctr[tf]
+streamlit>=0.65.0
+PyMuPDF>=1.16.0,!=1.18.11,!=1.18.12,!=1.19.5
+tf2onnx==1.13.0
+Pillow==9.0.1
+pytesseract==0.3.10
+torch==1.12.0
+torchvision==0.13.0
+numpy==1.21.6

tessdata/eng.traineddata

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8280aed0782fe27257a68ea10fe7ef324ca0f8d85bd2fd145d1c2b560bcb66ba
+size 15400601

weights/structure_wts.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46121ab2f4aba48a7d38624c861658ffeaacd0f305e95efcf66cb017e588b700
+size 14371957