from collections import OrderedDict, defaultdict import json import argparse import sys import xml.etree.ElementTree as ET import os import random import io import torch from torchvision import transforms from PIL import Image from fitz import Rect import numpy as np import pandas as pd import matplotlib #matplotlib.use('TkAgg') import matplotlib.pyplot as plt import matplotlib.patches as patches from matplotlib.patches import Patch from main import get_model import postprocess sys.path.append("../detr") from models import build_model 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]) ]) def get_class_map(data_type): if data_type == 'structure': class_map = { 'table': 0, 'table column': 1, 'table row': 2, 'table column header': 3, 'table projected row header': 4, 'table spanning cell': 5, 'no object': 6 } elif data_type == 'detection': class_map = {'table': 0, 'table rotated': 1, 'no object': 2} return class_map detection_class_thresholds = { "table": 0.5, "table rotated": 0.5, "no object": 10 } 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 get_args(): parser = argparse.ArgumentParser() parser.add_argument('--image_dir', help="Directory for input images") parser.add_argument('--words_dir', help="Directory for input words") parser.add_argument('--out_dir', help="Output directory") parser.add_argument('--mode', help="The processing to apply to the input image and tokens", choices=['detect', 'recognize', 'extract']) parser.add_argument('--structure_config_path', help="Filepath to the structure model config file") parser.add_argument('--structure_model_path', help="The path to the structure model") parser.add_argument('--detection_config_path', help="Filepath to the detection model config file") parser.add_argument('--detection_model_path', help="The path to the detection model") parser.add_argument('--detection_device', default="cuda") parser.add_argument('--structure_device', default="cuda") parser.add_argument('--crops', '-p', action='store_true', help='Output cropped data from table detections') parser.add_argument('--objects', '-o', action='store_true', help='Output objects') parser.add_argument('--cells', '-l', action='store_true', help='Output cells list') parser.add_argument('--html', '-m', action='store_true', help='Output HTML') parser.add_argument('--csv', '-c', action='store_true', help='Output CSV') parser.add_argument('--verbose', '-v', action='store_true', help='Verbose output') parser.add_argument('--visualize', '-z', action='store_true', help='Visualize output') parser.add_argument('--crop_padding', type=int, default=10, help="The amount of padding to add around a detected table when cropping.") return parser.parse_args() # 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): img_w, img_h = size b = box_cxcywh_to_xyxy(out_bbox) b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32) return b def iob(bbox1, bbox2): """ Compute the intersection area over box area, for bbox1. """ intersection = Rect(bbox1).intersect(bbox2) bbox1_area = Rect(bbox1).get_area() if bbox1_area > 0: return intersection.get_area() / bbox1_area return 0 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['column 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['column 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 refine_table_structure(table_structure, class_thresholds): """ Apply operations to the detected table structure objects such as thresholding, NMS, and alignment. """ rows = table_structure["rows"] columns = table_structure['columns'] # Process the headers column_headers = table_structure['column headers'] column_headers = postprocess.apply_threshold(column_headers, class_thresholds["table column header"]) column_headers = postprocess.nms(column_headers) column_headers = align_headers(column_headers, rows) # Process spanning cells spanning_cells = [elem for elem in table_structure['spanning cells'] if not elem['projected row header']] projected_row_headers = [elem for elem in table_structure['spanning cells'] if elem['projected row header']] spanning_cells = postprocess.apply_threshold(spanning_cells, class_thresholds["table spanning cell"]) projected_row_headers = postprocess.apply_threshold(projected_row_headers, class_thresholds["table projected row header"]) spanning_cells += projected_row_headers # Align before NMS for spanning cells because alignment brings them into agreement # with rows and columns first; if spanning cells still overlap after this operation, # the threshold for NMS can basically be lowered to just above 0 spanning_cells = postprocess.align_supercells(spanning_cells, rows, columns) spanning_cells = postprocess.nms_supercells(spanning_cells) postprocess.header_supercell_tree(spanning_cells) table_structure['columns'] = columns table_structure['rows'] = rows table_structure['spanning cells'] = spanning_cells table_structure['column headers'] = column_headers return table_structure def outputs_to_objects(outputs, img_size, class_idx2name): m = outputs['pred_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 = class_idx2name[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 objects_to_crops(img, tokens, objects, class_thresholds, padding=10): """ Process the bounding boxes produced by the table detection model into cropped table images and cropped tokens. """ table_crops = [] for obj in objects: if obj['score'] < class_thresholds[obj['label']]: continue cropped_table = {} bbox = obj['bbox'] bbox = [bbox[0]-padding, bbox[1]-padding, bbox[2]+padding, bbox[3]+padding] cropped_img = img.crop(bbox) table_tokens = [token for token in tokens if iob(token['bbox'], bbox) >= 0.5] for token in table_tokens: token['bbox'] = [token['bbox'][0]-bbox[0], token['bbox'][1]-bbox[1], token['bbox'][2]-bbox[0], token['bbox'][3]-bbox[1]] # If table is predicted to be rotated, rotate cropped image and tokens/words: if obj['label'] == 'table rotated': cropped_img = cropped_img.rotate(270, expand=True) for token in table_tokens: bbox = token['bbox'] bbox = [cropped_img.size[0]-bbox[3]-1, bbox[0], cropped_img.size[0]-bbox[1]-1, bbox[2]] token['bbox'] = bbox cropped_table['image'] = cropped_img cropped_table['tokens'] = table_tokens table_crops.append(cropped_table) return table_crops def objects_to_structures(objects, tokens, class_thresholds): """ Process the bounding boxes produced by the table structure recognition model into a *consistent* set of table structures (rows, columns, spanning cells, headers). This entails resolving conflicts/overlaps, and ensuring the boxes meet certain alignment conditions (for example: rows should all have the same width, etc.). """ tables = [obj for obj in objects if obj['label'] == 'table'] table_structures = [] for table in tables: table_objects = [obj for obj in objects if iob(obj['bbox'], table['bbox']) >= 0.5] table_tokens = [token for token in tokens if iob(token['bbox'], table['bbox']) >= 0.5] structure = {} columns = [obj for obj in table_objects if obj['label'] == 'table column'] rows = [obj for obj in table_objects if obj['label'] == 'table row'] column_headers = [obj for obj in table_objects if obj['label'] == 'table column header'] spanning_cells = [obj for obj in table_objects if obj['label'] == 'table spanning cell'] for obj in spanning_cells: obj['projected row header'] = False projected_row_headers = [obj for obj in table_objects if obj['label'] == 'table projected row header'] for obj in projected_row_headers: obj['projected row header'] = True spanning_cells += projected_row_headers for obj in rows: obj['column header'] = False for header_obj in column_headers: if iob(obj['bbox'], header_obj['bbox']) >= 0.5: obj['column header'] = True # Refine table structures rows = postprocess.refine_rows(rows, table_tokens, class_thresholds['table row']) columns = postprocess.refine_columns(columns, table_tokens, 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['row_column_bbox'] = [column_rect[0], row_rect[1], column_rect[2], row_rect[3]] table['bbox'] = table['row_column_bbox'] # Process the rows and columns into a complete segmented table columns = postprocess.align_columns(columns, table['row_column_bbox']) rows = postprocess.align_rows(rows, table['row_column_bbox']) structure['rows'] = rows structure['columns'] = columns structure['column headers'] = column_headers structure['spanning cells'] = spanning_cells if len(rows) > 0 and len(columns) > 1: structure = refine_table_structure(structure, class_thresholds) table_structures.append(structure) return table_structures def structure_to_cells(table_structure, tokens): """ Assuming the row, column, spanning cell, 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_structure['columns'] rows = table_structure['rows'] spanning_cells = table_structure['spanning cells'] 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 = 'column header' in row and row['column header'] cell = {'bbox': list(cell_rect), 'column_nums': [column_num], 'row_nums': [row_num], 'column header': header} cell['subcell'] = False for spanning_cell in spanning_cells: spanning_cell_rect = Rect(list(spanning_cell['bbox'])) if (spanning_cell_rect.intersect(cell_rect).get_area() / cell_rect.get_area()) > 0.5: 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['projected row header'] = False cells.append(cell) for spanning_cell in spanning_cells: spanning_cell_rect = Rect(list(spanning_cell['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() if (subcell_rect.intersect(spanning_cell_rect).get_area() / 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 spanning cell to be classified as a header cell; # otherwise, this could lead to a non-rectangular header region header = header and 'column header' in subcell and subcell['column 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), 'column header': header, 'projected row header': spanning_cell['projected row header']} 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 = postprocess.slot_into_containers(cells, tokens) 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, _, _ = postprocess.slot_into_containers(cells, tokens, overlap_threshold=0.001, unique_assignment=True, forced_assignment=False) for cell, cell_span_nums in zip(cells, span_nums_by_cell): cell_spans = [tokens[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'] = postprocess.extract_text_from_spans(cell_spans, remove_integer_superscripts=False) cell['spans'] = cell_spans # Adjust the row, column, and cell bounding boxes to reflect the extracted text num_rows = len(rows) rows = postprocess.sort_objects_top_to_bottom(rows) num_columns = len(columns) columns = postprocess.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: cell['bbox'] = list(cell_rect) pass return cells, confidence_score def cells_to_csv(cells): if len(cells) > 0: num_columns = max([max(cell['column_nums']) for cell in cells]) + 1 num_rows = max([max(cell['row_nums']) for cell in cells]) + 1 else: return header_cells = [cell for cell in cells if cell['column header']] if len(header_cells) > 0: max_header_row = max([max(cell['row_nums']) for cell in header_cells]) else: max_header_row = -1 table_array = np.empty([num_rows, num_columns], dtype="object") if len(cells) > 0: for cell in cells: for row_num in cell['row_nums']: for column_num in cell['column_nums']: table_array[row_num, column_num] = cell["cell text"] header = table_array[:max_header_row+1,:] flattened_header = [] for col in header.transpose(): flattened_header.append(' | '.join(OrderedDict.fromkeys(col))) df = pd.DataFrame(table_array[max_header_row+1:,:], index=None, columns=flattened_header) return df.to_csv(index=None) def cells_to_html(cells): cells = sorted(cells, key=lambda k: min(k['column_nums'])) cells = sorted(cells, key=lambda k: min(k['row_nums'])) table = ET.Element("table") current_row = -1 for cell in cells: this_row = min(cell['row_nums']) attrib = {} colspan = len(cell['column_nums']) if colspan > 1: attrib['colspan'] = str(colspan) rowspan = len(cell['row_nums']) if rowspan > 1: attrib['rowspan'] = str(rowspan) if this_row > current_row: current_row = this_row if cell['column header']: cell_tag = "th" row = ET.SubElement(table, "thead") else: cell_tag = "td" row = ET.SubElement(table, "tr") tcell = ET.SubElement(row, cell_tag, attrib=attrib) tcell.text = cell['cell text'] return str(ET.tostring(table, encoding="unicode", short_empty_elements=False)) def visualize_detected_tables(img, det_tables, out_path): plt.imshow(img, interpolation="lanczos") plt.gcf().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') plt.savefig(out_path, bbox_inches='tight', dpi=150) plt.close() return def visualize_cells(img, cells, out_path): plt.imshow(img, interpolation="lanczos") plt.gcf().set_size_inches(20, 20) ax = plt.gca() for cell in cells: bbox = cell['bbox'] if cell['column header']: facecolor = (1, 0, 0.45) edgecolor = (1, 0, 0.45) alpha = 0.3 linewidth = 2 hatch='//////' elif cell['projected row header']: facecolor = (0.95, 0.6, 0.1) edgecolor = (0.95, 0.6, 0.1) alpha = 0.3 linewidth = 2 hatch='//////' else: facecolor = (0.3, 0.74, 0.8) edgecolor = (0.3, 0.7, 0.6) alpha = 0.3 linewidth = 2 hatch='\\\\\\\\\\\\' 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=(0.3, 0.74, 0.8), edgecolor=(0.3, 0.7, 0.6), label='Data cell', hatch='\\\\\\\\\\\\', alpha=0.3), Patch(facecolor=(1, 0, 0.45), edgecolor=(1, 0, 0.45), label='Column header cell', hatch='//////', alpha=0.3), Patch(facecolor=(0.95, 0.6, 0.1), edgecolor=(0.95, 0.6, 0.1), label='Projected row header cell', hatch='//////', alpha=0.3)] plt.legend(handles=legend_elements, bbox_to_anchor=(0.5, -0.02), loc='upper center', borderaxespad=0, fontsize=10, ncol=3) plt.gcf().set_size_inches(10, 10) plt.axis('off') plt.savefig(out_path, bbox_inches='tight', dpi=150) plt.close() return class TableExtractionPipeline(object): def __init__(self, det_device=None, str_device=None, det_model=None, str_model=None, det_model_path=None, str_model_path=None, det_config_path=None, str_config_path=None): self.det_device = det_device self.str_device = str_device self.det_class_name2idx = get_class_map('detection') self.det_class_idx2name = {v:k for k, v in self.det_class_name2idx.items()} self.det_class_thresholds = detection_class_thresholds self.str_class_name2idx = get_class_map('structure') self.str_class_idx2name = {v:k for k, v in self.str_class_name2idx.items()} self.str_class_thresholds = structure_class_thresholds if not det_config_path is None: with open(det_config_path, 'r') as f: det_config = json.load(f) det_args = type('Args', (object,), det_config) det_args.device = det_device self.det_model, _, _ = build_model(det_args) print("Detection model initialized.") if not det_model_path is None: self.det_model.load_state_dict(torch.load(det_model_path, map_location=torch.device(det_device))) self.det_model.to(det_device) self.det_model.eval() print("Detection model weights loaded.") else: self.det_model = None if not str_config_path is None: with open(str_config_path, 'r') as f: str_config = json.load(f) str_args = type('Args', (object,), str_config) str_args.device = str_device self.str_model, _, _ = build_model(str_args) print("Structure model initialized.") if not str_model_path is None: self.str_model.load_state_dict(torch.load(str_model_path, map_location=torch.device(str_device))) self.str_model.to(str_device) self.str_model.eval() print("Structure model weights loaded.") else: self.str_model = None def __call__(self, page_image, page_tokens=None): return self.extract(self, page_image, page_tokens) def detect(self, img, tokens=None, out_objects=True, out_crops=False, crop_padding=10): out_formats = {} if self.det_model is None: print("No detection model loaded.") return out_formats # Transform the image how the model expects it img_tensor = detection_transform(img) # Run input image through the model outputs = self.det_model([img_tensor.to(self.det_device)]) # Post-process detected objects, assign class labels objects = outputs_to_objects(outputs, img.size, self.det_class_idx2name) if out_objects: out_formats['objects'] = objects if not out_crops: return out_formats # Crop image and tokens for detected table if out_crops: tables_crops = objects_to_crops(img, tokens, objects, self.det_class_thresholds, padding=crop_padding) out_formats['crops'] = tables_crops return out_formats def recognize(self, img, tokens=None, out_objects=False, out_cells=False, out_html=False, out_csv=False): out_formats = {} if self.str_model is None: print("No structure model loaded.") return out_formats if not (out_objects or out_cells or out_html or out_csv): print("No output format specified") return out_formats # Transform the image how the model expects it img_tensor = structure_transform(img) # Run input image through the model outputs = self.str_model([img_tensor.to(self.str_device)]) # Post-process detected objects, assign class labels objects = outputs_to_objects(outputs, img.size, self.str_class_idx2name) if out_objects: out_formats['objects'] = objects if not (out_cells or out_html or out_csv): return out_formats # Further process the detected objects so they correspond to a consistent table tables_structure = objects_to_structures(objects, tokens, self.str_class_thresholds) # Enumerate all table cells: grid cells and spanning cells tables_cells = [structure_to_cells(structure, tokens)[0] for structure in tables_structure] if out_cells: out_formats['cells'] = tables_cells if not (out_html or out_csv): return out_formats # Convert cells to HTML if out_html: tables_htmls = [cells_to_html(cells) for cells in tables_cells] out_formats['html'] = tables_htmls # Convert cells to CSV, including flattening multi-row column headers to a single row if out_csv: tables_csvs = [cells_to_csv(cells) for cells in tables_cells] out_formats['csv'] = tables_csvs return out_formats def extract(self, img, tokens=None, out_objects=True, out_crops=False, out_cells=False, out_html=False, out_csv=False, crop_padding=10): detect_out = self.detect(img, tokens=tokens, out_objects=False, out_crops=True, crop_padding=crop_padding) cropped_tables = detect_out['crops'] extracted_tables = [] for table in cropped_tables: img = table['image'] tokens = table['tokens'] extracted_table = self.recognize(img, tokens=tokens, out_objects=out_objects, out_cells=out_cells, out_html=out_html, out_csv=out_csv) extracted_table['image'] = img extracted_table['tokens'] = tokens extracted_tables.append(extracted_table) return extracted_tables def output_result(key, val, args, img, img_file): if key == 'objects': if args.verbose: print(val) out_file = img_file.replace(".jpg", "_objects.json") with open(os.path.join(args.out_dir, out_file), 'w') as f: json.dump(val, f) if args.visualize: out_file = img_file.replace(".jpg", "_fig_tables.jpg") out_path = os.path.join(args.out_dir, out_file) visualize_detected_tables(img, val, out_path) elif not key == 'image' and not key == 'tokens': for idx, elem in enumerate(val): if key == 'crops': for idx, cropped_table in enumerate(val): out_img_file = img_file.replace(".jpg", "_table_{}.jpg".format(idx)) cropped_table['image'].save(os.path.join(args.out_dir, out_img_file)) out_words_file = out_img_file.replace(".jpg", "_words.json") with open(os.path.join(args.out_dir, out_words_file), 'w') as f: json.dump(cropped_table['tokens'], f) elif key == 'cells': out_file = img_file.replace(".jpg", "_{}_objects.json".format(idx)) with open(os.path.join(args.out_dir, out_file), 'w') as f: json.dump(elem, f) if args.verbose: print(elem) if args.visualize: out_file = img_file.replace(".jpg", "_fig_cells.jpg") out_path = os.path.join(args.out_dir, out_file) visualize_cells(img, elem, out_path) else: out_file = img_file.replace(".jpg", "_{}.{}".format(idx, key)) with open(os.path.join(args.out_dir, out_file), 'w') as f: f.write(elem) if args.verbose: print(elem) def main(): args = get_args() print(args.__dict__) print('-' * 100) if not args.out_dir is None and not os.path.exists(args.out_dir): os.makedirs(args.out_dir) # Create inference pipeline print("Creating inference pipeline") pipe = TableExtractionPipeline(det_device=args.detection_device, str_device=args.structure_device, det_config_path=args.detection_config_path, det_model_path=args.detection_model_path, str_config_path=args.structure_config_path, str_model_path=args.structure_model_path) # Load images img_files = os.listdir(args.image_dir) num_files = len(img_files) random.shuffle(img_files) for count, img_file in enumerate(img_files): print("({}/{})".format(count+1, num_files)) img_path = os.path.join(args.image_dir, img_file) img = Image.open(img_path) print("Image loaded.") if not args.words_dir is None: tokens_path = os.path.join(args.words_dir, img_file.replace(".jpg", "_words.json")) with open(tokens_path, 'r') as f: tokens = json.load(f) # Handle dictionary format if type(tokens) is dict and 'words' in tokens: tokens = tokens['words'] # 'tokens' is a list of tokens # Need to be in a relative reading order # If no order is provided, use current order for idx, token in enumerate(tokens): if not 'span_num' in token: token['span_num'] = idx if not 'line_num' in token: token['line_num'] = 0 if not 'block_num' in token: token['block_num'] = 0 else: tokens = [] if args.mode == 'recognize': extracted_table = pipe.recognize(img, tokens, out_objects=args.objects, out_cells=args.csv, out_html=args.html, out_csv=args.csv) print("Table(s) recognized.") for key, val in extracted_table.items(): output_result(key, val, args, img, img_file) if args.mode == 'detect': detected_tables = pipe.detect(img, tokens, out_objects=args.objects, out_crops=args.crops) print("Table(s) detected.") for key, val in detected_tables.items(): output_result(key, val, args, img, img_file) if args.mode == 'extract': extracted_tables = pipe.extract(img, tokens, out_objects=args.objects, out_cells=args.csv, out_html=args.html, out_csv=args.csv, crop_padding=args.crop_padding) print("Table(s) extracted.") for table_idx, extracted_table in enumerate(extracted_tables): for key, val in extracted_table.items(): output_result(key, val, args, extracted_table['image'], img_file.replace('.jpg', '_{}.jpg'.format(table_idx))) if __name__ == "__main__": main()