[AUR-432] Add layout-aware table parsing PDF reader (#27)

* add OCRReader, MathPixReader and ExcelReader

* update test case for ocr reader

* reformat

* minor fix

knowledgehub/loaders/utils/table.py

@@ -0,0 +1,335 @@
+import csv
+from io import StringIO
+from typing import List, Optional, Tuple
+
+
+def check_col_conflicts(
+    col_a: List[str], col_b: List[str], thres: float = 0.15
+) -> bool:
+    """Check if 2 columns A and B has non-empty content in the same row
+    (to be used with merge_cols)
+
+    Args:
+        col_a: column A (list of str)
+        col_b: column B (list of str)
+        thres: percentage of overlapping allowed
+    Returns:
+        if number of overlapping greater than threshold
+    """
+    num_rows = len([cell for cell in col_a if cell])
+    assert len(col_a) == len(col_b)
+    conflict_count = 0
+    for cell_a, cell_b in zip(col_a, col_b):
+        if cell_a and cell_b:
+            conflict_count += 1
+    return conflict_count > num_rows * thres
+
+
+def merge_cols(col_a: List[str], col_b: List[str]) -> List[str]:
+    """Merge column A and B if they do not have conflict rows
+
+    Args:
+        col_a: column A (list of str)
+        col_b: column B (list of str)
+    Returns:
+        merged column
+    """
+    for r_id in range(len(col_a)):
+        if col_b[r_id]:
+            col_a[r_id] = col_a[r_id] + " " + col_b[r_id]
+    return col_a
+
+
+def add_index_col(csv_rows: List[List[str]]) -> List[List[str]]:
+    """Add index column as the first column of the table csv_rows
+
+    Args:
+        csv_rows: input table
+    Returns:
+        output table with index column
+    """
+    new_csv_rows = [["row id"] + [""] * len(csv_rows[0])]
+    for r_id, row in enumerate(csv_rows):
+        new_csv_rows.append([str(r_id + 1)] + row)
+    return new_csv_rows
+
+
+def compress_csv(csv_rows: List[List[str]]) -> List[List[str]]:
+    """Compress table csv_rows by merging sparse columns (merge_cols)
+
+    Args:
+        csv_rows: input table
+    Returns:
+        output: compressed table
+    """
+    csv_cols = [[r[c_id] for r in csv_rows] for c_id in range(len(csv_rows[0]))]
+    to_remove_col_ids = []
+    last_c_id = 0
+    for c_id in range(1, len(csv_cols)):
+        if not check_col_conflicts(csv_cols[last_c_id], csv_cols[c_id]):
+            to_remove_col_ids.append(c_id)
+            csv_cols[last_c_id] = merge_cols(csv_cols[last_c_id], csv_cols[c_id])
+        else:
+            last_c_id = c_id
+
+    csv_cols = [r for c_id, r in enumerate(csv_cols) if c_id not in to_remove_col_ids]
+    csv_rows = [[c[r_id] for c in csv_cols] for r_id in range(len(csv_cols[0]))]
+    return csv_rows
+
+
+def _get_rect_iou(gt_box: List[tuple], pd_box: List[tuple], iou_type=0) -> int:
+    """Intersection over union on layout rectangle
+
+    Args:
+        gt_box: List[tuple]
+            A list contains bounding box coordinates of ground truth
+        pd_box: List[tuple]
+            A list contains bounding box coordinates of prediction
+        iou_type: int
+            0: intersection / union, normal IOU
+            1: intersection / min(areas), useful when boxes are under/over-segmented
+
+        Input format: [(x1, y1), (x2, y1), (x2, y2), (x1, y2)]
+        Annotation for each element in bbox:
+        (x1, y1)        (x2, y1)
+            +-------+
+            |       |
+            |       |
+            +-------+
+        (x1, y2)        (x2, y2)
+
+    Returns:
+        Intersection over union value
+    """
+
+    assert iou_type in [0, 1], "Only support 0: origin iou, 1: intersection / min(area)"
+
+    # determine the (x, y)-coordinates of the intersection rectangle
+    # gt_box: [(x1, y1), (x2, y1), (x2, y2), (x1, y2)]
+    # pd_box: [(x1, y1), (x2, y1), (x2, y2), (x1, y2)]
+    x_left = max(gt_box[0][0], pd_box[0][0])
+    y_top = max(gt_box[0][1], pd_box[0][1])
+    x_right = min(gt_box[2][0], pd_box[2][0])
+    y_bottom = min(gt_box[2][1], pd_box[2][1])
+
+    # compute the area of intersection rectangle
+    interArea = max(0, x_right - x_left) * max(0, y_bottom - y_top)
+
+    # compute the area of both the prediction and ground-truth
+    # rectangles
+    gt_area = (gt_box[2][0] - gt_box[0][0]) * (gt_box[2][1] - gt_box[0][1])
+    pd_area = (pd_box[2][0] - pd_box[0][0]) * (pd_box[2][1] - pd_box[0][1])
+
+    # compute the intersection over union by taking the intersection
+    # area and dividing it by the sum of prediction + ground-truth
+    # areas - the interesection area
+    if iou_type == 0:
+        iou = interArea / float(gt_area + pd_area - interArea)
+    elif iou_type == 1:
+        iou = interArea / max(min(gt_area, pd_area), 1)
+
+    # return the intersection over union value
+    return iou
+
+
+def get_table_from_ocr(ocr_list: List[dict], table_list: List[dict]):
+    """Get list of text lines belong to table regions specified by table_list
+
+    Args:
+        ocr_list: list of OCR output in Casia format (Flax)
+        table_list: list of table output in Casia format (Flax)
+
+    Returns:
+        _type_: _description_
+    """
+    table_texts = []
+    for table in table_list:
+        if table["type"] != "table":
+            continue
+        cur_table_texts = []
+        for ocr in ocr_list:
+            _iou = _get_rect_iou(table["location"], ocr["location"], iou_type=1)
+            if _iou > 0.8:
+                cur_table_texts.append(ocr["text"])
+        table_texts.append(cur_table_texts)
+
+    return table_texts
+
+
+def make_markdown_table(array: List[List[str]]) -> str:
+    """Convert table rows in list format to markdown string
+
+    Args:
+        Python list with rows of table as lists
+        First element as header.
+        Example Input:
+                [["Name", "Age", "Height"],
+                ["Jake", 20, 5'10],
+                ["Mary", 21, 5'7]]
+    Returns:
+        String to put into a .md file
+    """
+    array = compress_csv(array)
+    array = add_index_col(array)
+    markdown = "\n" + str("| ")
+
+    for e in array[0]:
+        to_add = " " + str(e) + str(" |")
+        markdown += to_add
+    markdown += "\n"
+
+    markdown += "| "
+    for i in range(len(array[0])):
+        markdown += str("--- | ")
+    markdown += "\n"
+
+    for entry in array[1:]:
+        markdown += str("| ")
+        for e in entry:
+            to_add = str(e) + str(" | ")
+            markdown += to_add
+        markdown += "\n"
+
+    return markdown + "\n"
+
+
+def parse_csv_string_to_list(csv_str: str) -> List[List[str]]:
+    """Convert CSV string to list of rows
+
+    Args:
+        csv_str: input CSV string
+
+    Returns:
+        Output table in list format
+    """
+    io = StringIO(csv_str)
+    csv_reader = csv.reader(io, delimiter=",")
+    rows = [row for row in csv_reader]
+    return rows
+
+
+def format_cell(cell: str, length_limit: Optional[int] = None) -> str:
+    """Format cell content by remove redundant character and enforce length limit
+
+    Args:
+        cell: input cell text
+        length_limit: limit of text length.
+
+    Returns:
+        new cell text
+    """
+    cell = cell.replace("\n", " ")
+    if length_limit:
+        cell = cell[:length_limit]
+    return cell
+
+
+def extract_tables_from_csv_string(
+    csv_content: str, table_texts: List[List[str]]
+) -> Tuple[List[str], str]:
+    """Extract list of table from FullOCR output
+    (csv_content) with the specified table_texts
+
+    Args:
+        csv_content: CSV output from FullOCR pipeline
+        table_texts: list of table texts extracted
+        from get_table_from_ocr()
+
+    Returns:
+        List of tables and non-text content
+    """
+    rows = parse_csv_string_to_list(csv_content)
+    used_row_ids = []
+    table_csv_list = []
+    for table in table_texts:
+        cur_rows = []
+        for row_id, row in enumerate(rows):
+            scores = [
+                any(cell in cell_reference for cell in table)
+                for cell_reference in row
+                if cell_reference
+            ]
+            score = sum(scores) / len(scores)
+            if score > 0.5 and row_id not in used_row_ids:
+                used_row_ids.append(row_id)
+                cur_rows.append([format_cell(cell) for cell in row])
+        if cur_rows:
+            table_csv_list.append(make_markdown_table(cur_rows))
+        else:
+            print("table not matched", table)
+
+    non_table_rows = [
+        row for row_id, row in enumerate(rows) if row_id not in used_row_ids
+    ]
+    non_table_text = "\n".join(
+        " ".join(format_cell(cell) for cell in row) for row in non_table_rows
+    )
+    return table_csv_list, non_table_text
+
+
+def strip_special_chars_markdown(text: str) -> str:
+    """Strip special characters from input text in markdown table format"""
+    return text.replace("|", "").replace(":---:", "").replace("---", "")
+
+
+def markdown_to_list(markdown_text: str, pad_to_max_col: Optional[bool] = True):
+    rows = []
+    lines = markdown_text.split("\n")
+    markdown_lines = [line.strip() for line in lines if " | " in line]
+
+    for row in markdown_lines:
+        tmp = row
+        # Get rid of leading and trailing '|'
+        if tmp.startswith("|"):
+            tmp = tmp[1:]
+        if tmp.endswith("|"):
+            tmp = tmp[:-1]
+
+        # Split line and ignore column whitespace
+        clean_line = tmp.split("|")
+        if not all(c == "" for c in clean_line):
+            # Append clean row data to rows variable
+            rows.append(clean_line)
+
+    # Get rid of syntactical sugar to indicate header (2nd row)
+    rows = [row for row in rows if "---" not in " ".join(row)]
+    max_cols = max(len(row) for row in rows)
+    if pad_to_max_col:
+        rows = [row + [""] * (max_cols - len(row)) for row in rows]
+    return rows
+
+
+def parse_markdown_text_to_tables(text: str) -> Tuple[List[str], List[str]]:
+    """Convert markdown text to list of non-table spans and table spans
+
+    Args:
+        text: input markdown text
+
+    Returns:
+        list of table spans and non-table spans
+    """
+    # init empty tables and texts list
+    tables = []
+    texts = []
+
+    # split input by line break
+    lines = text.split("\n")
+    cur_table = []
+    cur_text: List[str] = []
+    for line in lines:
+        line = line.strip()
+        if line.startswith("|"):
+            if len(cur_text) > 0:
+                texts.append(cur_text)
+                cur_text = []
+            cur_table.append(line)
+        else:
+            # add new table to the list
+            if len(cur_table) > 0:
+                tables.append(cur_table)
+                cur_table = []
+            cur_text.append(line)
+
+    table_texts = ["\n".join(table) for table in tables]
+    non_table_texts = ["\n".join(text) for text in texts]
+    return table_texts, non_table_texts