process_icdar2013.py

"""
Copyright (C) 2023 Microsoft Corporation

Script to edit, filter, and canonicalize ICDAR-2013 to align it with PubTables-1M.

We still need to verify that this script works correctly.

If you use this code in your published work, we request that you cite our papers
and table-transformer GitHub repo.

As a prerequisite, need to manually correct some of the original ICDAR-2013 samples.
Here is a log of the corrections we made:
 - Edited eu-dataset/eu-032-str.xml to change one of the end-col from 0 to 9
 - Fixed the end-col for two cells in us-gov-dataset/us-006-str.xml
 - Fixed the x1 coordinate for a cell in competition-dataset-us/us-018-str.xml
 - Fixed the end-col for a cell in eu-dataset/eu-023-str.xml, table 0
 - Fixed the end-col for four cells in us-gov-dataset/us-014-str.xml, table 0
 - Fixed the end-col for three cells in us-gov-dataset/us-021-str.xml, table 0
 - Fixed the start-col for a cell in us-gov-dataset/us-012-str.xml, table 0
 - Fixed the end-col for two cells; split the table into two regions/tables in us-gov-dataset/us-020-str.xml, table 0
 - Fixed the end-col for two cells; split the table into two regions/tables in us-gov-dataset/us-020-str.xml, table 1
 - Fixed the end-col for a cell in us-gov-dataset/us-012-str.xml, table 2
 - Fixed the end-col for a cell; fixed the start-row for several cells in us-gov-dataset/us-012-str.xml, table 3
 - Fixed the start-row for many cells in us-gov-dataset/us-025-str.xml, table 0
 - Fixed many start-rows, at least one start-col, and at least one text content in us-gov-dataset/us-019-str.xml, table 0
 - Fixed the y2 coordinate for a cell in us-gov-dataset/us-001-str.xml
 - Fixed the y2 coordinate for a cell in competition-dataset-us/us-035a-str.xml
 - Fixed the start-col for a cell in eu-dataset/eu-014-str.xml
 - Incremented all of the row counts by one (they started at -1) in competition-dataset-us/us-019-str.xml
 - Changed x1 from 550 to 519 for a cell in competition-dataset-us/us-001-str.xml, table 1
 - Edited the XML for us-gov-dataset/us-011-str.xml to make sure it doesn't break the parser (although it didn't break the parser so this could be reverted)
 - 18 total tables manually corrected

 Additional odd or potentially ambiguous cases to note:
- us-gov-dataset/us-008-str.xml
  - Unclear how the second row should be labeled
- competition-dataset-us/us-019-str.xml
  - Possible projected row headers but their alignment is consistent with column headers
- us-gov-dataset/us-016-str.xml
  - The last row also looks like it could be two separate rows
- us-gov-dataset/us-015-str.xml
  - The last row could be split into two rows
 """

import json
import os
from collections import defaultdict
import traceback
import xml.etree.ElementTree as ET
from xml.dom import minidom
import argparse

import fitz
from fitz import Rect
from PIL import Image
import numpy as np
from tqdm import tqdm
import editdistance

def adjust_bbox_coordinates(table_dict, doc):
    # Change bbox coordinates to be relative to PyMuPDF page.rect coordinate space
    mediabox = doc[table_dict['pdf_page_index']].mediabox
    mat = doc[table_dict['pdf_page_index']].transformation_matrix

    for cell in table_dict['cells']:
        if not 'pdf_text_tight_bbox' in cell or len(cell['pdf_text_tight_bbox']) == 0:
            continue
        bbox = list(Rect(cell['pdf_text_tight_bbox']))
        bbox = [bbox[0] + mediabox[0],
                bbox[1] - mediabox[3] + mediabox[2],
                bbox[2] + mediabox[0],
                bbox[3] - mediabox[3] + mediabox[2]]
        cell['pdf_text_tight_bbox'] = bbox

def create_document_page_image(doc, page_num, zoom=None, output_image_max_dim=1000):
    page = doc[page_num]
    
    if zoom is None:
        zoom = output_image_max_dim / max(page.rect)
        
    mat = fitz.Matrix(zoom, zoom)
    pix = page.get_pixmap(matrix = mat, alpha = False)
    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
    
    return img

def create_table_dict(doc, region, partial_table=False):
    table_dict = {}
    table_dict['reject'] = []
    table_dict['fix'] = []
    table_dict['partial_table'] = partial_table
    table_dict['cells'] = []

    page_num = int(region.get("page")) - 1
    page = doc[page_num]

    max_row = -1
    max_column = -1
    cell_presence_by_grid_location = defaultdict(bool)
    for cell in region.iter('cell'):
        table_cell = {}
        bbox = cell.find('bounding-box')
        bbox = [int(bbox.get('x1')),
                page.mediabox[3] - int(bbox.get('y2')) + 1,
                int(bbox.get('x2')),
                page.mediabox[3] - int(bbox.get('y1')) + 1]
        table_cell['pdf_text_tight_bbox'] = bbox
        table_cell['text_content'] = cell.find('content').text

        start_row = int(cell.get('start-row'))
        start_column = int(cell.get('start-col'))

        try:
            end_row = int(cell.get('end-row'))
        except:
            end_row = start_row

        try:
            end_column = int(cell.get('end-col'))
        except:
            end_column = start_column
        if end_column < start_column:
            print("END COLUMN BEFORE START COLUMN")
            end_column = 10

        table_cell['row_nums'] = list(range(start_row, end_row + 1))
        table_cell['column_nums'] = list(range(start_column, end_column + 1))
        table_cell['is_column_header'] = False

        table_dict['cells'].append(table_cell)

        if len(table_cell['column_nums']) == 0 or len(table_cell['row_nums']) == 0:
            print("Empty or negative grid locations: {}".format(table_cell))
            #raise Exception("Empty or negative grid locations: {}".format(table_cell))

        max_row = max(max_row, table_cell['row_nums'][-1])
        max_column = max(max_column, table_cell['column_nums'][-1])
        overlap = False
        for row_num in table_cell['row_nums']:
            for column_num in table_cell['column_nums']:
                if cell_presence_by_grid_location[(row_num, column_num)]:
                    overlap = True
                    print("Bad table, some cells intersect/overlap within the grid: ({}, {})".format(row_num, column_num))
                cell_presence_by_grid_location[(row_num, column_num)] = True
        if overlap:
            table_dict['reject'].append('HTML grid cells overlap')

    # Create missing cells
    for row_num in range(max_row + 1):
        for column_num in range(max_column + 1):
            if not (row_num, column_num) in cell_presence_by_grid_location:
                table_cell = {}
                table_cell['pdf_text_tight_bbox'] = []
                table_cell['text_content'] = ''
                table_cell['row_nums'] = list(range(row_num, row_num + 1))
                table_cell['column_nums'] = list(range(column_num, column_num + 1))
                table_cell['is_column_header'] = False
                
                table_dict['cells'].append(table_cell)

    num_rows = max([max(cell['row_nums']) for cell in table_dict['cells']]) + 1
    num_columns = max([max(cell['column_nums']) for cell in table_dict['cells']]) + 1
    expected_num_cells = num_rows * num_columns
    #actual_num_cells = len(grid_cell_locations)
        
    # Make sure all grid locations are present
    # Could be bad data or bad parsing algorithm
    #if not expected_num_cells == actual_num_cells:
    #    table_dict['reject'].append("HTML missing grid cells")
        
    table_dict['rows'] = {row_num: {'is_column_header': False} for row_num in range(num_rows)}
    table_dict['columns'] = {column_num: {} for column_num in range(num_columns)}
                
    table_dict['pdf_page_index'] = page_num
    return table_dict

def read_xml(xml_file, doc):

    tree = ET.parse(xml_file)
    root = tree.getroot()

    tables = []

    for table_tree in root.iter('table'):
        table_id = table_tree.attrib['id']
        page_nums = [int(region.get("page"))-1 for region in table_tree.iter('region')]
        for region_num, region in enumerate(table_tree.iter('region')):
            region_id = region.attrib['id']
            if region_num > 0 and not page_nums[region_num] == page_nums[region_num-1]:
                partial_table = True
            else:
                partial_table = False
            table_dict = create_table_dict(doc, region, partial_table=partial_table)
            table_dict['icdar2013_source_table_id'] = table_id
            table_dict['icdar2013_source_region_id'] = region_id
            tables.append(table_dict)

    return tables

def complete_table_grid(table_dict):
    rects_by_row = defaultdict(lambda: [None, None, None, None])
    rects_by_column = defaultdict(lambda: [None, None, None, None])
    table_rect = Rect()

    # Determine bounding box for rows and columns
    for cell in table_dict['cells']:
        if not 'pdf_text_tight_bbox' in cell or len(cell['pdf_text_tight_bbox']) == 0:
            continue

        bbox = cell['pdf_text_tight_bbox'] 

        table_rect.include_rect(list(bbox))
        
        min_row = min(cell['row_nums'])
        if rects_by_row[min_row][1] is None:
            rects_by_row[min_row][1] = bbox[1]
        else:
            rects_by_row[min_row][1] = min(rects_by_row[min_row][1], bbox[1])
            
        max_row = max(cell['row_nums'])
        if rects_by_row[max_row][3] is None:
            rects_by_row[max_row][3] = bbox[3]
        else:
            rects_by_row[max_row][3] = max(rects_by_row[max_row][3], bbox[3])
            
        min_column = min(cell['column_nums'])
        if rects_by_column[min_column][0] is None:
            rects_by_column[min_column][0] = bbox[0]
        else:
            rects_by_column[min_column][0] = min(rects_by_column[min_column][0], bbox[0])
            
        max_column = max(cell['column_nums'])
        if rects_by_column[max_column][2] is None:
            rects_by_column[max_column][2] = bbox[2]
        else:
            rects_by_column[max_column][2] = max(rects_by_column[max_column][2], bbox[2])

    table_bbox = list(table_rect)
    table_dict['pdf_table_bbox'] = table_bbox

    for row_num, row_rect in rects_by_row.items():
        row_rect[0] = table_bbox[0]
        row_rect[2] = table_bbox[2]

    for col_num, col_rect in rects_by_column.items():
        col_rect[1] = table_bbox[1]
        col_rect[3] = table_bbox[3]
        
    for k, row in table_dict['rows'].items():
        v = rects_by_row[k]
        table_dict['rows'][k]['pdf_row_bbox'] = list(v)
    for k, column in table_dict['columns'].items():
        v = rects_by_column[k]
        table_dict['columns'][k]['pdf_column_bbox'] = list(v)
        
    for k, row in table_dict['rows'].items():    
        for elem in row['pdf_row_bbox']:
            if elem is None:
                table_dict['reject'].append("undetermined row boundary")
    for k, column in table_dict['columns'].items():
        for elem in column['pdf_column_bbox']:
            if elem is None:
                table_dict['reject'].append("undetermined column boundary")
                
    # Adjust bounding boxes if minor overlap
    fixed_overlap = False
    num_rows = len(table_dict['rows'])
    for row_num in range(num_rows-1):
        row1_bbox = table_dict['rows'][row_num]['pdf_row_bbox']
        row2_bbox = table_dict['rows'][row_num+1]['pdf_row_bbox']
        overlap1 = overlap(row1_bbox, row2_bbox)
        overlap2 = overlap(row2_bbox, row1_bbox)
        
        if overlap1 > 0 and overlap2 > 0:
            if overlap1 < 0.22 and overlap2 < 0.22:
                fixed_overlap = True
                midpoint = 0.5 * (row1_bbox[3] + row2_bbox[1])
                table_dict['rows'][row_num]['pdf_row_bbox'][3] = midpoint
                table_dict['rows'][row_num+1]['pdf_row_bbox'][1] = midpoint
                fixed_overlap = True
            else:
                table_dict['reject'].append("rows intersect: {} {} ({}), {} {} ({})".format(row_num, row1_bbox, overlap1, row_num+1, row2_bbox, overlap2))
        
    # Intersect each row and column to determine grid cell bounding boxes
    #page_words = page.get_text_words()
    for cell in table_dict['cells']:
        rows_rect = Rect()
        cols_rect = Rect()

        for row_num in cell['row_nums']:
            rows_rect.include_rect(table_dict['rows'][row_num]['pdf_row_bbox'])

        for col_num in cell['column_nums']:
            cols_rect.include_rect(table_dict['columns'][col_num]['pdf_column_bbox'])

        pdf_bbox = rows_rect.intersect(cols_rect)
        cell['pdf_bbox'] = list(pdf_bbox)

def identify_projected_row_headers(table_dict):
    num_cols = len(table_dict['columns'])
    cells_with_text_count_by_row = defaultdict(int)
    all_cells_in_row_only_in_one_row_by_row = defaultdict(lambda: True)
    has_first_column_cell_with_text_by_row = defaultdict(bool)
    for cell in table_dict['cells']:
        if len(cell['text_content']) > 0:
            for row_num in cell['row_nums']:
                cells_with_text_count_by_row[row_num] += 1

            if 0 in cell['column_nums']:
                has_first_column_cell_with_text_by_row[row_num] = True

        one_row_only = len(cell['row_nums']) == 1
        for row_num in cell['row_nums']:
            all_cells_in_row_only_in_one_row_by_row[row_num] = all_cells_in_row_only_in_one_row_by_row[row_num] and one_row_only

    projected_row_header_rows = set()
    for row_num, row in table_dict['rows'].items():
        if (not row['is_column_header'] and cells_with_text_count_by_row[row_num] == 1
                and all_cells_in_row_only_in_one_row_by_row[row_num]
                and has_first_column_cell_with_text_by_row[row_num]):
            projected_row_header_rows.add(row_num)
            
    return projected_row_header_rows

def annotate_projected_row_headers(table_dict, exclude=[]):    
    num_cols = len(table_dict['columns'])
    projected_row_header_rows = identify_projected_row_headers(table_dict)
    
    projected_row_header_rows -= set(exclude)

    cells_to_remove = []
    for cell in table_dict['cells']:
        if len(set(cell['row_nums']).intersection(projected_row_header_rows)) > 0:
            if len(cell['text_content']) > 0:
                cell['column_nums'] = list(range(num_cols))
                cell['is_projected_row_header'] = True
            else:
                cells_to_remove.append(cell) # Consolidate blank cells after the first cell into the projected row header
        else:
            cell['is_projected_row_header'] = False

    for cell in cells_to_remove:
        table_dict['fix'].append('merged projected row header')
        table_dict['cells'].remove(cell)
        
    for row_num, row in table_dict['rows'].items():
        if row_num in projected_row_header_rows:
            row['is_projected_row_header'] = True
        else:
            row['is_projected_row_header'] = False
            
    # Delete projected row headers in first rows
    num_rows = len(table_dict['rows'])
    row_nums_to_delete = []
    for row_num in range(0, num_rows):
        if table_dict['rows'][row_num]['is_projected_row_header']:
            row_nums_to_delete.append(row_num)
        else:
            break
            
    if len(row_nums_to_delete) > 0:
        for row_num in row_nums_to_delete:
            del table_dict['rows'][row_num]
            table_dict['fix'].append('removed projected row header at top of table')
            for cell in table_dict['cells'][:]:
                if row_num in cell['row_nums']:
                    table_dict['cells'].remove(cell)
                    
        table_dict['rows'] = {k-len(row_nums_to_delete): v for k, v in table_dict['rows'].items()}
        for cell in table_dict['cells']:
            cell['row_nums'] = [row_num-len(row_nums_to_delete) for row_num in cell['row_nums']]
            
    # Delete projected row headers in last rows
    num_rows = len(table_dict['rows'])
    row_nums_to_delete = []
    for row_num in range(num_rows-1, -1, -1):
        if table_dict['rows'][row_num]['is_projected_row_header']:
            row_nums_to_delete.append(row_num)
        else:
            break
            
    if len(row_nums_to_delete) > 0:
        for row_num in row_nums_to_delete:
            del table_dict['rows'][row_num]
            table_dict['fix'].append('removed projected row header at bottom of table')
            for cell in table_dict['cells'][:]:
                if row_num in cell['row_nums']:
                    table_dict['cells'].remove(cell)

def merge_group(table_dict, group):
    cells_to_delete = []
    if len(group) == 1:
        return table_dict
    group = sorted(group, key=lambda k: min(k['row_nums'])) 
    cell = group[0]
    try:
        cell_text_rect = Rect(cell['pdf_text_tight_bbox'])
    except:
        cell_text_rect = Rect()
    for cell2 in group[1:]:
        cell['row_nums'] = list(set(sorted(cell['row_nums'] + cell2['row_nums'])))
        cell['column_nums'] = list(set(sorted(cell['column_nums'] + cell2['column_nums'])))
        cell['text_content'] = (cell['text_content'].strip() + " " + cell2['text_content'].strip()).strip()
        try:
            cell2_text_rect = Rect(cell2['pdf_text_tight_bbox'])
        except:
            cell2_text_rect = Rect()
        cell_text_rect = cell_text_rect.include_rect(list(cell2_text_rect))
        if cell_text_rect.get_area() == 0:
            cell['pdf_text_tight_bbox'] = []
        else:
            cell['pdf_text_tight_bbox'] = list(cell_text_rect)
        cell['is_projected_row_header'] = False
        cells_to_delete.append(cell2)
        
    try:
        for cell in cells_to_delete:
            table_dict['cells'].remove(cell)
            table_dict['fix'].append('merged oversegmented spanning cell')
    except:
        table_dict['reject'].append("ambiguous spanning cell")
        #raise AmbiguousSpanningCellException

def remove_empty_rows(table_dict):
    num_rows = len(table_dict['rows'])
    num_columns = len(table_dict['columns'])
    has_content_by_row = defaultdict(bool)
    for cell in table_dict['cells']:
        has_content = len(cell['text_content'].strip()) > 0
        for row_num in cell['row_nums']:
            has_content_by_row[row_num] = has_content_by_row[row_num] or has_content
    row_num_corrections = np.cumsum([int(not has_content_by_row[row_num]) for row_num in range(num_rows)]).tolist()
    
    # Delete cells in empty rows and renumber other cells
    cells_to_delete = []
    for cell in table_dict['cells']:
        new_row_nums = []
        for row_num in cell['row_nums']:
            if has_content_by_row[row_num]:
                new_row_nums.append(row_num - row_num_corrections[row_num])
        cell['row_nums'] = new_row_nums
        if len(new_row_nums) == 0:
            cells_to_delete.append(cell)
    for cell in cells_to_delete:
        table_dict['fix'].append('removed empty row')
        table_dict['cells'].remove(cell)
    
    rows = {}
    for row_num, has_content in has_content_by_row.items():
        if has_content:
            new_row_num = row_num - row_num_corrections[row_num]
            rows[new_row_num] = table_dict['rows'][row_num]
    table_dict['rows'] = rows
    
def merge_rows(table_dict):
    num_rows = len(table_dict['rows'])
    num_columns = len(table_dict['columns'])
    co_occurrence_matrix = np.zeros((num_rows, num_rows))
    for cell in table_dict['cells']:
        for row_num1 in cell['row_nums']:
            for row_num2 in cell['row_nums']:
                if row_num1 >= row_num2:
                    continue
                co_occurrence_matrix[row_num1, row_num2] += len(cell['column_nums'])
                
    new_row_num = 0
    current_row_group = 0
    keep_row = [True]
    row_grouping = [current_row_group]
    for row_num in range(num_rows-1):
        if not co_occurrence_matrix[row_num, row_num+1] == num_columns:
            keep_row.append(True)
            new_row_num += 1
        else:
            table_dict['fix'].append('merged rows spanned together in every column')
            keep_row.append(False)
        row_grouping.append(new_row_num)

    for cell in table_dict['cells']:
        cell['row_nums'] = [row_grouping[row_num] for row_num in cell['row_nums'] if keep_row[row_num]]
        
    table_dict['rows'] = {row_grouping[row_num]: table_dict['rows'][row_num] for row_num in range(num_rows) if keep_row[row_num]} 
            
        
def remove_empty_columns(table_dict):
    num_rows = len(table_dict['rows'])
    num_columns = len(table_dict['columns'])
    has_content_by_column = defaultdict(bool)
    for cell in table_dict['cells']:
        has_content = len(cell['text_content'].strip()) > 0
        for column_num in cell['column_nums']:
            has_content_by_column[column_num] = has_content_by_column[column_num] or has_content
    column_num_corrections = np.cumsum([int(not has_content_by_column[column_num]) for column_num in range(num_columns)]).tolist()
    
    # Delete cells in empty columns and renumber other cells
    cells_to_delete = []
    for cell in table_dict['cells']:
        new_column_nums = []
        for column_num in cell['column_nums']:
            if has_content_by_column[column_num]:
                new_column_nums.append(column_num - column_num_corrections[column_num])
        cell['column_nums'] = new_column_nums
        if len(new_column_nums) == 0:
            cells_to_delete.append(cell)
    for cell in cells_to_delete:
        table_dict['fix'].append('removed empty column')
        table_dict['cells'].remove(cell)
    
    columns = {}
    for column_num, has_content in has_content_by_column.items():
        if has_content:
            new_column_num = column_num - column_num_corrections[column_num]
            columns[new_column_num] = table_dict['columns'][column_num]
    table_dict['columns'] = columns
    
def merge_columns(table_dict):
    num_rows = len(table_dict['rows'])
    num_columns = len(table_dict['columns'])
    co_occurrence_matrix = np.zeros((num_columns, num_columns))
    for cell in table_dict['cells']:
        for column_num1 in cell['column_nums']:
            for column_num2 in cell['column_nums']:
                if column_num1 >= column_num2:
                    continue
                co_occurrence_matrix[column_num1, column_num2] += len(cell['row_nums'])
                
    new_column_num = 0
    current_column_group = 0
    keep_column = [True]
    column_grouping = [current_column_group]
    for column_num in range(num_columns-1):
        if not co_occurrence_matrix[column_num, column_num+1] == num_rows:
            keep_column.append(True)
            new_column_num += 1
        else:
            table_dict['fix'].append('merged columns spanned together in every row')
            keep_column.append(False)
        column_grouping.append(new_column_num)

    for cell in table_dict['cells']:
        cell['column_nums'] = [column_grouping[column_num] for column_num in cell['column_nums'] if keep_column[column_num]]
        
    table_dict['columns'] = {column_grouping[column_num]: table_dict['columns'][column_num] for column_num in range(num_columns) if keep_column[column_num]}

# Look for tables with blank cells to merge in the first column
def merge_spanning_cells_in_first_column(table_dict, first_column_merge_exclude=[]):
    numeric_count_by_column = defaultdict(int)
    alpha_count_by_column = defaultdict(int)
    for cell in table_dict['cells']:
        if cell['is_column_header'] or cell['is_projected_row_header']:
            continue
        numeric_count = sum([1 for ch in cell['text_content'] if ch.isnumeric()])
        alpha_count = sum([1 for ch in cell['text_content'] if ch.isalpha()])
        for column_num in cell['column_nums']:
            numeric_count_by_column[column_num] += numeric_count
            alpha_count_by_column[column_num] += alpha_count
    if not alpha_count_by_column[1] > numeric_count_by_column[1]:
        return

    first_column_cells = [cell for cell in table_dict['cells'] if 0 in cell['column_nums']]
    first_column_cells = sorted(first_column_cells, key=lambda item: max(item['row_nums']))
    
    current_filled_cell = None
    groups = defaultdict(list)
    group_num = -1
    for cell in first_column_cells:
        if len(set(cell['row_nums']).intersection(set(first_column_merge_exclude))) > 0:
            group_num += 1
        elif len(cell['text_content']) > 0:
            group_num += 1
        if group_num >= 0:
            groups[group_num].append(cell)
        
    for group_num, group in groups.items():
        if len(group) > 1 and not group[0]['is_projected_row_header'] and not group[0]['is_column_header']:
            merge_group(table_dict, group)

# STANDARDS:
# 1. Column header, if it exists, is a tree structure. FinTabNet contains no header annotation so we can only
#    infer the header given some assumptions. If the top row does not contain all leaf nodes, complete the tree down to the
#    leaf nodes.
# 2. There should be no blank cells in the column header. Blank cells should be aggregated into supercells where possible.
#    - First, blank supercells should be split into blank grid cells.
#    - If a column header cell has only blank grid cells directly below it, extend the cell downward to consume
#      any rows of entirely blank cells directly below it.
#    - After doing this for all column header cells, if a column header cell has only blank cells above it, consume 
#      any rows of entirely blank cells directly above it.
#    - Blank supercells that occur after this grouping are arguably rightly annotated as supercells, but we will not
#      annotate these as supercells at the moment for consistency with previous datasets that annotated these as blank
#      cells only. Detecting a blank supercell would not impact the structure inferred for the table.
#    - Any remaining blank cells are ambiguous, and while it is not good table design to have these, they're not likely
#      to be a nuisance.
# 3. There should be no blank cells in the row header. This is trickier because the row header is not explicitly
#    annotated and must be inferred. See below for more on determining which columns are in the row header.
#    - For columns in the row header, blank cells should be aggregated under the first cell that does not span the
#      entire row. This assumes "top" vertical alignment for text. "Middle" vertical alignment is normally already
#      associated with supercells and is already explicit.
# 4. Inferring the row header.
#    - The row header is explicit whenever the first N columns do not have a column header. In other words, when
#      the stub header is blank. Otherwise it is implicit which columns, if any, correspond to the column header.
#    - If a row header exists, it is also a tree just like the column header and must end at a column of leaf nodes.
#      Not only does this mean supercells cannot be the final column of a row header, but repeated values in a column
#      mean that the column cannot be the final column of a row header.
#    - A column that is not part of the row header (possibly the first column) can have repeated values. Having repeated
#      values is an indication of row header continuation but not of row header status to begin with.
#    - In most cases, numeric values are data. If the numeric values are integer and sorted, this may be part of the row
#      header.
#    - Rows where only one cell has content, either left justified or centered across the table, are part of an implicit
#      first column that begins a row header. The stub header belongs in this first column if there is not a row cell.
# 5. A row cell at the top of the table is either the title of the table (if there are no other row cells in the table),
#    or part of the row header if there are additional row cells below, and belongs in an implicit column.
# 6. Tables have at least one row and two columns. A table with only one column is a list.

def correct_header(table_dict, n=2, assume_header_if_more_than_n_columns=True, last_header_row=None):
    num_columns = len(table_dict['columns'])
    num_rows = len(table_dict['rows'])
    
    if last_header_row is None:

        if num_columns < 2 or num_rows < 1:
            table_dict['reject'].append("small table")
            #raise SmallTableException("Table does not have at least one row and two columns") 

        #---DETERMINE FULL EXTENT OF COLUMN HEADER
        # - Each of the below steps determines different rows that must be in the column header.
        # - The final column header includes all rows that are originally annotated as being in the column
        #   header plus any additional rows determined to be in the column header by the following steps.

        table_has_column_header = False

        # First determine if there is definitely a column header. Four cases:

        # 1. We specify that we want to assume there is one for all tables with more than two columns:
        if assume_header_if_more_than_n_columns and num_columns > n:
            table_has_column_header = True

        # 2. An annotator says there is
        if not table_has_column_header:
            header_rows = [row_num for row_num, row in table_dict['rows'].items() if row['is_column_header']]
            if 0 in header_rows:
                table_has_column_header = True

        # 3. The cell occupying the first row and column is blank
        if not table_has_column_header:
            for cell in table_dict['cells']:
                if 0 in cell['column_nums'] and 0 in cell['row_nums'] and len(cell['text_content'].strip()) == 0:
                    table_has_column_header = True
                    break

        # 4. There is a horizontal spanning cell in the first row
        if not table_has_column_header:
            for cell in table_dict['cells']:
                if 0 in cell['row_nums'] and len(cell['column_nums']) > 1:
                    table_has_column_header = True
                    break

        # Then determine if the column header needs to be extended past its current annotated extent.
        #  1. A header that already is annotated in at least one row continues at least until each column
        #     has a cell occupying only that column
        #  2. A header with a column with a blank cell must continue at least as long as the blank cells continue
        #     (unless rule #1 is satisfied and a possible projected row header is reached?)
        if table_has_column_header:
            first_column_filled_by_row = defaultdict(bool)
            for cell in table_dict['cells']:
                if 0 in cell['column_nums']:
                    if len(cell['text_content']) > 0:
                        for row_num in cell['row_nums']:
                            first_column_filled_by_row[row_num] = True

            first_single_node_row_by_column = defaultdict(lambda: len(table_dict['rows'])-1)
            for cell in table_dict['cells']:
                if len(cell['column_nums']) == 1:
                    first_single_node_row_by_column[cell['column_nums'][0]] = min(first_single_node_row_by_column[cell['column_nums'][0]],
                                                                                   max(cell['row_nums']))

            first_filled_single_node_row_by_column = defaultdict(lambda: len(table_dict['rows'])-1)
            for cell in table_dict['cells']:
                if len(cell['column_nums']) == 1 and len(cell['text_content'].strip()) > 0:
                    first_filled_single_node_row_by_column[cell['column_nums'][0]] = min(first_filled_single_node_row_by_column[cell['column_nums'][0]],
                                                                                   max(cell['row_nums']))

            first_filled_cell_by_column = defaultdict(lambda: len(table_dict['rows'])-1)
            for cell in table_dict['cells']:
                if len(cell['text_content']) > 0:
                    min_row_num = min(cell['row_nums'])
                    for column_num in cell['column_nums']:
                        first_filled_cell_by_column[column_num] = min(first_filled_cell_by_column[column_num],
                                                                      min_row_num)

            projected_row_header_rows = identify_projected_row_headers(table_dict)
            if 0 in projected_row_header_rows:
                table_dict['reject'].append("bad projected row header")
                #raise BadProjectedRowHeaderException('Starting with PRH')
            #for row_num in range(num_rows):
            #    if row_num in projected_row_header_rows:
            #        projected_row_header_rows.remove(row_num)
            #    else:
            #        break

            # Header must continue until at least this row
            minimum_grid_cell_single_node_row = max(first_single_node_row_by_column.values())

            # Header can stop prior to the first of these rows that occurs after the above row
            minimum_first_body_row = min(num_rows-1, max(first_filled_cell_by_column.values()))

            # Determine the max row for which a column N has been single and filled but column N+1 has not
            minimum_all_following_filled = -1
            for row_num in range(num_rows):
                for column_num1 in range(num_columns-1):
                    for column_num2 in range(column_num1+1, num_columns):
                        if (first_filled_single_node_row_by_column[column_num2] > row_num
                            and first_filled_single_node_row_by_column[column_num1] < first_filled_single_node_row_by_column[column_num2]):
                            minimum_all_following_filled = row_num + 1

            #minimum_projected_row_header_row = min([num_rows-1] + [elem for elem in projected_row_header_rows if elem > minimum_grid_cell_single_node_row])
            if len(projected_row_header_rows) > 0:
                minimum_projected_row_header_row = min(projected_row_header_rows)
            else:
                minimum_projected_row_header_row = num_rows

            #first_possible_last_header_row = min(minimum_first_body_row, minimum_projected_row_header_row) - 1
            first_possible_last_header_row = minimum_first_body_row - 1

            last_header_row = max(minimum_all_following_filled,
                                  minimum_grid_cell_single_node_row,
                                  first_possible_last_header_row)

            x = last_header_row
            while(last_header_row < num_rows and not first_column_filled_by_row[last_header_row+1]):
                last_header_row += 1

            #incomplete_header = False # temp for debugging
            if minimum_projected_row_header_row <= last_header_row:
                last_header_row = minimum_projected_row_header_row - 1
                #incomplete_header = True
                
    else:
        table_has_column_header = True
            
    if last_header_row == 0 and table_dict['partial_table']:
        last_header_row = -1
    del table_dict['partial_table']

    for cell in table_dict['cells']:
        if max(cell['row_nums']) <= last_header_row:
            cell['is_column_header'] = True

    for row_num, row in table_dict['rows'].items():
        if row_num <= last_header_row:
            row['is_column_header'] = True

    #if not x == last_header_row:
    #    raise DebugException("Header extended")

    #if minimum_all_following_filled == last_header_row:
    #    raise DebugException

    #if incomplete_header:
    #    raise IncompleteHeaderException("Set last header row to be just before minimum projected row header row".format(last_header_row, minimum_projected_row_header_row))
    
    if not table_has_column_header and num_columns == 2:
        table_dict['reject'].append("ambiguous header")
        #raise AmbiguousHeaderException("Missing header annotation for table with two columns; cannot unambiguously determine header")

def canonicalize(table_dict, first_column_merge_exclude=[]):
    # Preprocessing step: Split every blank spanning cell in the column header into blank grid cells.
    cells_to_delete = []
    try:
        for cell in table_dict['cells']:
            if (cell['is_column_header'] and len(cell['text_content'].strip()) == 0
                    and (len(cell['column_nums']) > 1 or len(cell['row_nums']) > 1)):
                cells_to_delete.append(cell)
                # Split this blank spanning cell into blank grid cells
                for column_num in cell['column_nums']:
                    for row_num in cell['row_nums']:
                        #row_bbox = table_dict['rows'][row_num]['pdf_row_bbox']
                        #column_bbox = table_dict['columns'][column_num]['pdf_column_bbox']
                        #bbox = list(Rect(row_bbox).intersect(list(column_bbox)))
                        new_cell = {'text_content': '',
                                    'column_nums': [column_num],
                                    'row_nums': [row_num],
                                    'is_column_header': cell['is_column_header'],
                                    'pdf_text_tight_bbox': [],
                                    'is_projected_row_header': False}
                        table_dict['cells'].append(new_cell)
    except:
        print(traceback.format_exc())
    for cell in cells_to_delete:
        table_dict['cells'].remove(cell)
        
    # Index cells by row-column position
    cell_grid_index = {}
    for cell in table_dict['cells']:
        for column_num in cell['column_nums']:
            for row_num in cell['row_nums']:
                cell_grid_index[(row_num, column_num)] = cell
        
    # Go bottom up, try to extend non-blank cells up to absorb blank cells
    header_groups = []
    for cell in table_dict['cells']:
        if not cell['is_column_header'] or len(cell['text_content']) == 0:
            continue
        header_group = [cell]
        next_row_num = min(cell['row_nums']) - 1
        for row_num in range(next_row_num, -1, -1):
            all_are_blank = True
            for column_num in cell['column_nums']:
                cell2 = cell_grid_index[(row_num, column_num)]
                all_are_blank = all_are_blank and len(cell2['text_content']) == 0
            if all_are_blank:
                for column_num in cell['column_nums']:
                    header_group.append(cell_grid_index[(row_num, column_num)])
            else:
                break # Stop looking; must be contiguous
        if len(header_group) > 1:
            header_groups.append(header_group)
    for group in header_groups:
        merge_group(table_dict, group)
            
    # Index cells by row-column position
    cell_grid_index = {}
    for cell in table_dict['cells']:
        for column_num in cell['column_nums']:
            for row_num in cell['row_nums']:
                cell_grid_index[(row_num, column_num)] = cell
                
    num_rows = len(table_dict['rows'])
    # Go top down, try to extend non-blank cells down to absorb blank cells
    header_groups = []
    for cell in table_dict['cells']:
        if not cell['is_column_header'] or len(cell['text_content']) == 0:
            continue
        header_group = [cell]
        next_row_num = max(cell['row_nums']) + 1
        for row_num in range(next_row_num, num_rows):
            if not table_dict['rows'][row_num]['is_column_header']:
                break
            all_are_blank = True
            for column_num in cell['column_nums']:
                cell2 = cell_grid_index[(row_num, column_num)]
                all_are_blank = all_are_blank and len(cell2['text_content']) == 0
            if all_are_blank:
                for column_num in cell['column_nums']:
                    header_group.append(cell_grid_index[(row_num, column_num)])
            else:
                break # Stop looking; must be contiguous
        if len(header_group) > 1:
            header_groups.append(header_group)
    for group in header_groups:
        merge_group(table_dict, group)
    
    # Index cells by row-column position
    cell_grid_index = {}
    for cell in table_dict['cells']:
        for column_num in cell['column_nums']:
            for row_num in cell['row_nums']:
                cell_grid_index[(row_num, column_num)] = cell
        
    # Go top down, merge any neighboring cells occupying the same columns, whether they are blank or not
    header_groups_by_row_column = defaultdict(list)
    header_groups = []
    do_full_break = False
    for row_num in table_dict['rows']:
        for column_num in table_dict['columns']:
            cell = cell_grid_index[(row_num, column_num)]
            if not cell['is_column_header']:
                do_full_break = True
                break
            if len(header_groups_by_row_column[(row_num, column_num)]) > 0:
                continue
            if not row_num == min(cell['row_nums']) and column_num == min(cell['column_nums']):
                continue
            # Start new header group
            header_group = [cell]
            next_row_num = max(cell['row_nums']) + 1
            while next_row_num < num_rows:
                cell2 = cell_grid_index[(next_row_num, column_num)]
                if cell2['is_column_header'] and set(cell['column_nums']) == set(cell2['column_nums']):
                    header_group.append(cell2)
                    for row_num2 in cell2['row_nums']:
                        for column_num2 in cell2['column_nums']:
                            header_groups_by_row_column[(row_num2, column_num2)] = header_group
                else:
                    break
                next_row_num = max(cell2['row_nums']) + 1
            for row_num2 in cell['row_nums']:
                for column_num2 in cell['column_nums']:
                    header_groups_by_row_column[(row_num2, column_num2)] = header_group
            if len(header_group) > 1:
                header_groups.append(header_group)
        if do_full_break:
            break
    for group in header_groups:
        merge_group(table_dict, group)
        
    # Merge spanning cells in the row header
    merge_spanning_cells_in_first_column(table_dict, first_column_merge_exclude=first_column_merge_exclude)

def is_all_dots(text):
    if len(text) > 0 and len(text.replace('.','')) == 0:
        return True
    return False

def extract_pdf_text(table_dict, page_words, threshold=0.5):
    adjusted_text_tight_bbox = False
    for cell in table_dict['cells']:
        pdf_text_tight_bbox = cell['pdf_text_tight_bbox']
        pdf_bbox = cell['pdf_bbox']
        
        cell_page_words = [w for w in page_words if Rect(w[:4]).intersect(list(pdf_bbox)).get_area() / Rect(w[:4]).get_area() > threshold]
        cell_words = [w[4] for w in cell_page_words]
        cell_text = ''.join(cell_words)
        
        # Remove trailing dots from cell_page_words
        # Some of the original annotations include dots in the pdf_text_tight_bbox when they shouldn't
        # This code ensures that those are fixed, plus that dots are not added by extracting text from the
        # entire grid cell
        if len(cell_text) > 2 and cell_text[-1] == '.' and cell_text[-2] == '.':
            for page_word in cell_page_words[::-1]:
                if is_all_dots(page_word[4]):
                    table_dict['fix'].append('removed dots from text cell')
                    cell_page_words.remove(page_word)
                else:
                    break
        
        cell_words_rect = Rect()
        for w in cell_page_words:
            cell_words_rect.include_rect(w[:4])
        cell_words = [w[4] for w in cell_page_words]
        cell_text = ' '.join(cell_words)
        cell_text = cell_text.replace(' .', '.').replace(' ,', ',')
        if cell_text.endswith('...'):
            table_dict['reject'].append("dots retained")
            #raise DotsRetainedException("Dots retained in text [{}] '{}'".format(cell_words, cell_text))
        cell['pdf_text_content'] = cell_text
        if cell_words_rect.get_area() > 0:
            new_pdf_text_tight_bbox = list(cell_words_rect)
            if not pdf_text_tight_bbox == new_pdf_text_tight_bbox:
                adjusted_text_tight_bbox = True
                cell['pdf_text_tight_bbox'] = new_pdf_text_tight_bbox
                
    return adjusted_text_tight_bbox

def overlap(bbox1, bbox2):
    try:
        return Rect(bbox1).intersect(list(bbox2)).get_area() / Rect(bbox1).get_area()
    except:
        return 1

def table_text_edit_distance(cells):
    if len(cells) == 0:
        return 0
    
    D = 0
    for cell in cells:
        # Remove spaces and trailing periods
        xml_text = ''.join(cell['text_content'].split()).strip('.')
        pdf_text = ''.join(cell['pdf_text_content'].split()).strip('.')
        L = max(len(xml_text), len(pdf_text))
        if L > 0:
            D += editdistance.eval(xml_text, pdf_text) / L
            
    return D / len(cells)

def quality_control(table_dict, page_words):
    for row_num1, row1 in table_dict['rows'].items():
        for row_num2, row2, in table_dict['rows'].items():
            if row_num1 == row_num2 - 1:
                if row1['pdf_row_bbox'][3] > row2['pdf_row_bbox'][1] + 1:
                    table_dict['reject'].append("rows intersect")
                    #raise RowsIntersectException
                    
    for column_num1, column1 in table_dict['columns'].items():
        for column_num2, column2, in table_dict['columns'].items():
            if column_num1 == column_num2 - 1:
                if column1['pdf_column_bbox'][2] > column2['pdf_column_bbox'][0] + 1:
                    table_dict['reject'].append("columns intersect")
                    #raise ColumnsIntersectException
    
    D = table_text_edit_distance(table_dict['cells'])
    if D > 0.05:
        table_dict['reject'].append("text annotation quality {}".format(D))
        
    word_overlaps = []
    table_bbox = table_dict['pdf_table_bbox']
    for w in page_words:
        if w[4] == '.':
            continue
        if overlap(w[:4], table_bbox) < 0.5:
            continue
        word_overlaps.append(max([overlap(w[:4], cell['pdf_bbox']) for cell in table_dict['cells']]))
    C = sum(word_overlaps) / len(word_overlaps)
    if C < 0.9:
        table_dict['reject'].append("poor text cell fit")

def is_good_bbox(bbox, page_bbox):
    if (not bbox[0] is None and not bbox[1] is None and not bbox[2] is None and not bbox[3] is None
            and bbox[0] >= 0 and bbox[1] >= 0 and bbox[2] <= page_bbox[2] and bbox[3] <= page_bbox[3]
            and bbox[0] < bbox[2]-1 and bbox[1] < bbox[3]-1):
        return True
    return False

def create_document_page_image(doc, page_num, output_image_max_dim=1000):
    page = doc[page_num]
    page_width = page.rect[2]
    page_height = page.rect[3]
    
    if page_height > page_width:
        zoom = output_image_max_dim / page_height
        output_image_height = output_image_max_dim
        output_image_width = int(round(output_image_max_dim * page_width / page_height))
    else:
        zoom = output_image_max_dim / page_width
        output_image_width = output_image_max_dim
        output_image_height = int(round(output_image_max_dim * page_height / page_width))
        
    mat = fitz.Matrix(zoom, zoom)
    pix = page.get_pixmap(matrix = mat, alpha = False)
    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
    
    return img


def create_pascal_voc_page_element(image_filename, output_image_width, output_image_height, database):
    # Create XML of tables on PDF page in PASCAL VOC format
    annotation = ET.Element("annotation")

    folder = ET.SubElement(annotation, "folder").text = ""
    filename = ET.SubElement(annotation, "filename").text = image_filename
    path = ET.SubElement(annotation, "path").text = image_filename
    source = ET.SubElement(annotation, "source")
    database = ET.SubElement(source, "database").text = database
    size = ET.SubElement(annotation, "size")
    width = ET.SubElement(size, "width").text = str(output_image_width)
    height = ET.SubElement(size, "height").text = str(output_image_height)
    depth = ET.SubElement(size, "depth").text = "3"
    segmented = ET.SubElement(annotation, "segmented").text = "0"
    
    return annotation


def create_pascal_voc_object_element(class_name, bbox, page_bbox, output_image_max_dim=1000):
    bbox_area = fitz.Rect(bbox).get_area()
    if bbox_area == 0:
        raise Exception
    intersect_area = fitz.Rect(page_bbox).intersect(fitz.Rect(bbox)).get_area()
    if abs(intersect_area - bbox_area) > 0.1:
        print(bbox)
        print(bbox_area)
        print(page_bbox)
        print(intersect_area)
        raise Exception
    
    object_ = ET.Element("object")
    name = ET.SubElement(object_, "name").text = class_name
    pose = ET.SubElement(object_, "pose").text = "Frontal"
    truncated = ET.SubElement(object_, "truncated").text = "0"
    difficult = ET.SubElement(object_, "difficult").text = "0"
    occluded = ET.SubElement(object_, "occluded").text = "0"
    bndbox = ET.SubElement(object_, "bndbox")
    
    page_width = page_bbox[2] - page_bbox[0]
    page_height = page_bbox[3] - page_bbox[1]
    
    if page_width > page_height:
        output_image_width = output_image_max_dim
        output_image_height = int(output_image_max_dim * page_height / page_width)
    else:
        output_image_height = output_image_max_dim
        output_image_width = int(output_image_max_dim * page_width / page_height)

    xmin = (bbox[0] - page_bbox[0]) * output_image_width / page_width
    ymin = (bbox[1] - page_bbox[1]) * output_image_height / page_height
    xmax = (bbox[2] - page_bbox[0]) * output_image_width / page_width
    ymax = (bbox[3] - page_bbox[1]) * output_image_height / page_height
    
    ET.SubElement(bndbox, "xmin").text = str(xmin)
    ET.SubElement(bndbox, "ymin").text = str(ymin)
    ET.SubElement(bndbox, "xmax").text = str(xmax)
    ET.SubElement(bndbox, "ymax").text = str(ymax)
    
    return object_


def save_xml_pascal_voc(page_annotation, filepath):
    xmlstr = minidom.parseString(ET.tostring(page_annotation)).toprettyxml(indent="   ")
    with open(filepath, "w") as f:
        f.write(xmlstr)
        
        
def bbox_pdf_to_image(bbox, page_bbox, output_image_max_dim=1000):
    page_width = page_bbox[2] - page_bbox[0]
    page_height = page_bbox[3] - page_bbox[1]
    
    if page_width > page_height:
        output_image_width = output_image_max_dim
        output_image_height = int(output_image_max_dim * page_height / page_width)
    else:
        output_image_height = output_image_max_dim
        output_image_width = int(output_image_max_dim * page_width / page_height)

    xmin = (bbox[0] - page_bbox[0]) * output_image_width / page_width
    ymin = (bbox[1] - page_bbox[1]) * output_image_height / page_height
    xmax = (bbox[2] - page_bbox[0]) * output_image_width / page_width
    ymax = (bbox[3] - page_bbox[1]) * output_image_height / page_height
    
    return [xmin, ymin, xmax, ymax]

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 get_tokens_in_table_img(page_words, table_img_bbox):
    tokens = []
    for word_num, word in enumerate(page_words):
        word['flags'] = 0
        word['span_num'] = word_num
        word['line_num'] = 0
        word['block_num'] = 0
        tokens.append(word)

    tokens_in_table = [token for token in tokens if iob(token['bbox'], table_img_bbox) >= 0.5]
    
    return tokens_in_table

def get_args():
    parser = argparse.ArgumentParser()

    parser.add_argument('--data_dir',
                        help="Root directory for source data to process")
    parser.add_argument('--output_dir',
                        help="Root directory for output data")
    parser.add_argument('--train_padding', type=int, default=30,
                        help="The amount of padding to add around a table in the training set when cropping.")
    parser.add_argument('--test_padding', type=int, default=5,
                        help="The amount of padding to add around a table in the val and test sets when cropping.")
    parser.add_argument('--skip_large', action='store_true')
    return parser.parse_args()


def main():
    args = get_args()

    data_directory = args.data_dir

    output_json_directory = os.path.join(args.output_dir, "ICDAR-2013.c_PDF_Annotations_JSON")
    if not os.path.exists(output_json_directory):
        os.makedirs(output_json_directory)

    output_subdirs = ['images', 'train', 'test', 'val']
    output_structure_directory = os.path.join(args.output_dir, "ICDAR-2013.c_Image_Structure_PASCAL_VOC")
    if not os.path.exists(output_structure_directory):
        os.makedirs(output_structure_directory)
    for subdir in output_subdirs:
        subdirectory = os.path.join(output_structure_directory, subdir)
        if not os.path.exists(subdirectory):
            os.makedirs(subdirectory)

    output_table_words_directory = os.path.join(args.output_dir, "ICDAR-2013.c_Image_Table_Words_JSON")
    if not os.path.exists(output_table_words_directory):
        os.makedirs(output_table_words_directory)

    subdirectories = [elem for elem in os.listdir(data_directory)]

    pdf_filepaths = []
    for subdirectory in subdirectories:
        pdf_directory = os.path.join(data_directory, subdirectory)
        for elem in os.listdir(pdf_directory):
            if elem.endswith(".pdf"):
                if os.path.exists(os.path.join(pdf_directory, elem.replace(".pdf", "-str.xml"))):
                    pdf_filepaths.append(os.path.join(pdf_directory, elem))
                    
    num_tables = 0
    for pdf_filepath in pdf_filepaths:
        xml_filepath = pdf_filepath.replace(".pdf", "-str.xml")
        tree = ET.parse(xml_filepath)
        root = tree.getroot()
        tables = []
        for table in root.iter('table'):
            num_tables += len(list(table.iter('region')))

    print("Num tables for structure recognition: {}".format(num_tables))

    processed_count = 0
    good_count = 0
    reject_count = 0
    reject_reasons = defaultdict(list)
    fixes = defaultdict(list)
    kept_as_is_count = 0
    save_count = 0

    output_image_max_dim = 1000

    do_save = True
    do_break = False

    L = len(pdf_filepaths)
    for idx, pdf_filepath in enumerate(pdf_filepaths):
        print("{}/{}".format(idx+1, L), end='\r', flush=True)
        xml_filepath = pdf_filepath.replace(".pdf", "-str.xml")
        
        if 'competition-dataset' in xml_filepath:
            split = 'test'
        else:
            split = 'val'
        
        doc = fitz.open(pdf_filepath)
        error_files = []
        try:
            tables = read_xml(xml_filepath, doc)
        except:
            print("ERROR FOR FILE {}".format(xml_filepath))
            print(traceback.format_exc())
        
        document_tables = []
        for table_index, table_dict in enumerate(tables):
            try:
                #table_dict['document_table_index'] = table_index
                page_num = table_dict['pdf_page_index']
                
                page = doc[page_num]
                if (('eu-dataset/eu-014' in xml_filepath or 'eu-dataset/eu-015' in xml_filepath)
                    and page.rotation == 90):
                    adjust_bbox_coordinates(table_dict, doc)
                page_words = page.get_text_words()
                page_words = [tuple(Rect(w[:4]) * page.rotation_matrix)+w[4:] for w in page_words]
                page_bbox = list(page.rect)
                for w in page_words[:]:
                    if Rect(w[:4]).get_area() == 0 or overlap(w[:4], page_bbox) < 1:
                        page_words.remove(w)
                
                exclude_for_structure = False
                exclude_for_detection = False

                table_dict['exclude_for_structure'] = exclude_for_structure
                table_dict['exclude_for_detection'] = exclude_for_detection
                
                table_dict['split'] = split
                table_dict['pdf_file_name'] = pdf_filepath.split("/")[-1]
                table_dict['pdf_folder'] = pdf_filepath.split("/")[-2] + "/"
                table_dict['document_id'] = "_".join(pdf_filepath.split(".")[0].split("/")[-2:])
                table_dict['icdar2013_source_file_name'] = xml_filepath.split("/")[-1]
                table_dict['pdf_full_page_bbox'] = list(page.rect)
                table_dict['document_table_index'] = table_index
                table_dict['structure_id'] = "{}_{}".format(table_dict['document_id'], table_dict['document_table_index'])

                merged = False
                debug = False

                remove_empty_columns(table_dict)
                merge_columns(table_dict)      
                remove_empty_rows(table_dict)
                merge_rows(table_dict)

                for cell in table_dict['cells']:
                    if cell['text_content'] in ['¢', '$']:
                        table_dict['reject'].append("oversegmented columns")
                        
                # Handle special cases
                if xml_filepath.endswith('competition-dataset-eu/eu-004-str.xml') and page_num == 9 and table_index == 8:
                    exclude = [7]
                elif xml_filepath.endswith('competition-dataset-us/us-007-str.xml') and page_num == 1 and table_index == 0:
                    exclude = [21, 24]
                else:
                    exclude = []
                annotate_projected_row_headers(table_dict, exclude=exclude)

                # Handle special cases
                if xml_filepath.endswith('competition-dataset-eu/eu-021-str.xml') and page_num == 2 and table_index == 0:        
                    correct_header(table_dict, n=1, assume_header_if_more_than_n_columns=True, last_header_row=0)
                elif xml_filepath.endswith('competition-dataset-eu/eu-021-str.xml') and page_num == 6 and table_index == 1:        
                    correct_header(table_dict, n=1, assume_header_if_more_than_n_columns=True, last_header_row=0)
                else:
                    correct_header(table_dict, n=1, assume_header_if_more_than_n_columns=True)
                    
                annotate_projected_row_headers(table_dict, exclude=exclude)
                    
                # Look for tables with multiple headers
                num_rows = len(table_dict['rows'])
                num_columns = len(table_dict['columns'])
                cell_grid = np.zeros((num_rows, num_columns)).astype('str').tolist()
                for cell in table_dict['cells']:
                    for row_num in cell['row_nums']:
                        for column_num in cell['column_nums']:
                            cell_grid[row_num][column_num] = cell['text_content']
                for row_num1 in range(num_rows-1):
                    row1 = table_dict['rows'][row_num1]
                    if not row1['is_column_header']:
                        continue
                    for row_num2 in range(row_num1+1, num_rows):
                        row2 = table_dict['rows'][row_num2]
                        if row2['is_column_header']:
                            continue
                        if cell_grid[row_num1] == cell_grid[row_num2]:
                            table_dict['reject'].append("multiple column headers")

                # Handle special cases
                if xml_filepath.endswith('competition-dataset-us/us-019-str.xml') and page_num == 3 and table_index == 2:
                    first_column_merge_exclude = [4, 9]
                elif xml_filepath.endswith('competition-dataset-us/us-019-str.xml') and page_num == 3 and table_index == 3:
                    first_column_merge_exclude = [7]
                elif xml_filepath.endswith('competition-dataset-eu/eu-009a-str.xml') and page_num == 0 and table_index == 0:
                    first_column_merge_exclude = [3, 5, 7]
                else:
                    first_column_merge_exclude = []
                canonicalize(table_dict, first_column_merge_exclude=first_column_merge_exclude)

                remove_empty_columns(table_dict)
                merge_columns(table_dict)      
                remove_empty_rows(table_dict)
                merge_rows(table_dict)

                for row_num, row in table_dict['rows'].items():
                    if row['is_column_header'] and row_num > 4:
                        table_dict['reject'].append("oversized header")

                # Iterative process because a grid cell bounding box depends on surrounding text, which can
                # change the bounding box for the cell, which can change the text that falls in the bounding box,
                # which can change the bounding boxes for other cells, and so on...
                adjust_text = True
                iterations = 0
                while(adjust_text and iterations < 3):
                    #look_for_dots_in_text_tight_bbox(table_dict, page_words, threshold=0.5)
                    complete_table_grid(table_dict)
                    adjust_text = extract_pdf_text(table_dict, page_words)
                    iterations += 1
                if adjust_text:
                    table_dict['reject'].append("runaway text adjustment")
                    
                for word in page_words:
                    cell_set = set()
                    for cell_num, cell in enumerate(table_dict['cells']):
                        if iob(word[:4], cell['pdf_bbox']) >= 0.5:
                            cell_set.add(cell_num)
                    if len(cell_set) > 1:
                        table_dict['reject'].append('overlapping cells')
                        break

                num_rows = len(table_dict['rows'])
                num_cells_in_last_row = 0
                for cell in table_dict['cells']:
                    if num_rows-1 in cell['row_nums']:
                        num_cells_in_last_row += 1

                # Do manual visual inspection instead
                #quality_control(table_dict, page_words)

                has_body = False
                for row_num, row in table_dict['rows'].items():
                    if not row['is_column_header']:
                        has_body = True
                        break
                if not has_body:
                    table_dict['reject'].append("no table body")

                #if table_dict['rows'][0]['is_projected_row_header']:
                #    table_dict['reject'].append("bad projected row header")
                num_rows = len(table_dict['rows'])
                if table_dict['rows'][num_rows-1]['is_projected_row_header']:
                    table_dict['reject'].append("bad projected row header")
                    
                if (Rect(table_dict['pdf_table_bbox']).get_area() == 0
                    or overlap(table_dict['pdf_table_bbox'], list(page.rect)) < 1):
                    table_dict['reject'].append("bad table bbox")
                    raise Exception("Bad table bbox: {}".format(table_dict['pdf_table_bbox']))
                #print("ok")
                
                if xml_filepath.endswith('us-gov-dataset/us-011-str.xml') and page_num == 1 and table_index == 0:
                    table_dict['exclude_for_detection'] = True
            except KeyboardInterrupt:
                do_break = True
                break
            except:
                print(traceback.format_exc())
                table_dict['reject'].append('unknown exception')
                print('not ok')
                #continue
                
            processed_count += 1

            if len(table_dict['reject']) > 0:
                reject_count += 1

                for reject_reason in set(table_dict['reject']):
                    reject_reasons[reject_reason].append(table_dict['structure_id'])

                table_dict['exclude_for_detection'] = True
                table_dict['exclude_for_structure'] = True
            else:
                good_count += 1

                if len(table_dict['fix']) > 0:
                    for fix in set(table_dict['fix']):
                        fixes[fix].append(table_dict['structure_id'])
                else:
                    kept_as_is_count += 1
                    
                document_tables.append(table_dict)

            del table_dict['reject']
            del table_dict['fix']
            
            if do_break:
                break
        
        if do_break:
            break
            
        # If not all tables present and included for detection, then exclude all for detection
        if not sum([1 for elem in document_tables if not elem['exclude_for_detection']]) == len(tables):
            for table_dict in document_tables:
                table_dict['exclude_for_detection'] = True
                
        if len(document_tables) == 0:
            continue
            
        save_filename = "_".join(pdf_filepath.split(".")[0].split("/")[-2:]).replace(".pdf", "") + "_tables.json"
        save_filepath = os.path.join(output_json_directory, save_filename)
        if do_save:
            with open(save_filepath, 'w') as out_file:
                json.dump(document_tables, out_file, ensure_ascii=False, indent=4)
        
        # Create structure PASCAL VOC data
        # output_structure_directory
        for table_dict in document_tables:
            if table_dict['exclude_for_structure']:
                continue
            page_num = table_dict['pdf_page_index']
            page_rect = list(doc[page_num].rect)
            scale = output_image_max_dim / max(page_rect)
            page_img = create_document_page_image(doc, page_num, output_image_max_dim=output_image_max_dim)
            
            table_num = table_dict['document_table_index']
            table_boxes = []      

            # Create structure recognition data
            class_label = 'table'
            dict_entry = {'class_label': class_label, 'bbox': table_dict['pdf_table_bbox']}
            table_boxes.append(dict_entry)

            rows = table_dict['rows'].values()
            rows = sorted(rows, key=lambda k: k['pdf_row_bbox'][1]) 
            if len(rows) > 1:
                for row1, row2 in zip(rows[:-1], rows[1:]):
                    mid_point = (row1['pdf_row_bbox'][3] + row2['pdf_row_bbox'][1]) / 2
                    row1['pdf_row_bbox'][3] = mid_point
                    row2['pdf_row_bbox'][1] = mid_point
            columns = table_dict['columns'].values()
            columns = sorted(columns, key=lambda k: k['pdf_column_bbox'][0]) 
            for col1, col2 in zip(columns[:-1], columns[1:]):
                mid_point = (col1['pdf_column_bbox'][2] + col2['pdf_column_bbox'][0]) / 2
                col1['pdf_column_bbox'][2] = mid_point
                col2['pdf_column_bbox'][0] = mid_point
            for cell in table_dict['cells']:
                column_nums = cell['column_nums']
                row_nums = cell['row_nums']
                column_rect = Rect()
                row_rect = Rect()
                for column_num in column_nums:
                    column_rect.include_rect(columns[column_num]['pdf_column_bbox'])
                for row_num in row_nums:
                    row_rect.include_rect(rows[row_num]['pdf_row_bbox'])
                cell_rect = column_rect.intersect(row_rect)
                cell['pdf_bbox'] = list(cell_rect)

            header_rect = Rect()
            for cell in table_dict['cells']:
                cell_bbox = cell['pdf_bbox']
                is_blank = len(cell['text_content'].strip()) == 0
                is_spanning_cell = len(cell['row_nums']) > 1 or len(cell['column_nums']) > 1
                is_column_header = cell['is_column_header']
                is_projected_row_header = cell['is_projected_row_header']
                if is_projected_row_header:
                    dict_entry = {'class_label': 'table projected row header', 'bbox': cell['pdf_bbox']}
                    table_boxes.append(dict_entry)                      
                elif is_spanning_cell and not is_blank:
                    dict_entry = {'class_label': 'table spanning cell', 'bbox': cell['pdf_bbox']}
                    table_boxes.append(dict_entry)                     

                if is_column_header:
                    header_rect.include_rect(cell_bbox)

            if header_rect.get_area() > 0:
                dict_entry = {'class_label': 'table column header', 'bbox': list(header_rect)}
                table_boxes.append(dict_entry)

            for row in rows:
                row_bbox = row['pdf_row_bbox']
                dict_entry = {'class_label': 'table row', 'bbox': row_bbox}
                table_boxes.append(dict_entry) 

            # table_entry['columns']
            for column in columns:
                dict_entry = {'class_label': 'table column', 'bbox': column['pdf_column_bbox']}
                table_boxes.append(dict_entry) 

            # Crop
            table_bbox = table_dict['pdf_table_bbox']

            # Convert to image coordinates
            crop_bbox = [int(round(scale * elem)) for elem in table_bbox]

            if split == 'val' or split == 'test':
                padding = args.test_padding
            else:
                padding = args.train_padding
            
            # Pad
            crop_bbox = [crop_bbox[0]-padding,
                        crop_bbox[1]-padding,
                        crop_bbox[2]+padding,
                        crop_bbox[3]+padding]

            # Keep within image
            crop_bbox = [max(0, crop_bbox[0]),
                        max(0, crop_bbox[1]),
                        min(page_img.size[0], crop_bbox[2]),
                        min(page_img.size[1], crop_bbox[3])]

            table_img = page_img.crop(crop_bbox)                    
            for entry in table_boxes:
                bbox = entry['bbox']
                bbox = [scale*elem for elem in bbox]
                bbox = [max(0, bbox[0]-crop_bbox[0]-1),
                        max(0, bbox[1]-crop_bbox[1]-1),
                        min(table_img.size[0], bbox[2]-crop_bbox[0]-1),
                        min(table_img.size[1], bbox[3]-crop_bbox[1]-1)]
                entry['bbox'] = bbox

            # Initialize PASCAL VOC XML
            document_id = table_dict['document_id']
            table_image_filename = document_id + "_table_" + str(table_num) + ".jpg"
            table_image_filepath = os.path.join(output_structure_directory, "images", table_image_filename)
            table_annotation = create_pascal_voc_page_element(table_image_filename, table_img.width, table_img.height,
                                                            database="ICDAR-2013-C-Structure")


            for entry in table_boxes:
                bbox = entry['bbox']
                # Add to PASCAl VOC
                element = create_pascal_voc_object_element(entry['class_label'],
                                                        entry['bbox'], [0, 0, table_img.size[0], table_img.size[1]],
                                                        output_image_max_dim=max(table_img.size))  
                table_annotation.append(element)              

            if do_save:
                table_img.save(table_image_filepath)

                xml_filename = table_image_filename.replace(".jpg", ".xml")
                xml_filepath = os.path.join(output_structure_directory, split, xml_filename)
                save_xml_pascal_voc(table_annotation, xml_filepath)

            # Table words
            # output_table_words_directory
            tokens = []
            for word_num, word in enumerate(doc[page_num].get_text_words()):
                token = {}
                token['flags'] = 0
                token['span_num'] = word_num
                token['line_num'] = 0
                token['block_num'] = 0
                word = tuple(Rect(word[:4]) * doc[page_num].rotation_matrix)+word[4:]
                bbox = [round(scale * v, 5) for v in word[:4]]
                if overlap(bbox, crop_bbox) > 0.75:
                    bbox = [max(0, bbox[0]-crop_bbox[0]-1),
                            max(0, bbox[1]-crop_bbox[1]-1),
                            min(table_img.size[0], bbox[2]-crop_bbox[0]-1),
                            min(table_img.size[1], bbox[3]-crop_bbox[1]-1)]
                    if Rect(bbox).get_area() > 0:
                        token['bbox'] = bbox
                        token['text'] = word[4]
                        tokens.append(token)
                    else:
                        print("REMOVED BAD TABLE WORD")

            if do_save:
                words_save_filepath = os.path.join(output_table_words_directory, table_image_filename.replace(".jpg", "_words.json"))
                with open(words_save_filepath, 'w', encoding='utf8') as f:
                    json.dump(tokens, f)
                save_count += 1
        
        del doc # Just removes from memory, not from disk

    reject_counts = defaultdict(int)
    reject_counts.update({k: len(v) for k, v in reject_reasons.items()})
    for reason, idxs in reject_reasons.items():
        print("{}: {}".format(reason, len(idxs)))
    print("----")
    for reason, idxs in fixes.items():
        print("{}: {}".format(reason, len(idxs)))

    for reason, count in reject_counts.items():
        print("{}: {}".format(reason, count))

    print("{} processed tables".format(processed_count))
    print("{} good tables".format(good_count))
    print("{} rejected tables".format(reject_count))
    print("{} fixed tables".format(good_count - kept_as_is_count))
    print("{} kept as-is tables".format(kept_as_is_count))
    print("{} saved tables".format(save_count))

if __name__ == "__main__":
    main()