| |
| |
|
|
| import re |
| from typing import List, Union, Optional, Literal |
| import numpy as np |
| import numpy.typing as npt |
| import pandas as pd |
|
|
|
|
| def assign_boxes( |
| box: Union[List[float], npt.NDArray[np.float64]], |
| candidate_boxes: npt.NDArray[np.float64], |
| delta: float = 2.0, |
| min_overlap: float = 0.25, |
| mode: Literal["cell", "row", "column"] = "cell", |
| ) -> npt.NDArray[np.int_]: |
| """ |
| Assigns the best candidate boxes to a reference `box` based on overlap. |
| |
| If mode is "cell", the overlap is calculated using surface area overlap. |
| If mode is "row", the overlap is calculated using row height overlap. |
| If mode is "column", the overlap is calculated using column width overlap. |
| |
| If delta > 1, it will look for multiple matches, |
| using candidates with score >= max_overlap / delta. |
| |
| Args: |
| box (list or numpy.ndarray): Reference bounding box [x_min, y_min, x_max, y_max]. |
| candidate_boxes (numpy.ndarray [N, 4]): Array of candidate bounding boxes. |
| delta (float, optional): Factor for matches relative to the best overlap. Defaults to 2.0. |
| min_overlap (float, optional): Minimum required overlap for a match. Defaults to 0.25. |
| mode (str, optional): Mode to assign boxes ("cell", "row", or "column"). Defaults to "cell". |
| |
| Returns: |
| numpy.ndarray [M]: Indices of the matched boxes sorted by decreasing overlap. |
| Returns an empty array if no matches are found. |
| """ |
| if not len(candidate_boxes): |
| return np.array([], dtype=np.int_) |
|
|
| x0_1, y0_1, x1_1, y1_1 = box |
| x0_2, y0_2, x1_2, y1_2 = ( |
| candidate_boxes[:, 0], |
| candidate_boxes[:, 1], |
| candidate_boxes[:, 2], |
| candidate_boxes[:, 3], |
| ) |
|
|
| |
| inter_y0 = np.maximum(y0_1, y0_2) |
| inter_y1 = np.minimum(y1_1, y1_2) |
| inter_x0 = np.maximum(x0_1, x0_2) |
| inter_x1 = np.minimum(x1_1, x1_2) |
|
|
| if mode == "cell": |
| inter_area = np.maximum(0, inter_y1 - inter_y0) * np.maximum(0, inter_x1 - inter_x0) |
| box_area = (y1_1 - y0_1) * (x1_1 - x0_1) |
| overlap = inter_area / (box_area + 1e-6) |
| elif mode == "row": |
| inter_area = np.maximum(0, inter_y1 - inter_y0) |
| box_area = y1_1 - y0_1 |
| overlap = inter_area / (box_area + 1e-6) |
| elif mode == "column": |
| inter_area = np.maximum(0, inter_x1 - inter_x0) |
| box_area = x1_1 - x0_1 |
| overlap = inter_area / (box_area + 1e-6) |
| else: |
| raise ValueError(f"Invalid mode: {mode}") |
|
|
| max_overlap = np.max(overlap) |
| if max_overlap <= min_overlap: |
| return np.array([], dtype=np.int_) |
|
|
| n = len(np.where(overlap >= (max_overlap / delta))[0]) if delta > 1 else 1 |
| matches = np.argsort(-overlap)[:n] |
| return matches |
|
|
|
|
| def merge_text_in_cell(df_cell: pd.DataFrame) -> pd.DataFrame: |
| """ |
| Merges text from multiple rows into a single cell and recalculates its bounding box. |
| Values are sorted by rounded (y, x) coordinates. |
| |
| Args: |
| df_cell (pandas.DataFrame): DataFrame containing cells to merge. |
| |
| Returns: |
| pandas.DataFrame: Updated DataFrame with merged text and a single bounding box. |
| """ |
| boxes = np.stack(df_cell["box"].values) |
|
|
| df_cell["x"] = (boxes[:, 0] - boxes[:, 0].min()) // 10 |
| df_cell["y"] = (boxes[:, 1] - boxes[:, 1].min()) // 10 |
| df_cell = df_cell.sort_values(["y", "x"]) |
|
|
| text = " ".join(df_cell["text"].values.tolist()) |
| df_cell["text"] = text |
| df_cell = df_cell.head(1) |
| df_cell["box"] = df_cell["cell"] |
| df_cell.drop(["x", "y"], axis=1, inplace=True) |
|
|
| return df_cell |
|
|
|
|
| def remove_empty_row(mat: List[List[str]]) -> List[List[str]]: |
| """ |
| Remove empty rows from a matrix. |
| |
| Args: |
| mat (list[list]): The matrix to remove empty rows from. |
| |
| Returns: |
| list[list]: The matrix with empty rows removed. |
| """ |
| mat_filter = [] |
| for row in mat: |
| if max([len(c) for c in row]): |
| mat_filter.append(row) |
| return mat_filter |
|
|
|
|
| def build_markdown( |
| df: pd.DataFrame, |
| remove_empty: bool = True, |
| n_rows: Optional[int] = None, |
| repeat_single: bool = False, |
| ) -> Union[List[List[str]], npt.NDArray[np.str_]]: |
| """ |
| Convert a dataframe into a markdown table. |
| |
| Args: |
| df (pandas.DataFrame): The dataframe to convert with columns 'col_ids', |
| 'row_ids', and 'text'. |
| remove_empty (bool, optional): Whether to remove empty rows & cols. Defaults to True. |
| n_rows (int, optional): Number of rows. Inferred from df if None. Defaults to None. |
| repeat_single (bool, optional): Whether to repeat single element in rows. |
| Defaults to False. |
| |
| Returns: |
| list[list[str]] or numpy.ndarray: A list of lists or array representing the markdown table. |
| """ |
| df = df.reset_index(drop=True) |
| n_cols = max([np.max(c) for c in df['col_ids'].values]) |
| if n_rows is None: |
| n_rows = max([np.max(c) for c in df['row_ids'].values]) |
| else: |
| n_rows = max( |
| n_rows - 1, |
| max([np.max(c) for c in df['row_ids'].values]) |
| ) |
|
|
| mat = np.empty((n_rows + 1, n_cols + 1), dtype=str).tolist() |
|
|
| for i in range(len(df)): |
| if isinstance(df["row_ids"][i], int) or isinstance(df["col_ids"][i], int): |
| continue |
| for r in df["row_ids"][i]: |
| for c in df["col_ids"][i]: |
| mat[r][c] = (mat[r][c] + " " + df["text"][i]).strip() |
|
|
| |
| if remove_empty: |
| mat = remove_empty_row(mat) |
| mat = np.array(remove_empty_row(np.array(mat).T.tolist())).T.tolist() |
|
|
| if repeat_single: |
| new_mat = [] |
| for row in mat: |
| if sum([len(c) > 0 for c in row]) == 1: |
| txt = [c for c in row if len(c)][0] |
| new_mat.append([txt for _ in range(len(row))]) |
| else: |
| new_mat.append(row) |
| mat = np.array(new_mat) |
|
|
| return mat |
|
|
|
|
| def display_markdown( |
| data: List[List[str]], show: bool = True, use_header: bool = True |
| ) -> str: |
| """ |
| Convert a list of lists of strings into a markdown table. |
| If show is True, use_header will be set to True. |
| |
| Args: |
| data (list[list[str]]): The table data. The first sublist should contain headers. |
| show (bool, optional): Whether to display the table. Defaults to True. |
| use_header (bool, optional): Whether to use the first sublist as headers. Defaults to True. |
| |
| Returns: |
| str: A markdown-formatted table as a string. |
| """ |
| if show: |
| use_header = True |
| data = [[re.sub(r'\n', ' ', c) for c in row] for row in data] |
|
|
| if not len(data): |
| return "EMPTY TABLE" |
|
|
| max_cols = max(len(row) for row in data) |
| data = [row + [""] * (max_cols - len(row)) for row in data] |
|
|
| if use_header: |
| header = "| " + " | ".join(data[0]) + " |" |
| separator = "| " + " | ".join(["---"] * max_cols) + " |" |
| body = "\n".join("| " + " | ".join(row) + " |" for row in data[1:]) |
| markdown_table = ( |
| f"{header}\n{separator}\n{body}" if body else f"{header}\n{separator}" |
| ) |
|
|
| if show: |
| from IPython.display import display, Markdown |
| markdown_table = re.sub(r'\$', r'\\$', markdown_table) |
| markdown_table = re.sub(r'\%', r'\\%', markdown_table) |
| display(Markdown(markdown_table)) |
|
|
| else: |
| markdown_table = "\n".join("| " + " | ".join(row) + " |" for row in data) |
|
|
| return markdown_table |
|
|
