init commit

.gitattributes

@@ -33,3 +33,5 @@ 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
+*.ttc filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

SimSong.ttc

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

app.py

@@ -0,0 +1,266 @@
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from matplotlib.patches import Patch
+import io
+import cv2
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+import csv
+import pandas as pd
+
+from ultralytics import YOLO
+import torch
+
+from paddleocr import PaddleOCR
+import postprocess
+
+import gradio as gr
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+detection_model = YOLO('yolov8/runs/detect/yolov8s-custom-detection/weights/best.pt').to(device)
+structure_model = YOLO('yolov8/runs/detect/yolov8s-custom-structure-all/weights/best.pt').to(device)
+ocr_model = PaddleOCR(use_angle_cls=True, lang="ch", det_limit_side_len=1920)  # TODO use large det_limit_side_len to get better OCR result
+
+detection_class_names = ['table', 'table rotated']
+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.5,
+    "table projected row header": 0.5,
+    "table spanning cell": 0.5,
+    "no object": 10
+}
+
+
+def table_detection(image):
+    imgsz = 800
+    pred = detection_model.predict(image, imgsz=imgsz)
+    pred = pred[0].boxes
+    result = pred.cpu().numpy()
+    result_list = [list(result.xywhn[i]) + [result.conf[i], result.cls[i]] for i in range(result.shape[0])]
+    return result_list
+
+
+def table_structure(image):
+    imgsz = 1024
+    pred = structure_model.predict(image, imgsz=imgsz)
+    pred = pred[0].boxes
+    result = pred.cpu().numpy()
+    result_list = [list(result.xywhn[i]) + [result.conf[i], result.cls[i]] for i in range(result.shape[0])]
+    return result_list
+
+
+def crop_image(image, detection_result):
+    # crop_filenames = []
+    width = image.shape[1]
+    height = image.shape[0]
+    # print(width, height)
+    for i, result in enumerate(detection_result[:1]):  # TODO only return first detected table
+        class_id = int(result[5])
+        score = float(result[4])
+        min_x = result[0]
+        min_y = result[1]
+        w = result[2]
+        h = result[3]
+        
+        # x1 = max(0, int((min_x-w/2-0.02)*width))  # TODO expand 2%
+        # y1 = max(0, int((min_y-h/2-0.02)*height))  # TODO expand 2%
+        # x2 = min(width, int((min_x+w/2+0.02)*width))  # TODO expand 2%
+        # y2 = min(height, int((min_y+h/2+0.02)*height))  # TODO expand 2%
+        x1 = max(0, int((min_x-w/2)*width)-10)  # TODO expand 10px
+        y1 = max(0, int((min_y-h/2)*height)-10)  # TODO expand 10px
+        x2 = min(width, int((min_x+w/2)*width)+10)  # TODO expand 10px
+        y2 = min(height, int((min_y+h/2)*height)+10)  # TODO expand 10px
+        # print(x1, y1, x2, y2)
+        crop_image = image[y1:y2, x1:x2, :]
+        # crop_filename = filename[:-4]+'_'+str(i)+'_'+detection_class_names[class_id]+filename[-4:]
+        # crop_filenames.append(crop_filename)
+        # cv2.imwrite(crop_filename, crop_image)
+    return crop_image
+
+
+def convert_stucture(ocr_result, image, structure_result):
+    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)
+    
+    page_tokens = ocr_result
+    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(image, table_structures, cells):
+    width = image.shape[1]
+    height = image.shape[0]
+    # print(width, height)
+    empty_image = np.zeros((height, width, 3), np.uint8)
+    empty_image.fill(255)
+    empty_image = Image.fromarray(cv2.cvtColor(empty_image, cv2.COLOR_BGR2RGB))
+    draw = ImageDraw.Draw(empty_image)
+    fontStyle = ImageFont.truetype("SimSong.ttc", 10, encoding="utf-8")
+    
+    num_cols = len(table_structures['columns'])
+    num_rows = len(table_structures['rows'])
+    data_rows = [['' for _ in range(num_cols)] for _ in range(num_rows)]
+    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])
+        col_num = cell['column_nums'][0]
+        row_num = cell['row_nums'][0]
+        spans = cell['spans']
+        text = ''
+        for span in spans:
+            if 'text' in span:
+                text += span['text']     
+        data_rows[row_num][col_num] = text
+        
+        # print('text:', text)
+        text_len = len(text)
+        # print('text_len:', text_len)
+        cell_width = x2-x1
+        # print('cell_width:', cell_width)
+        num_per_line = cell_width//10
+        # print('num_per_line:', num_per_line)
+        if num_per_line != 0:
+            line_num = text_len//num_per_line
+        else:
+            line_num = 0
+        # print('line_num:', line_num)
+        new_text = text[:num_per_line]+'\n'
+        for j in range(line_num):
+            new_text += text[(j+1)*num_per_line:(j+2)*num_per_line]+'\n'
+        # print('new_text:', new_text)
+        text = new_text
+        
+        cv2.rectangle(image, (x1, y1), (x2, y2), color=(0,255,0))
+        cv2.putText(image, str(row_num)+'-'+str(col_num), (x1, y1+30), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(0,0,255))
+        
+        # cv2.rectangle(empty_image, (x1, y1), (x2, y2), color=(0,0,255))
+        # cv2.putText(empty_image, str(row_num)+'-'+str(col_num), (x1-10, y1), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(0,0,255))
+        # cv2.putText(empty_image, text, (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(0,0,255))
+        draw.rectangle([(x1, y1), (x2, y2)], (255,255,255), (0,255,0))
+        draw.text((x1-20, y1), str(row_num)+'-'+str(col_num), (255,0,0), font=fontStyle)
+        draw.text((x1, y1), text, (0,0,255), font=fontStyle)
+
+    df = pd.DataFrame(data_rows[1:], columns=data_rows[0])
+    return image, df, df.to_json()
+
+
+def ocr(image):
+    result = ocr_model.ocr(image, cls=True)
+    result = result[0]
+    new_result = []
+    if result is not None:
+        bounding_boxes = [line[0] for line in result]
+        txts = [line[1][0] for line in result]
+        scores = [line[1][1] for line in result]
+        # print('txts:', txts)
+        # print('scores:', scores)
+        # print('bounding_boxes:', bounding_boxes)
+        for label, bbox in zip(txts, bounding_boxes):
+            new_result.append({'bbox': [bbox[0][0], bbox[0][1], bbox[2][0], bbox[2][1]], 'text': label})
+    
+    return new_result
+
+
+def detect_and_crop_table(image):
+    detection_result = table_detection(image)
+    # print('detection_result:', detection_result)
+    cropped_table = crop_image(image, detection_result)
+
+    return cropped_table
+
+
+def recognize_table(image, ocr_result):
+    structure_result = table_structure(image)
+    print('structure_result:', structure_result)
+    table_structures, cells, confidence_score = convert_stucture(ocr_result, image, structure_result)
+    print('table_structures:', table_structures)
+    print('cells:', cells)
+    print('confidence_score:', confidence_score)
+    image, df, data = visualize_cells(image, table_structures, cells)
+        
+    return image, df, data
+
+
+def process_pdf(image):
+    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    
+    cropped_table = detect_and_crop_table(image)
+    
+    ocr_result = ocr(cropped_table)
+    # print('ocr_result:', ocr_result)
+
+    image, df, data = recognize_table(cropped_table, ocr_result)
+    print('df:', df)
+    
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    return image, df, data
+    
+
+title = "Demo: table detection & recognition with Table Structure Recognition (Yolov8)."
+description = """Demo for table extraction with the Table Structure Recognition (Yolov8)."""
+examples = [['image.png'], ['mistral_paper.png']]
+
+app = gr.Interface(fn=process_pdf, 
+                     inputs=gr.Image(type="numpy"), 
+                     outputs=[gr.Image(type="numpy", label="Detected table"), gr.Dataframe(label="Table as CSV"), gr.JSON(label="Data as JSON")],
+                     title=title,
+                     description=description,
+                     examples=examples)
+app.queue()
+# app.launch(debug=True, share=True)
+app.launch()

postprocess.py

@@ -0,0 +1,887 @@
+"""
+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.
+    """
+
+    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.
+    """
+
+    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,13 @@
+torch
+#git+https://github.com/nielsrogge/transformers.git@convert_table_transformer_new_checkpoints
+transformers
+easyocr
+matplotlib
+Pillow
+pandas
+ultralytics
+PyMuPDF
+opencv-python
+gradio
+paddlepaddle-gpu
+paddleocr

tatr-app.py

@@ -0,0 +1,294 @@
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from matplotlib.patches import Patch
+import io
+from PIL import Image, ImageDraw
+import numpy as np
+import csv
+import pandas as pd
+
+from torchvision import transforms
+
+from transformers import AutoModelForObjectDetection
+import torch
+
+import easyocr
+
+import gradio as gr
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+class MaxResize(object):
+    def __init__(self, max_size=800):
+        self.max_size = max_size
+
+    def __call__(self, image):
+        width, height = image.size
+        current_max_size = max(width, height)
+        scale = self.max_size / current_max_size
+        resized_image = image.resize((int(round(scale*width)), int(round(scale*height))))
+        
+        return resized_image
+
+detection_transform = transforms.Compose([
+    MaxResize(800),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+structure_transform = transforms.Compose([
+    MaxResize(1000),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+# load table detection model
+# processor = TableTransformerImageProcessor(max_size=800)
+model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-detection", revision="no_timm").to(device)
+
+# load table structure recognition model
+# structure_processor = TableTransformerImageProcessor(max_size=1000)
+structure_model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-structure-recognition-v1.1-all").to(device)
+
+# load EasyOCR reader
+reader = easyocr.Reader(['en'])
+
+
+# for output bounding box post-processing
+def box_cxcywh_to_xyxy(x):
+    x_c, y_c, w, h = x.unbind(-1)
+    b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)]
+    return torch.stack(b, dim=1)
+
+
+def rescale_bboxes(out_bbox, size):
+    width, height = size
+    boxes = box_cxcywh_to_xyxy(out_bbox)
+    boxes = boxes * torch.tensor([width, height, width, height], dtype=torch.float32)
+    return boxes
+
+
+def outputs_to_objects(outputs, img_size, id2label):
+    m = outputs.logits.softmax(-1).max(-1)
+    pred_labels = list(m.indices.detach().cpu().numpy())[0]
+    pred_scores = list(m.values.detach().cpu().numpy())[0]
+    pred_bboxes = outputs['pred_boxes'].detach().cpu()[0]
+    pred_bboxes = [elem.tolist() for elem in rescale_bboxes(pred_bboxes, img_size)]
+
+    objects = []
+    for label, score, bbox in zip(pred_labels, pred_scores, pred_bboxes):
+        class_label = id2label[int(label)]
+        if not class_label == 'no object':
+            objects.append({'label': class_label, 'score': float(score),
+                            'bbox': [float(elem) for elem in bbox]})
+
+    return objects
+
+
+def fig2img(fig):
+    """Convert a Matplotlib figure to a PIL Image and return it"""
+    buf = io.BytesIO()
+    fig.savefig(buf)
+    buf.seek(0)
+    image = Image.open(buf)
+    return image
+
+
+def visualize_detected_tables(img, det_tables):
+    plt.imshow(img, interpolation="lanczos")
+    fig = plt.gcf()
+    fig.set_size_inches(20, 20)
+    ax = plt.gca()
+
+    for det_table in det_tables:
+        bbox = det_table['bbox']
+
+        if det_table['label'] == 'table':
+            facecolor = (1, 0, 0.45)
+            edgecolor = (1, 0, 0.45)
+            alpha = 0.3
+            linewidth = 2
+            hatch='//////'
+        elif det_table['label'] == 'table rotated':
+            facecolor = (0.95, 0.6, 0.1)
+            edgecolor = (0.95, 0.6, 0.1)
+            alpha = 0.3
+            linewidth = 2
+            hatch='//////'
+        else:
+            continue
+
+        rect = patches.Rectangle(bbox[:2], bbox[2]-bbox[0], bbox[3]-bbox[1], linewidth=linewidth,
+                                    edgecolor='none',facecolor=facecolor, alpha=0.1)
+        ax.add_patch(rect)
+        rect = patches.Rectangle(bbox[:2], bbox[2]-bbox[0], bbox[3]-bbox[1], linewidth=linewidth,
+                                    edgecolor=edgecolor,facecolor='none',linestyle='-', alpha=alpha)
+        ax.add_patch(rect)
+        rect = patches.Rectangle(bbox[:2], bbox[2]-bbox[0], bbox[3]-bbox[1], linewidth=0,
+                                    edgecolor=edgecolor,facecolor='none',linestyle='-', hatch=hatch, alpha=0.2)
+        ax.add_patch(rect)
+
+    plt.xticks([], [])
+    plt.yticks([], [])
+
+    legend_elements = [Patch(facecolor=(1, 0, 0.45), edgecolor=(1, 0, 0.45),
+                                label='Table', hatch='//////', alpha=0.3),
+                        Patch(facecolor=(0.95, 0.6, 0.1), edgecolor=(0.95, 0.6, 0.1),
+                                label='Table (rotated)', hatch='//////', alpha=0.3)]
+    plt.legend(handles=legend_elements, bbox_to_anchor=(0.5, -0.02), loc='upper center', borderaxespad=0,
+                    fontsize=10, ncol=2)
+    plt.gcf().set_size_inches(10, 10)
+    plt.axis('off')
+
+    return fig
+
+
+def detect_and_crop_table(image):
+    # prepare image for the model
+    # pixel_values = processor(image, return_tensors="pt").pixel_values
+    pixel_values = detection_transform(image).unsqueeze(0).to(device)
+
+    # forward pass
+    with torch.no_grad():
+        outputs = model(pixel_values)
+
+    # postprocess to get detected tables
+    id2label = model.config.id2label
+    id2label[len(model.config.id2label)] = "no object"
+    detected_tables = outputs_to_objects(outputs, image.size, id2label)
+
+    # visualize
+    # fig = visualize_detected_tables(image, detected_tables)
+    # image = fig2img(fig)
+
+    # crop first detected table out of image
+    cropped_table = image.crop(detected_tables[0]["bbox"])
+
+    return cropped_table
+
+
+def recognize_table(image):
+    # prepare image for the model
+    # pixel_values = structure_processor(images=image, return_tensors="pt").pixel_values
+    pixel_values = structure_transform(image).unsqueeze(0).to(device)
+
+    # forward pass
+    with torch.no_grad():
+        outputs = structure_model(pixel_values)
+
+    # postprocess to get individual elements
+    id2label = structure_model.config.id2label
+    id2label[len(structure_model.config.id2label)] = "no object"
+    cells = outputs_to_objects(outputs, image.size, id2label)
+
+    # visualize cells on cropped table
+    draw = ImageDraw.Draw(image)
+
+    for cell in cells:
+        draw.rectangle(cell["bbox"], outline="red")
+        
+    return image, cells
+
+
+def get_cell_coordinates_by_row(table_data):
+    # Extract rows and columns
+    rows = [entry for entry in table_data if entry['label'] == 'table row']
+    columns = [entry for entry in table_data if entry['label'] == 'table column']
+
+    # Sort rows and columns by their Y and X coordinates, respectively
+    rows.sort(key=lambda x: x['bbox'][1])
+    columns.sort(key=lambda x: x['bbox'][0])
+
+    # Function to find cell coordinates
+    def find_cell_coordinates(row, column):
+        cell_bbox = [column['bbox'][0], row['bbox'][1], column['bbox'][2], row['bbox'][3]]
+        return cell_bbox
+
+    # Generate cell coordinates and count cells in each row
+    cell_coordinates = []
+
+    for row in rows:
+        row_cells = []
+        for column in columns:
+            cell_bbox = find_cell_coordinates(row, column)
+            row_cells.append({'column': column['bbox'], 'cell': cell_bbox})
+
+        # Sort cells in the row by X coordinate
+        row_cells.sort(key=lambda x: x['column'][0])
+
+        # Append row information to cell_coordinates
+        cell_coordinates.append({'row': row['bbox'], 'cells': row_cells, 'cell_count': len(row_cells)})
+
+    # Sort rows from top to bottom
+    cell_coordinates.sort(key=lambda x: x['row'][1])
+
+    return cell_coordinates
+
+
+def apply_ocr(cell_coordinates, cropped_table):
+    # let's OCR row by row
+    data = dict()
+    max_num_columns = 0
+    for idx, row in enumerate(cell_coordinates):
+      row_text = []
+      for cell in row["cells"]:
+        # crop cell out of image
+        cell_image = np.array(cropped_table.crop(cell["cell"]))
+        # apply OCR
+        result = reader.readtext(np.array(cell_image))
+        if len(result) > 0:
+          text = " ".join([x[1] for x in result])
+          row_text.append(text)
+
+      if len(row_text) > max_num_columns:
+          max_num_columns = len(row_text)
+      
+      data[str(idx)] = row_text
+
+    # pad rows which don't have max_num_columns elements
+    # to make sure all rows have the same number of columns
+    for idx, row_data in data.copy().items():
+        if len(row_data) != max_num_columns:
+          row_data = row_data + ["" for _ in range(max_num_columns - len(row_data))]
+        data[str(idx)] = row_data
+
+    # write to csv
+    with open('output.csv','w') as result_file:
+        wr = csv.writer(result_file, dialect='excel')
+    
+        for row, row_text in data.items():
+            wr.writerow(row_text)
+
+    # return as Pandas dataframe
+    df = pd.read_csv('output.csv')
+
+    return df, data
+
+
+def process_pdf(image):
+    cropped_table = detect_and_crop_table(image)
+
+    image, cells = recognize_table(cropped_table)
+
+    cell_coordinates = get_cell_coordinates_by_row(cells)
+
+    df, data = apply_ocr(cell_coordinates, image)
+
+    return image, df, data
+    
+
+title = "Demo: table detection & recognition with Table Transformer (TATR)."
+description = """Demo for table extraction with the Table Transformer. First, table detection is performed on the input image using https://huggingface.co/microsoft/table-transformer-detection,
+after which the detected table is extracted and https://huggingface.co/microsoft/table-transformer-structure-recognition-v1.1-all is leveraged to recognize the individual rows, columns and cells. OCR is then performed per cell, row by row."""
+examples = [['image.png'], ['mistral_paper.png']]
+
+app = gr.Interface(fn=process_pdf, 
+                     inputs=gr.Image(type="pil"), 
+                     outputs=[gr.Image(type="pil", label="Detected table"), gr.Dataframe(label="Table as CSV"), gr.JSON(label="Data as JSON")],
+                     title=title,
+                     description=description,
+                     examples=examples)
+app.queue()
+app.launch(debug=True)

yolov8/runs/detect/yolov8s-custom-detection/weights/best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d7c77d9723839f582e02377794dccee2ed745f72e08f4818d1ed5a7f7c3e591
+size 22520345

yolov8/runs/detect/yolov8s-custom-structure-all/weights/best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:41877a8b980bd7c28139c267b3c3b1f0ffefb6babeaee8340c1680d2d623a794
+size 22527577