| import json |
| import cv2 |
| import os |
|
|
| def get_visible_cell_coords(table_properties, table_x, table_y): |
| rows = table_properties.get("rows", 0) |
| columns = table_properties.get("columns", 0) |
| column_widths = table_properties.get("columnWidths", {}) |
| row_heights = table_properties.get("rowHeights", {}) |
| merged_cells = table_properties.get("mergedCells", {}) |
| hidden_cells = table_properties.get("hiddenCells", {}) |
|
|
| DEFAULT_WIDTH = 100 |
| DEFAULT_HEIGHT = 30 |
|
|
| def get_col_width(col): |
| return column_widths.get(str(col), DEFAULT_WIDTH) |
|
|
| def get_row_height(row): |
| return row_heights.get(str(row), DEFAULT_HEIGHT) |
|
|
| |
| merged_spanned_cells = set() |
| for key, merge_info in merged_cells.items(): |
| base_row, base_col = map(int, key.split('-')) |
| rowspan = merge_info.get('rowspan', 1) |
| colspan = merge_info.get('colspan', 1) |
| for r in range(base_row, base_row + rowspan): |
| for c in range(base_col, base_col + colspan): |
| if (r, c) != (base_row, base_col): |
| merged_spanned_cells.add((r, c)) |
|
|
| result = {} |
|
|
| for row in range(rows): |
| for col in range(columns): |
| coord_key = f"{row}-{col}" |
| if hidden_cells.get(coord_key): |
| continue |
| if (row, col) in merged_spanned_cells: |
| continue |
|
|
| |
| x = sum(get_col_width(c) for c in range(col)) |
| y = sum(get_row_height(r) for r in range(row)) |
|
|
| |
| if coord_key in merged_cells: |
| colspan = merged_cells[coord_key].get("colspan", 1) |
| rowspan = merged_cells[coord_key].get("rowspan", 1) |
| else: |
| colspan = 1 |
| rowspan = 1 |
|
|
| width = sum(get_col_width(c) for c in range(col, col + colspan)) |
| height = sum(get_row_height(r) for r in range(row, row + rowspan)) |
|
|
| result[(row, col)] = { |
| "x": int(x + table_x), |
| "y": int(y + table_y), |
| "width": int(width), |
| "height": int(height) |
| } |
|
|
| return result |
|
|
| folder_path = "/Users/tuvn18/Desktop/tuvn18/dev/KIAI/dev/trace/train_table_1209" |
| save_folder = "cell_vis" |
| os.makedirs(save_folder, exist_ok=True) |
| for name in os.listdir(folder_path): |
| if name.endswith("json"): |
| json_file = os.path.join(folder_path, name) |
| with open(json_file, "r") as f: |
| data = json.load(f) |
| |
| img_name = name.split('.')[0] + ".png" |
| print(img_name) |
| image = cv2.imread(os.path.join(folder_path,img_name)) |
| items = data.get('items') |
| for index in range(len(items)): |
|
|
| table = items[index] |
| if table.get('type') != 'table': |
| continue |
| table_x = table.get('x') |
| table_y = table.get('y') |
|
|
| table_prob = table.get('properties') |
| boxes = get_visible_cell_coords(table_prob, table_x, table_y) |
|
|
| for key, box in boxes.items(): |
| row, col = key |
| x, y, w, h = int(box["x"]), int(box["y"]), int(box["width"]), int(box["height"]) |
|
|
| top_left = (x, y) |
| bottom_right = (x + w, y + h) |
|
|
| cv2.rectangle(image, top_left, bottom_right, (0, 255, 0), 2) |
| |
| |
|
|
| |
| cv2.imwrite(os.path.join(save_folder, img_name), image) |
