mugdha99/Table-Structure-Recognition-Demo
0
1"""2Copyright (C) 2021 Microsoft Corporation3"""4from collections import defaultdict5 6from fitz import Rect7 8 9def apply_threshold(objects, threshold):10 """11 Filter out objects below a certain score.12 """13 return [obj for obj in objects if obj['score'] >= threshold]14 15 16def apply_class_thresholds(bboxes, labels, scores, class_names, class_thresholds):17 """18 Filter out bounding boxes whose confidence is below the confidence threshold for19 its associated class label. 20 """21 # Apply class-specific thresholds22 indices_above_threshold = [idx for idx, (score, label) in enumerate(zip(scores, labels))23 if score >= class_thresholds[24 class_names[label]25 ]26 ]27 bboxes = [bboxes[idx] for idx in indices_above_threshold]28 scores = [scores[idx] for idx in indices_above_threshold]29 labels = [labels[idx] for idx in indices_above_threshold]30 31 return bboxes, scores, labels32 33 34def iou(bbox1, bbox2):35 """36 Compute the intersection-over-union of two bounding boxes.37 """38 intersection = Rect(bbox1).intersect(bbox2)39 union = Rect(bbox1).include_rect(bbox2)40 41 union_area = union.get_area() # getArea()42 if union_area > 0:43 return intersection.get_area() / union.get_area() # .getArea()44 45 return 046 47 48def iob(bbox1, bbox2):49 """50 Compute the intersection area over box area, for bbox1.51 """52 intersection = Rect(bbox1).intersect(bbox2)53 54 bbox1_area = Rect(bbox1).get_area() # .getArea()55 if bbox1_area > 0:56 return intersection.get_area() / bbox1_area # getArea()57 58 return 059 60 61def objects_to_cells(table, objects_in_table, tokens_in_table, class_map, class_thresholds):62 """63 Process the bounding boxes produced by the table structure recognition model64 and the token/word/span bounding boxes into table cells.65 66 Also return a confidence score based on how well the text was able to be67 uniquely slotted into the cells detected by the table model.68 """69 70 table_structures = objects_to_table_structures(table, objects_in_table, tokens_in_table, class_map,71 class_thresholds)72 73 # Check for a valid table74 if len(table_structures['columns']) < 1 or len(table_structures['rows']) < 1:75 cells = []#None76 confidence_score = 077 else:78 cells, confidence_score = table_structure_to_cells(table_structures, tokens_in_table, table['bbox'])79 80 return table_structures, cells, confidence_score81 82 83def objects_to_table_structures(table_object, objects_in_table, tokens_in_table, class_names, class_thresholds):84 """85 Process the bounding boxes produced by the table structure recognition model into86 a *consistent* set of table structures (rows, columns, supercells, headers).87 This entails resolving conflicts/overlaps, and ensuring the boxes meet certain alignment88 conditions (for example: rows should all have the same width, etc.).89 """90 91 page_num = table_object['page_num']92 93 table_structures = {}94 95 columns = [obj for obj in objects_in_table if class_names[obj['label']] == 'table column']96 rows = [obj for obj in objects_in_table if class_names[obj['label']] == 'table row']97 headers = [obj for obj in objects_in_table if class_names[obj['label']] == 'table column header']98 supercells = [obj for obj in objects_in_table if class_names[obj['label']] == 'table spanning cell']99 for obj in supercells:100 obj['subheader'] = False101 subheaders = [obj for obj in objects_in_table if class_names[obj['label']] == 'table projected row header']102 for obj in subheaders:103 obj['subheader'] = True104 supercells += subheaders105 for obj in rows:106 obj['header'] = False107 for header_obj in headers:108 if iob(obj['bbox'], header_obj['bbox']) >= 0.5:109 obj['header'] = True110 111 for row in rows:112 row['page'] = page_num113 114 for column in columns:115 column['page'] = page_num116 117 #Refine table structures118 rows = refine_rows(rows, tokens_in_table, class_thresholds['table row'])119 columns = refine_columns(columns, tokens_in_table, class_thresholds['table column'])120 121 # Shrink table bbox to just the total height of the rows122 # and the total width of the columns123 row_rect = Rect()124 for obj in rows:125 row_rect.include_rect(obj['bbox'])126 column_rect = Rect() 127 for obj in columns:128 column_rect.include_rect(obj['bbox'])129 table_object['row_column_bbox'] = [column_rect[0], row_rect[1], column_rect[2], row_rect[3]]130 table_object['bbox'] = table_object['row_column_bbox']131 132 # Process the rows and columns into a complete segmented table133 columns = align_columns(columns, table_object['row_column_bbox'])134 rows = align_rows(rows, table_object['row_column_bbox'])135 136 table_structures['rows'] = rows137 table_structures['columns'] = columns138 table_structures['headers'] = headers139 table_structures['supercells'] = supercells140 141 if len(rows) > 0 and len(columns) > 1:142 table_structures = refine_table_structures(table_object['bbox'], table_structures, tokens_in_table, class_thresholds)143 144 return table_structures145 146 147def refine_rows(rows, page_spans, score_threshold):148 """149 Apply operations to the detected rows, such as150 thresholding, NMS, and alignment.151 """152 153 rows = nms_by_containment(rows, page_spans, overlap_threshold=0.5)154 # remove_objects_without_content(page_spans, rows) # TODO155 if len(rows) > 1:156 rows = sort_objects_top_to_bottom(rows)157 158 return rows159 160 161def refine_columns(columns, page_spans, score_threshold):162 """163 Apply operations to the detected columns, such as164 thresholding, NMS, and alignment.165 """166 167 columns = nms_by_containment(columns, page_spans, overlap_threshold=0.5)168 # remove_objects_without_content(page_spans, columns) # TODO169 if len(columns) > 1:170 columns = sort_objects_left_to_right(columns)171 172 return columns173 174 175def nms_by_containment(container_objects, package_objects, overlap_threshold=0.5):176 """177 Non-maxima suppression (NMS) of objects based on shared containment of other objects.178 """179 container_objects = sort_objects_by_score(container_objects)180 num_objects = len(container_objects)181 suppression = [False for obj in container_objects]182 183 packages_by_container, _, _ = slot_into_containers(container_objects, package_objects, overlap_threshold=overlap_threshold,184 unique_assignment=True, forced_assignment=False)185 186 for object2_num in range(1, num_objects):187 object2_packages = set(packages_by_container[object2_num])188 if len(object2_packages) == 0:189 suppression[object2_num] = True190 for object1_num in range(object2_num):191 if not suppression[object1_num]:192 object1_packages = set(packages_by_container[object1_num])193 if len(object2_packages.intersection(object1_packages)) > 0:194 suppression[object2_num] = True195 196 final_objects = [obj for idx, obj in enumerate(container_objects) if not suppression[idx]]197 return final_objects198 199 200def slot_into_containers(container_objects, package_objects, overlap_threshold=0.5,201 unique_assignment=True, forced_assignment=False):202 """203 Slot a collection of objects into the container they occupy most (the container which holds the largest fraction of the object).204 """205 best_match_scores = []206 207 container_assignments = [[] for container in container_objects]208 package_assignments = [[] for package in package_objects]209 210 if len(container_objects) == 0 or len(package_objects) == 0:211 return container_assignments, package_assignments, best_match_scores212 213 match_scores = defaultdict(dict)214 for package_num, package in enumerate(package_objects):215 match_scores = []216 package_rect = Rect(package['bbox'])217 package_area = package_rect.get_area() # getArea() 218 for container_num, container in enumerate(container_objects):219 container_rect = Rect(container['bbox'])220 intersect_area = container_rect.intersect(package['bbox']).get_area() # getArea()221 overlap_fraction = intersect_area / package_area222 match_scores.append({'container': container, 'container_num': container_num, 'score': overlap_fraction})223 224 sorted_match_scores = sort_objects_by_score(match_scores)225 226 best_match_score = sorted_match_scores[0]227 best_match_scores.append(best_match_score['score'])228 if forced_assignment or best_match_score['score'] >= overlap_threshold:229 container_assignments[best_match_score['container_num']].append(package_num)230 package_assignments[package_num].append(best_match_score['container_num'])231 232 if not unique_assignment: # slot package into all eligible slots233 for match_score in sorted_match_scores[1:]:234 if match_score['score'] >= overlap_threshold:235 container_assignments[match_score['container_num']].append(package_num)236 package_assignments[package_num].append(match_score['container_num'])237 else:238 break239 240 return container_assignments, package_assignments, best_match_scores241 242 243def sort_objects_by_score(objects, reverse=True):244 """245 Put any set of objects in order from high score to low score.246 """247 if reverse:248 sign = -1249 else:250 sign = 1251 return sorted(objects, key=lambda k: sign*k['score'])252 253 254def remove_objects_without_content(page_spans, objects):255 """256 Remove any objects (these can be rows, columns, supercells, etc.) that don't257 have any text associated with them.258 """259 for obj in objects[:]:260 object_text, _ = extract_text_inside_bbox(page_spans, obj['bbox'])261 if len(object_text.strip()) == 0:262 objects.remove(obj)263 264 265def extract_text_inside_bbox(spans, bbox):266 """267 Extract the text inside a bounding box.268 """269 bbox_spans = get_bbox_span_subset(spans, bbox)270 bbox_text = extract_text_from_spans(bbox_spans, remove_integer_superscripts=True)271 272 return bbox_text, bbox_spans273 274 275def get_bbox_span_subset(spans, bbox, threshold=0.5):276 """277 Reduce the set of spans to those that fall within a bounding box.278 279 threshold: the fraction of the span that must overlap with the bbox.280 """281 span_subset = []282 for span in spans:283 if overlaps(span['bbox'], bbox, threshold):284 span_subset.append(span)285 return span_subset286 287 288def overlaps(bbox1, bbox2, threshold=0.5):289 """290 Test if more than "threshold" fraction of bbox1 overlaps with bbox2.291 """292 rect1 = Rect(list(bbox1))293 area1 = rect1.get_area() # .getArea()294 if area1 == 0:295 return False296 return rect1.intersect(list(bbox2)).get_area()/area1 >= threshold # getArea()297 298 299def extract_text_from_spans(spans, join_with_space=True, remove_integer_superscripts=True):300 """301 Convert a collection of page tokens/words/spans into a single text string.302 """303 304 if join_with_space:305 join_char = " "306 else:307 join_char = ""308 spans_copy = spans[:]309 310 if remove_integer_superscripts:311 for span in spans:312 flags = span['flags']313 if flags & 2**0: # superscript flag314 if is_int(span['text']):315 spans_copy.remove(span)316 else:317 span['superscript'] = True318 319 if len(spans_copy) == 0:320 return ""321 322 spans_copy.sort(key=lambda span: span['span_num'])323 spans_copy.sort(key=lambda span: span['line_num'])324 spans_copy.sort(key=lambda span: span['block_num'])325 326 # Force the span at the end of every line within a block to have exactly one space327 # unless the line ends with a space or ends with a non-space followed by a hyphen328 line_texts = []329 line_span_texts = [spans_copy[0]['text']]330 for span1, span2 in zip(spans_copy[:-1], spans_copy[1:]):331 if not span1['block_num'] == span2['block_num'] or not span1['line_num'] == span2['line_num']:332 line_text = join_char.join(line_span_texts).strip()333 if (len(line_text) > 0334 and not line_text[-1] == ' '335 and not (len(line_text) > 1 and line_text[-1] == "-" and not line_text[-2] == ' ')):336 if not join_with_space:337 line_text += ' '338 line_texts.append(line_text)339 line_span_texts = [span2['text']]340 else:341 line_span_texts.append(span2['text'])342 line_text = join_char.join(line_span_texts)343 line_texts.append(line_text)344 345 return join_char.join(line_texts).strip()346 347 348def sort_objects_left_to_right(objs):349 """350 Put the objects in order from left to right.351 """352 return sorted(objs, key=lambda k: k['bbox'][0] + k['bbox'][2])353 354 355def sort_objects_top_to_bottom(objs):356 """357 Put the objects in order from top to bottom.358 """359 return sorted(objs, key=lambda k: k['bbox'][1] + k['bbox'][3])360 361 362def align_columns(columns, bbox):363 """364 For every column, align the top and bottom boundaries to the final365 table bounding box.366 """367 try:368 for column in columns:369 column['bbox'][1] = bbox[1]370 column['bbox'][3] = bbox[3]371 except Exception as err:372 print("Could not align columns: {}".format(err))373 pass374 375 return columns376 377 378def align_rows(rows, bbox):379 """380 For every row, align the left and right boundaries to the final381 table bounding box.382 """383 try:384 for row in rows:385 row['bbox'][0] = bbox[0]386 row['bbox'][2] = bbox[2]387 except Exception as err:388 print("Could not align rows: {}".format(err))389 pass390 391 return rows392 393 394def refine_table_structures(table_bbox, table_structures, page_spans, class_thresholds):395 """396 Apply operations to the detected table structure objects such as397 thresholding, NMS, and alignment.398 """399 rows = table_structures["rows"]400 columns = table_structures['columns']401 402 #columns = fill_column_gaps(columns, table_bbox)403 #rows = fill_row_gaps(rows, table_bbox)404 405 # Process the headers406 headers = table_structures['headers']407 headers = apply_threshold(headers, class_thresholds["table column header"])408 headers = nms(headers)409 headers = align_headers(headers, rows)410 411 # Process supercells412 supercells = [elem for elem in table_structures['supercells'] if not elem['subheader']]413 subheaders = [elem for elem in table_structures['supercells'] if elem['subheader']]414 supercells = apply_threshold(supercells, class_thresholds["table spanning cell"])415 subheaders = apply_threshold(subheaders, class_thresholds["table projected row header"])416 supercells += subheaders417 # Align before NMS for supercells because alignment brings them into agreement418 # with rows and columns first; if supercells still overlap after this operation,419 # the threshold for NMS can basically be lowered to just above 0420 supercells = align_supercells(supercells, rows, columns)421 supercells = nms_supercells(supercells)422 423 header_supercell_tree(supercells)424 425 table_structures['columns'] = columns426 table_structures['rows'] = rows427 table_structures['supercells'] = supercells428 table_structures['headers'] = headers429 430 return table_structures431 432 433def nms(objects, match_criteria="object2_overlap", match_threshold=0.05, keep_metric="score", keep_higher=True):434 """435 A customizable version of non-maxima suppression (NMS).436 437 Default behavior: If a lower-confidence object overlaps more than 5% of its area438 with a higher-confidence object, remove the lower-confidence object.439 440 objects: set of dicts; each object dict must have a 'bbox' and a 'score' field441 match_criteria: how to measure how much two objects "overlap"442 match_threshold: the cutoff for determining that overlap requires suppression of one object443 keep_metric: which metric to use to determine the object to keep444 keep_higher: if True, keep the object with the higher metric; otherwise, keep the lower445 """446 if len(objects) == 0:447 return []448 449 if keep_metric=="score":450 objects = sort_objects_by_score(objects, reverse=keep_higher)451 elif keep_metric=="area":452 objects = sort_objects_by_area(objects, reverse=keep_higher)453 454 num_objects = len(objects)455 suppression = [False for obj in objects]456 457 for object2_num in range(1, num_objects):458 object2_rect = Rect(objects[object2_num]['bbox'])459 object2_area = object2_rect.get_area() # .getArea()460 for object1_num in range(object2_num):461 if not suppression[object1_num]:462 object1_rect = Rect(objects[object1_num]['bbox'])463 object1_area = object1_rect.get_area() # .getArea()464 intersect_area = object1_rect.intersect(object2_rect).get_area() # .getArea()465 try:466 if match_criteria=="object1_overlap":467 metric = intersect_area / object1_area468 elif match_criteria=="object2_overlap":469 metric = intersect_area / object2_area470 elif match_criteria=="iou":471 metric = intersect_area / (object1_area + object2_area - intersect_area)472 if metric >= match_threshold:473 suppression[object2_num] = True474 break475 except Exception:476 # Intended to recover from divide-by-zero477 pass478 479 return [obj for idx, obj in enumerate(objects) if not suppression[idx]]480 481 482def align_headers(headers, rows):483 """484 Adjust the header boundary to be the convex hull of the rows it intersects485 at least 50% of the height of.486 487 For now, we are not supporting tables with multiple headers, so we need to488 eliminate anything besides the top-most header.489 """490 491 aligned_headers = []492 493 for row in rows:494 row['header'] = False495 496 header_row_nums = []497 for header in headers:498 for row_num, row in enumerate(rows):499 row_height = row['bbox'][3] - row['bbox'][1]500 min_row_overlap = max(row['bbox'][1], header['bbox'][1])501 max_row_overlap = min(row['bbox'][3], header['bbox'][3])502 overlap_height = max_row_overlap - min_row_overlap503 if overlap_height / row_height >= 0.5:504 header_row_nums.append(row_num)505 506 if len(header_row_nums) == 0:507 return aligned_headers508 509 header_rect = Rect()510 if header_row_nums[0] > 0:511 header_row_nums = list(range(header_row_nums[0]+1)) + header_row_nums512 513 last_row_num = -1514 for row_num in header_row_nums:515 if row_num == last_row_num + 1:516 row = rows[row_num]517 row['header'] = True518 header_rect = header_rect.include_rect(row['bbox'])519 last_row_num = row_num520 else:521 # Break as soon as a non-header row is encountered.522 # This ignores any subsequent rows in the table labeled as a header.523 # Having more than 1 header is not supported currently.524 break525 526 header = {'bbox': list(header_rect)}527 aligned_headers.append(header)528 529 return aligned_headers530 531 532def align_supercells(supercells, rows, columns):533 """534 For each supercell, align it to the rows it intersects 50% of the height of,535 and the columns it intersects 50% of the width of.536 Eliminate supercells for which there are no rows and columns it intersects 50% with.537 """538 aligned_supercells = []539 540 for supercell in supercells:541 supercell['header'] = False542 row_bbox_rect = None543 col_bbox_rect = None544 intersecting_header_rows = set()545 intersecting_data_rows = set()546 for row_num, row in enumerate(rows):547 row_height = row['bbox'][3] - row['bbox'][1]548 supercell_height = supercell['bbox'][3] - supercell['bbox'][1]549 min_row_overlap = max(row['bbox'][1], supercell['bbox'][1])550 max_row_overlap = min(row['bbox'][3], supercell['bbox'][3])551 overlap_height = max_row_overlap - min_row_overlap552 if 'span' in supercell:553 overlap_fraction = max(overlap_height/row_height,554 overlap_height/supercell_height)555 else:556 overlap_fraction = overlap_height / row_height557 if overlap_fraction >= 0.5:558 if 'header' in row and row['header']:559 intersecting_header_rows.add(row_num)560 else:561 intersecting_data_rows.add(row_num)562 563 # Supercell cannot span across the header boundary; eliminate whichever564 # group of rows is the smallest565 supercell['header'] = False566 if len(intersecting_data_rows) > 0 and len(intersecting_header_rows) > 0:567 if len(intersecting_data_rows) > len(intersecting_header_rows):568 intersecting_header_rows = set()569 else:570 intersecting_data_rows = set()571 if len(intersecting_header_rows) > 0:572 supercell['header'] = True573 elif 'span' in supercell:574 continue # Require span supercell to be in the header575 intersecting_rows = intersecting_data_rows.union(intersecting_header_rows)576 # Determine vertical span of aligned supercell577 for row_num in intersecting_rows:578 if row_bbox_rect is None:579 row_bbox_rect = Rect(rows[row_num]['bbox'])580 else:581 row_bbox_rect = row_bbox_rect.include_rect(rows[row_num]['bbox'])582 if row_bbox_rect is None:583 continue584 585 intersecting_cols = []586 for col_num, col in enumerate(columns):587 col_width = col['bbox'][2] - col['bbox'][0]588 supercell_width = supercell['bbox'][2] - supercell['bbox'][0]589 min_col_overlap = max(col['bbox'][0], supercell['bbox'][0])590 max_col_overlap = min(col['bbox'][2], supercell['bbox'][2])591 overlap_width = max_col_overlap - min_col_overlap592 if 'span' in supercell:593 overlap_fraction = max(overlap_width/col_width,594 overlap_width/supercell_width)595 # Multiply by 2 effectively lowers the threshold to 0.25596 if supercell['header']:597 overlap_fraction = overlap_fraction * 2598 else:599 overlap_fraction = overlap_width / col_width600 if overlap_fraction >= 0.5:601 intersecting_cols.append(col_num)602 if col_bbox_rect is None:603 col_bbox_rect = Rect(col['bbox'])604 else:605 col_bbox_rect = col_bbox_rect.include_rect(col['bbox'])606 if col_bbox_rect is None:607 continue608 609 supercell_bbox = list(row_bbox_rect.intersect(col_bbox_rect))610 supercell['bbox'] = supercell_bbox611 612 # Only a true supercell if it joins across multiple rows or columns613 if (len(intersecting_rows) > 0 and len(intersecting_cols) > 0614 and (len(intersecting_rows) > 1 or len(intersecting_cols) > 1)):615 supercell['row_numbers'] = list(intersecting_rows)616 supercell['column_numbers'] = intersecting_cols617 aligned_supercells.append(supercell)618 619 # A span supercell in the header means there must be supercells above it in the header620 if 'span' in supercell and supercell['header'] and len(supercell['column_numbers']) > 1:621 for row_num in range(0, min(supercell['row_numbers'])):622 new_supercell = {'row_numbers': [row_num], 'column_numbers': supercell['column_numbers'],623 'score': supercell['score'], 'propagated': True}624 new_supercell_columns = [columns[idx] for idx in supercell['column_numbers']]625 new_supercell_rows = [rows[idx] for idx in supercell['row_numbers']]626 bbox = [min([column['bbox'][0] for column in new_supercell_columns]),627 min([row['bbox'][1] for row in new_supercell_rows]),628 max([column['bbox'][2] for column in new_supercell_columns]),629 max([row['bbox'][3] for row in new_supercell_rows])]630 new_supercell['bbox'] = bbox631 aligned_supercells.append(new_supercell)632 633 return aligned_supercells634 635 636def nms_supercells(supercells):637 """638 A NMS scheme for supercells that first attempts to shrink supercells to639 resolve overlap.640 If two supercells overlap the same (sub)cell, shrink the lower confidence641 supercell to resolve the overlap. If shrunk supercell is empty, remove it.642 """643 644 supercells = sort_objects_by_score(supercells)645 num_supercells = len(supercells)646 suppression = [False for supercell in supercells]647 648 for supercell2_num in range(1, num_supercells):649 supercell2 = supercells[supercell2_num]650 for supercell1_num in range(supercell2_num):651 supercell1 = supercells[supercell1_num]652 remove_supercell_overlap(supercell1, supercell2)653 if ((len(supercell2['row_numbers']) < 2 and len(supercell2['column_numbers']) < 2)654 or len(supercell2['row_numbers']) == 0 or len(supercell2['column_numbers']) == 0):655 suppression[supercell2_num] = True656 657 return [obj for idx, obj in enumerate(supercells) if not suppression[idx]]658 659 660def header_supercell_tree(supercells):661 """662 Make sure no supercell in the header is below more than one supercell in any row above it.663 The cells in the header form a tree, but a supercell with more than one supercell in a row664 above it means that some cell has more than one parent, which is not allowed. Eliminate665 any supercell that would cause this to be violated.666 """667 header_supercells = [supercell for supercell in supercells if 'header' in supercell and supercell['header']]668 header_supercells = sort_objects_by_score(header_supercells)669 670 for header_supercell in header_supercells[:]:671 ancestors_by_row = defaultdict(int)672 min_row = min(header_supercell['row_numbers'])673 for header_supercell2 in header_supercells:674 max_row2 = max(header_supercell2['row_numbers'])675 if max_row2 < min_row:676 if (set(header_supercell['column_numbers']).issubset(677 set(header_supercell2['column_numbers']))):678 for row2 in header_supercell2['row_numbers']:679 ancestors_by_row[row2] += 1680 for row in range(0, min_row):681 if not ancestors_by_row[row] == 1:682 supercells.remove(header_supercell)683 break684 685 686def table_structure_to_cells(table_structures, table_spans, table_bbox):687 """688 Assuming the row, column, supercell, and header bounding boxes have689 been refined into a set of consistent table structures, process these690 table structures into table cells. This is a universal representation691 format for the table, which can later be exported to Pandas or CSV formats.692 Classify the cells as header/access cells or data cells693 based on if they intersect with the header bounding box.694 """695 columns = table_structures['columns']696 rows = table_structures['rows']697 supercells = table_structures['supercells']698 cells = []699 subcells = []700 701 # Identify complete cells and subcells702 for column_num, column in enumerate(columns):703 for row_num, row in enumerate(rows):704 column_rect = Rect(list(column['bbox']))705 row_rect = Rect(list(row['bbox']))706 cell_rect = row_rect.intersect(column_rect)707 header = 'header' in row and row['header']708 cell = {'bbox': list(cell_rect), 'column_nums': [column_num], 'row_nums': [row_num],709 'header': header}710 711 cell['subcell'] = False712 for supercell in supercells:713 supercell_rect = Rect(list(supercell['bbox']))714 if (supercell_rect.intersect(cell_rect).get_area() # .getArea()715 / cell_rect.get_area()) > 0.5: # getArea()716 cell['subcell'] = True717 break718 719 if cell['subcell']:720 subcells.append(cell)721 else:722 #cell_text = extract_text_inside_bbox(table_spans, cell['bbox'])723 #cell['cell_text'] = cell_text724 cell['subheader'] = False725 cells.append(cell)726 727 for supercell in supercells:728 supercell_rect = Rect(list(supercell['bbox']))729 cell_columns = set()730 cell_rows = set()731 cell_rect = None732 header = True733 for subcell in subcells:734 subcell_rect = Rect(list(subcell['bbox']))735 subcell_rect_area = subcell_rect.get_area() # .getArea()736 if (subcell_rect.intersect(supercell_rect).get_area() # .getArea()737 / subcell_rect_area) > 0.5:738 if cell_rect is None:739 cell_rect = Rect(list(subcell['bbox']))740 else:741 cell_rect.include_rect(Rect(list(subcell['bbox'])))742 cell_rows = cell_rows.union(set(subcell['row_nums']))743 cell_columns = cell_columns.union(set(subcell['column_nums']))744 # By convention here, all subcells must be classified745 # as header cells for a supercell to be classified as a header cell;746 # otherwise, this could lead to a non-rectangular header region747 header = header and 'header' in subcell and subcell['header']748 if len(cell_rows) > 0 and len(cell_columns) > 0:749 cell = {'bbox': list(cell_rect), 'column_nums': list(cell_columns), 'row_nums': list(cell_rows),750 'header': header, 'subheader': supercell['subheader']}751 cells.append(cell)752 753 # Compute a confidence score based on how well the page tokens754 # slot into the cells reported by the model755 _, _, cell_match_scores = slot_into_containers(cells, table_spans)756 try:757 mean_match_score = sum(cell_match_scores) / len(cell_match_scores)758 min_match_score = min(cell_match_scores)759 confidence_score = (mean_match_score + min_match_score)/2760 except:761 confidence_score = 0762 763 # Dilate rows and columns before final extraction764 #dilated_columns = fill_column_gaps(columns, table_bbox)765 dilated_columns = columns766 #dilated_rows = fill_row_gaps(rows, table_bbox)767 dilated_rows = rows768 for cell in cells:769 column_rect = Rect()770 for column_num in cell['column_nums']:771 column_rect.include_rect(list(dilated_columns[column_num]['bbox']))772 row_rect = Rect()773 for row_num in cell['row_nums']:774 row_rect.include_rect(list(dilated_rows[row_num]['bbox']))775 cell_rect = column_rect.intersect(row_rect)776 cell['bbox'] = list(cell_rect)777 778 span_nums_by_cell, _, _ = slot_into_containers(cells, table_spans, overlap_threshold=0.001,779 unique_assignment=True, forced_assignment=False)780 781 for cell, cell_span_nums in zip(cells, span_nums_by_cell):782 cell_spans = [table_spans[num] for num in cell_span_nums]783 # TODO: Refine how text is extracted; should be character-based, not span-based;784 # but need to associate 785 # cell['cell_text'] = extract_text_from_spans(cell_spans, remove_integer_superscripts=False) # TODO786 cell['spans'] = cell_spans787 788 # Adjust the row, column, and cell bounding boxes to reflect the extracted text789 num_rows = len(rows)790 rows = sort_objects_top_to_bottom(rows)791 num_columns = len(columns)792 columns = sort_objects_left_to_right(columns)793 min_y_values_by_row = defaultdict(list)794 max_y_values_by_row = defaultdict(list)795 min_x_values_by_column = defaultdict(list)796 max_x_values_by_column = defaultdict(list)797 for cell in cells:798 min_row = min(cell["row_nums"])799 max_row = max(cell["row_nums"])800 min_column = min(cell["column_nums"])801 max_column = max(cell["column_nums"])802 for span in cell['spans']:803 min_x_values_by_column[min_column].append(span['bbox'][0])804 min_y_values_by_row[min_row].append(span['bbox'][1])805 max_x_values_by_column[max_column].append(span['bbox'][2])806 max_y_values_by_row[max_row].append(span['bbox'][3])807 for row_num, row in enumerate(rows):808 if len(min_x_values_by_column[0]) > 0:809 row['bbox'][0] = min(min_x_values_by_column[0])810 if len(min_y_values_by_row[row_num]) > 0:811 row['bbox'][1] = min(min_y_values_by_row[row_num])812 if len(max_x_values_by_column[num_columns-1]) > 0:813 row['bbox'][2] = max(max_x_values_by_column[num_columns-1])814 if len(max_y_values_by_row[row_num]) > 0:815 row['bbox'][3] = max(max_y_values_by_row[row_num])816 for column_num, column in enumerate(columns):817 if len(min_x_values_by_column[column_num]) > 0:818 column['bbox'][0] = min(min_x_values_by_column[column_num])819 if len(min_y_values_by_row[0]) > 0:820 column['bbox'][1] = min(min_y_values_by_row[0])821 if len(max_x_values_by_column[column_num]) > 0:822 column['bbox'][2] = max(max_x_values_by_column[column_num])823 if len(max_y_values_by_row[num_rows-1]) > 0:824 column['bbox'][3] = max(max_y_values_by_row[num_rows-1])825 for cell in cells:826 row_rect = Rect()827 column_rect = Rect()828 for row_num in cell['row_nums']:829 row_rect.include_rect(list(rows[row_num]['bbox']))830 for column_num in cell['column_nums']:831 column_rect.include_rect(list(columns[column_num]['bbox']))832 cell_rect = row_rect.intersect(column_rect)833 if cell_rect.get_area() > 0: # getArea()834 cell['bbox'] = list(cell_rect)835 pass836 837 return cells, confidence_score838 839 840def remove_supercell_overlap(supercell1, supercell2):841 """842 This function resolves overlap between supercells (supercells must be843 disjoint) by iteratively shrinking supercells by the fewest grid cells844 necessary to resolve the overlap.845 Example:846 If two supercells overlap at grid cell (R, C), and supercell #1 is less847 confident than supercell #2, we eliminate either row R from supercell #1848 or column C from supercell #1 by comparing the number of columns in row R849 versus the number of rows in column C. If the number of columns in row R850 is less than the number of rows in column C, we eliminate row R from851 supercell #1. This resolves the overlap by removing fewer grid cells from852 supercell #1 than if we eliminated column C from it.853 """854 common_rows = set(supercell1['row_numbers']).intersection(set(supercell2['row_numbers']))855 common_columns = set(supercell1['column_numbers']).intersection(set(supercell2['column_numbers']))856 857 # While the supercells have overlapping grid cells, continue shrinking the less-confident858 # supercell one row or one column at a time859 while len(common_rows) > 0 and len(common_columns) > 0:860 # Try to shrink the supercell as little as possible to remove the overlap;861 # if the supercell has fewer rows than columns, remove an overlapping column,862 # because this removes fewer grid cells from the supercell;863 # otherwise remove an overlapping row864 if len(supercell2['row_numbers']) < len(supercell2['column_numbers']):865 min_column = min(supercell2['column_numbers'])866 max_column = max(supercell2['column_numbers'])867 if max_column in common_columns:868 common_columns.remove(max_column)869 supercell2['column_numbers'].remove(max_column)870 elif min_column in common_columns:871 common_columns.remove(min_column)872 supercell2['column_numbers'].remove(min_column)873 else:874 supercell2['column_numbers'] = []875 common_columns = set()876 else:877 min_row = min(supercell2['row_numbers'])878 max_row = max(supercell2['row_numbers'])879 if max_row in common_rows:880 common_rows.remove(max_row)881 supercell2['row_numbers'].remove(max_row)882 elif min_row in common_rows:883 common_rows.remove(min_row)884 supercell2['row_numbers'].remove(min_row)885 else:886 supercell2['row_numbers'] = []887 common_rows = set()888