data_loader_json.py

# -*- coding: utf-8 -*-
# 2019-01 
# written by Xiaohui Zhao
# xiaohui.zhao@outlook.com
import json
import random
from collections import defaultdict
from os import walk
from os.path import join

import numpy as np

np_load_old = np.load
np.load = lambda *a, **k: np_load_old(*a, allow_pickle=True, **k)
import tokenization
import cv2

DEBUG = False  # True to show grid as image

import unicodedata


def is_number(s):
    try:
        float(s)
        return True
    except ValueError:
        pass

    try:
        unicodedata.numeric(s)
        return True
    except (TypeError, ValueError):
        pass
    return False


class DataLoader():
    """
    grid tables producer
    """

    def __init__(self, params, update_dict=True, load_dictionary=False, data_split=0.75):
        self.random = False
        self.data_laundry = False
        self.encoding_factor = 1  # ensures the size (rows/cols) of grid table compat with the network
        self.classes = ['O', 'DATE_SALE', 'ADDRESS_SELLER', 'DOC_NR', 'DATE_CREATION', 'ADDRESS_CONTRACTOR', 'VAT_ID_CONTRACTOR', 'PAYMENT_METHOD',
                        'VAT_ID_SELLER', 'PAYMENT_BANK_NR', 'TOTAL_PAY', 'TOTAL_CURRENCY', 'TOTAL_TAX', 'TOTAL_WITH_TAX', 'TOTAL_WITHOUT_TAX',
                        'DATE_PAYMENT', 'NAME_SELLER', 'NAME_CONTRACTOR']
        # self.classes = ['DontCare', 'Table'] # for table
        # self.classes = ['DontCare', 'Column0', 'Column1', 'Column2', 'Column3', 'Column4', 'Column5'] # for column
        # self.classes = ['DontCare', 'Column']
        # self.classes = ['DontCare', 'VendorName', 'VendorTaxID', 'InvoiceDate', 'InvoiceNumber', 'ExpenseAmount', 'BaseAmount', 'TaxAmount', 'TaxRate'] # for Spanish project

        self.doc_path = params.doc_path
        self.doc_test_path = params.test_path
        self.use_cutie2 = params.use_cutie2
        self.text_case = params.text_case
        self.tokenize = params.tokenize
        if self.tokenize:
            self.tokenizer = tokenization.FullTokenizer('dict/vocab.txt', do_lower_case=not self.text_case)

        self.rows = self.encoding_factor  # to be updated
        self.cols = self.encoding_factor  # to be updated
        self.segment_grid = params.segment_grid if hasattr(params,
                                                           'segment_grid') else False  # segment grid into two parts if grid is larger than cols_target
        self.augment_strategy = params.augment_strategy if hasattr(params, 'augment_strategy') else 1
        self.pm_strategy = params.positional_mapping_strategy if hasattr(params, 'positional_mapping_strategy') else 2
        self.rows_segment = params.rows_segment if hasattr(params, 'rows_segment') else 72
        self.cols_segment = params.cols_segment if hasattr(params, 'cols_segment') else 72
        self.rows_target = params.rows_target if hasattr(params, 'rows_target') else 64
        self.cols_target = params.cols_target if hasattr(params, 'cols_target') else 64
        self.rows_ulimit = params.rows_ulimit if hasattr(params, 'rows_ulimit') else 80  # handle OOM, must be multiple of self.encoding_factor
        self.cols_ulimit = params.cols_ulimit if hasattr(params, 'cols_ulimit') else 80  # handle OOM, must be multiple of self.encoding_factor

        self.fill_bbox = params.fill_bbox if hasattr(params,
                                                     'fill_bbox') else False  # fill bbox with labels or use one single lable for the entire bbox

        self.data_augmentation_dropout = params.data_augmentation_dropout if hasattr(params,
                                                                                     'data_augmentation_dropout') else False  # TBD: randomly dropout rows/cols
        self.data_augmentation_extra = params.data_augmentation_extra if hasattr(params,
                                                                                 'data_augmentation_extra') else False  # randomly expand rows/cols
        self.da_extra_rows = params.data_augmentation_extra_rows if hasattr(params,
                                                                            'data_augmentation_extra_rows') else 0  # randomly expand rows/cols
        self.da_extra_cols = params.data_augmentation_extra_cols if hasattr(params,
                                                                            'data_augmentation_extra_cols') else 0  # randomly expand rows/cols

        ## 0> parameters to be tuned
        self.load_dictionary = load_dictionary  # load dictionary from file rather than start from empty
        self.dict_path = params.load_dict_from_path if load_dictionary else params.dict_path
        if self.load_dictionary:
            self.dictionary = np.load(self.dict_path + '_dictionary.npy').item()
            self.word_to_index = np.load(self.dict_path + '_word_to_index.npy').item()
            self.index_to_word = np.load(self.dict_path + '_index_to_word.npy').item()
        else:
            self.dictionary = {'[PAD]': 0, '[UNK]': 0}  # word/counts. to be updated in self.load_data() and self._update_docs_dictionary()
            self.word_to_index = {}
            self.index_to_word = {}

        self.data_split = data_split  # split data to training/validation, 0 for all for validation
        self.data_mode = 2  # 0 to consider key and value as two different class, 1 the same class, 2 only value considered
        self.remove_lowfreq_words = False  # remove low frequency words when set as True

        self.num_classes = len(self.classes)
        self.batch_size = params.batch_size if hasattr(params, 'batch_size') else 1

        # TBD: build a special cared dictionary
        self.special_dict = {'*', '='}  # map texts to specific tokens

        ## 1.1> load words and their location/class as training/validation docs and labels 
        self.training_doc_files = self._get_filenames(self.doc_path)
        self.training_docs, self.training_labels = self.load_data(self.training_doc_files,
                                                                  update_dict=update_dict)  # TBD: optimize the update dict flag

        # polish and load dictionary/word_to_index/index_to_word as file
        self.num_words = len(self.dictionary)
        self._updae_word_to_index()
        self._update_docs_dictionary(self.training_docs, 3,
                                     self.remove_lowfreq_words)  # remove low frequency words and add it under the <unknown> key

        # save dictionary/word_to_index/index_to_word as file
        np.save(self.dict_path + '_dictionary.npy', self.dictionary)
        np.save(self.dict_path + '_word_to_index.npy', self.word_to_index)
        np.save(self.dict_path + '_index_to_word.npy', self.index_to_word)
        np.save(self.dict_path + '_classes.npy', self.classes)
        # sorted(self.dictionary.items(), key=lambda x:x[1], reverse=True)

        # split training / validation docs and show statistics
        num_training = int(len(self.training_docs) * self.data_split)
        data_to_be_fetched = [i for i in range(len(self.training_docs))]
        selected_training_index = data_to_be_fetched[:num_training]
        if self.random:
            selected_training_index = random.sample(data_to_be_fetched, num_training)
        selected_validation_index = list(set(data_to_be_fetched).difference(set(selected_training_index)))
        self.validation_docs = [self.training_docs[x] for x in selected_validation_index]
        self.training_docs = [self.training_docs[x] for x in selected_training_index]
        self.validation_labels = self.training_labels
        print('\n\nDATASET: %d vocabularies, %d target classes' % (len(self.dictionary), len(self.classes)))
        print('DATASET: %d for training, %d for validation' % (len(self.training_docs), len(self.validation_docs)))

        ## 1.2> load test files
        self.test_doc_files = self._get_filenames(params.test_path) if hasattr(params, 'test_path') else []
        self.test_docs, self.test_labels = self.load_data(self.test_doc_files, update_dict=update_dict)  # TBD: optimize the update dict flag
        print('DATASET: %d for test from %s \n' % (len(self.test_docs), params.test_path if hasattr(params, 'test_path') else '_'))

        self.data_shape_statistic()  # show data shape static
        if len(self.training_docs) > 0:  # adapt grid table size to all training dataset docs
            self.rows, self.cols, _, _ = self._cal_rows_cols(self.training_docs)
            print('\nDATASHAPE: data set with maximum grid table of ({},{}), updated.\n'.format(self.rows, self.cols))
        else:
            self.rows, self.cols = self.rows_ulimit, self.cols_ulimit

        ## 2> call self.next_batch() outside to generate a batch of grid tables data and labels
        self.training_data_tobe_fetched = [i for i in range(len(self.training_docs))]
        self.validation_data_tobe_fetched = [i for i in range(len(self.validation_docs))]
        self.test_data_tobe_fetched = [i for i in range(len(self.test_docs))]

    def _updae_word_to_index(self):
        if self.load_dictionary:
            max_index = len(self.word_to_index.keys())
            for word in self.dictionary:
                if word not in self.word_to_index:
                    max_index += 1
                    self.word_to_index[word] = max_index
                    self.index_to_word[max_index] = word
        else:
            self.word_to_index = dict(list(zip(self.dictionary.keys(), list(range(self.num_words)))))
            self.index_to_word = dict(list(zip(list(range(self.num_words)), self.dictionary.keys())))

    def _update_docs_dictionary(self, docs, lower_limit, remove_lowfreq_words):
        # assign docs words that appear less than @lower_limit times to word [UNK]
        if remove_lowfreq_words:
            for doc in docs:
                for line in doc:
                    [file_name, dressed_text, word_id, [x_left, y_top, x_right, y_bottom], \
                     [image_w, image_h], max_row_words, max_col_words] = line
                    if self.dictionary[dressed_text] < lower_limit:
                        line = [file_name, '[UNK]', self.word_to_index['[UNK]'], [x_left, y_top, x_right, y_bottom], \
                                [image_w, image_h], max_row_words, max_col_words]
                        self.dictionary[dressed_text] -= 1
                        self.dictionary['[UNK]'] += 1

    def next_batch(self):
        batch_size = self.batch_size

        while True:
            if len(self.training_data_tobe_fetched) < batch_size:
                self.training_data_tobe_fetched = [i for i in range(len(self.training_docs))]
            selected_index = random.sample(self.training_data_tobe_fetched, batch_size)
            self.training_data_tobe_fetched = list(set(self.training_data_tobe_fetched).difference(set(selected_index)))

            training_docs = [self.training_docs[x] for x in selected_index]

            ## data augmentation in each batch if self.data_augmentation==True
            rows, cols, pre_rows, pre_cols = self._cal_rows_cols(training_docs, extra_augmentation=self.data_augmentation_extra,
                                                                 dropout=self.data_augmentation_dropout)
            if self.data_augmentation_extra:
                print('Training grid AUGMENT size: ({},{}) from ({},{})' \
                      .format(rows, cols, pre_rows, pre_cols))

            grid_table, gt_classes, bboxes, label_mapids, bbox_mapids, file_names, updated_cols, ps_indices_x, ps_indices_y = \
                self._positional_mapping(training_docs, self.training_labels, rows, cols)
            if updated_cols > cols:
                print('Training grid EXPAND size: ({},{}) from ({},{})' \
                      .format(rows, updated_cols, rows, cols))
                grid_table, gt_classes, bboxes, label_mapids, bbox_mapids, file_names, _, ps_indices_x, ps_indices_y = \
                    self._positional_mapping(training_docs, self.training_labels, rows, updated_cols, update_col=False)

                ## load image and generate corresponding @ps_1dindices
            images, ps_1d_indices = [], []
            if self.use_cutie2:
                images, ps_1d_indices = self._positional_sampling(self.doc_path, file_names, ps_indices_x, ps_indices_y, updated_cols)
                # print("image fetched {}".format(len(images)))
                if len(images) == batch_size:
                    break
            else:
                break

        batch = {'grid_table': np.array(grid_table), 'gt_classes': np.array(gt_classes),
                 'data_image': np.array(images), 'ps_1d_indices': np.array(ps_1d_indices),  # @images and @ps_1d_indices are only used for CUTIEv2
                 'bboxes': bboxes, 'label_mapids': label_mapids, 'bbox_mapids': bbox_mapids,
                 'file_name': file_names, 'shape': [rows, cols]}
        return batch

    def fetch_validation_data(self):
        batch_size = 1

        while True:
            if len(self.validation_data_tobe_fetched) == 0:
                self.validation_data_tobe_fetched = [i for i in range(len(self.validation_docs))]
            selected_index = random.sample(self.validation_data_tobe_fetched, 1)
            self.validation_data_tobe_fetched = list(set(self.validation_data_tobe_fetched).difference(set(selected_index)))

            validation_docs = [self.validation_docs[x] for x in selected_index]

            ## fixed validation shape leads to better result (to be verified)
            real_rows, real_cols, _, _ = self._cal_rows_cols(validation_docs, extra_augmentation=False)
            rows = max(self.rows_target, real_rows)
            cols = max(self.rows_target, real_cols)

            grid_table, gt_classes, bboxes, label_mapids, bbox_mapids, file_names, updated_cols, ps_indices_x, ps_indices_y = \
                self._positional_mapping(validation_docs, self.validation_labels, rows, cols)
            if updated_cols > cols:
                print('Validation grid EXPAND size: ({},{}) from ({},{})' \
                      .format(rows, updated_cols, rows, cols))
                grid_table, gt_classes, bboxes, label_mapids, bbox_mapids, file_names, _, ps_indices_x, ps_indices_y = \
                    self._positional_mapping(validation_docs, self.validation_labels, rows, updated_cols, update_col=False)

                ## load image and generate corresponding @ps_1dindices
            images, ps_1d_indices = [], []
            if self.use_cutie2:
                images, ps_1d_indices = self._positional_sampling(self.doc_path, file_names, ps_indices_x, ps_indices_y, updated_cols)
                if len(images) == batch_size:
                    break
            else:
                break

        def build_gt_pyramid(self, gt_classes):
            gt_classes = np.array(gt_classes)

            rate = 4  # self.pooling_factor
            b, h, w = np.shape(gt_classes)
            same_padding_left = (rate - w % rate) // 2 if w % rate else 0
            same_padding_right = rate - (rate - w % rate) // 2 if w % rate else 0
            same_padding_top = (rate - h % rate) // 2 if h % rate else 0
            same_padding_bottom = rate - (rate - h % rate) // 2 if h % rate else 0
            for gt_class in gt_classes:
                pad_v = np.pad(gt_class, ((same_padding_top, same_padding_bottom), (0, 0)), 'constant', constant_values=((0, 0), (0, 0)))
                pad_h = np.pad(gt_class, ((0, 0), (same_padding_left, same_padding_right)), 'constant', constant_values=((0, 0), (0, 0)))

                ## find mask range for each single entity
                num_entities = np.max(gt_classes) / self.num_classes
                entity_ranges = [[] for _ in range(0, num_entities)]
                for i in range(1, num_entities):
                    if i % self.num_classes:  # only consider non <DontCare> classes
                        range_y, range_x = np.where(gt_classes == i)
                        # entity_ranges[i] = [top, left, bottom, right, height, width]
                        entity_ranges[i] = [min(range_y), min(range_x), max(range_y), max(range_x),
                                            max(range_y) - min(range_y), max(range_x) - min(range_x)]

        batch = {'grid_table': np.array(grid_table), 'gt_classes': np.array(gt_classes),
                 'data_image': np.array(images), 'ps_1d_indices': np.array(ps_1d_indices),  # @images and @ps_1d_indices are only used for CUTIEv2
                 'bboxes': bboxes, 'label_mapids': label_mapids, 'bbox_mapids': bbox_mapids,
                 'file_name': file_names, 'shape': [rows, cols]}
        return batch

    def fetch_test_data(self):
        batch_size = 1

        while True:
            if len(self.test_data_tobe_fetched) == 0:
                self.test_data_tobe_fetched = [i for i in range(len(self.test_docs))]
                return None

            selected_index = self.test_data_tobe_fetched[0]
            self.test_data_tobe_fetched = list(set(self.test_data_tobe_fetched).difference(set([selected_index])))

            test_docs = [self.test_docs[selected_index]]

            real_rows, real_cols, _, _ = self._cal_rows_cols(test_docs, extra_augmentation=False)
            rows = max(self.rows_target, real_rows)  # small shaped documents have better performance with shape 64
            cols = max(self.cols_target, real_cols)  # large shaped docuemnts have better performance with shape 80

            grid_table, gt_classes, bboxes, label_mapids, bbox_mapids, file_names, updated_cols, ps_indices_x, ps_indices_y = \
                self._positional_mapping(test_docs, self.test_labels, rows, cols)
            if updated_cols > cols:
                print('Test grid EXPAND size: ({},{}) from ({},{})' \
                      .format(rows, updated_cols, rows, cols))
                grid_table, gt_classes, bboxes, label_mapids, bbox_mapids, file_names, _, ps_indices_x, ps_indices_y = \
                    self._positional_mapping(test_docs, self.test_labels, rows, updated_cols, update_col=False)

                ## load image and generate corresponding @ps_1dindices
            images, ps_1d_indices = [], []
            if self.use_cutie2:
                images, ps_1d_indices = self._positional_sampling(self.doc_test_path, file_names, ps_indices_x, ps_indices_y, updated_cols)
                if len(images) == batch_size:
                    break
            else:
                break

        batch = {'grid_table': np.array(grid_table), 'gt_classes': np.array(gt_classes),
                 'data_image': np.array(images), 'ps_1d_indices': np.array(ps_1d_indices),  # @images and @ps_1d_indices are only used for CUTIEv2
                 'bboxes': bboxes, 'label_mapids': label_mapids, 'bbox_mapids': bbox_mapids,
                 'file_name': file_names, 'shape': [rows, cols]}
        return batch

    def _form_label_matrix(self, gt_classes, target_h, target_w):
        """
        build gt_classes and gt_masks with given target featuremap shape (height, width)
        by inspecting bboxes regions (x,y,w,h)
        for table / row / column identity segmentation
        """

        def has_entity_with_augmentation(entity_ranges, roi, use_jittering=False):
            ## find mask with maximum overlap
            max_iou = 0
            max_idx = None
            roi_t, roi_l, roi_b, roi_r = roi
            roi_h = roi_b - roi_t
            roi_w = roi_r - roi_l
            roi_cy = roi_t + roi_h / 2
            roi_cx = roi_l + roi_w / 2
            for idx, entity in enumerate(entity_ranges):
                if len(entity):
                    t, l, b, r, h, w = entity
                    if l > roi_l and r < roi_r and t > roi_t and b < roi_b:  # overlap 1
                        iou = h * w / (roi_h * roi_w)
                    elif l < roi_l and r > roi_r and t < roi_t and b > roi_b:  # overlap 2
                        iou = roi_h * roi_w / (h * w)
                    elif l > roi_r or t > roi_b or b < roi_t or r < roi_l:  # no intersection
                        continue
                    else:
                        iou = min(h * w, roi_h * roi_w) / max(h * w, roi_h * roi_w)

                    # TBD: add jittering augmentation method  
                    if use_jittering:
                        pass
                    if iou > max_iou:
                        max_idx = idx
                        max_iou = iou

            ## check centrality / containment / uniqueness
            t, l, b, r, h, w = entity[idx]
            cy = t + h / 2
            cx = l + w / 2
            if roi_t + h / 3 < cy and cy < toi_b - h / 3 and roi_l + w / 3 < cx and cx < roi_r - w / 3:  # centrality
                if (w > h and roi_w > w * 0.9) or (w < h and roi_h > h * 0.9):  # containment
                    if True:  # uniqueness is already checked with maixmum IOU
                        return True
            return False

        shape = gt_classes.shape
        rate_v = shape[0] / target_h
        rate_h = shape[1] / target_w
        dst_classes = [[[] for i in range(target_h)] for j in range(target_w)]
        dst_masks = [[[] for i in range(target_h)] for j in range(target_w)]
        for i in range(target_h):
            for j in range(target_w):
                roi = [rate_h * j, rate_v * i, rate_h * (j + 1), rate_v * (i + 1)]  # [top, left, bottom, right]

                dst_classes[i][j] = has_entity_with_augmentation(entity_ranges, roi, False)

                mask = gt_classes[roi[1]:roi[3], roi[0]:roi[2]]
                dst_masks[i][j] = mask if dst_classes[i][j] else np.zeros(np.shape(mask))

        return np.array(dst_classes), np.array(dst_masks)

    def data_shape_statistic(self):
        def shape_statistic(docs):
            res_all = defaultdict(int)
            res_row = defaultdict(int)
            res_col = defaultdict(int)
            for doc in docs:
                rows, cols, _, _ = self._cal_rows_cols([doc])
                res_all[rows] += 1
                res_all[cols] += 1
                res_row[rows] += 1
                res_col[cols] += 1
            res_all = sorted(res_all.items(), key=lambda x: x[0], reverse=True)
            res_row = sorted(res_row.items(), key=lambda x: x[0], reverse=True)
            res_col = sorted(res_col.items(), key=lambda x: x[0], reverse=True)
            return res_all, res_row, res_col

        tss, tss_r, tss_c = shape_statistic(self.training_docs)  # training shape static
        vss, vss_r, vss_c = shape_statistic(self.validation_docs)
        tess, tess_r, tess_c = shape_statistic(self.test_docs)
        print("Training statistic: ", tss)
        print("\t num: ", len(self.training_docs))
        print("\t rows statistic: ", tss_r)
        print("\t cols statistic: ", tss_c)
        print("\nValidation statistic: ", vss)
        print("\t num: ", len(self.validation_docs))
        print("\t rows statistic: ", vss_r)
        print("\t cols statistic: ", vss_c)
        print("\nTest statistic: ", tess)
        print("\t num: ", len(self.test_docs))
        print("\t rows statistic: ", tess_r)
        print("\t cols statistic: ", tess_c)

        ## remove data samples not matching the training principle
        def data_laundry(docs):
            idx = 0
            while idx < len(docs):
                rows, cols, _, _ = self._cal_rows_cols([docs[idx]])
                if rows > self.rows_ulimit or cols > self.cols_ulimit:
                    del docs[idx]
                else:
                    idx += 1

        if self.data_laundry:
            print("\nRemoving grids with shape larger than ({},{}).".format(self.rows_ulimit, self.cols_ulimit))
            data_laundry(self.training_docs)
            data_laundry(self.validation_docs)
            data_laundry(self.training_docs)

            tss, tss_r, tss_c = shape_statistic(self.training_docs)  # training shape static
            vss, vss_r, vss_c = shape_statistic(self.validation_docs)
            tess, tess_r, tess_c = shape_statistic(self.test_docs)
            print("Training statistic after laundary: ", tss)
            print("\t num: ", len(self.training_docs))
            print("\t rows statistic: ", tss_r)
            print("\t cols statistic: ", tss_c)
            print("Validation statistic after laundary: ", vss)
            print("\t num: ", len(self.validation_docs))
            print("\t rows statistic: ", vss_r)
            print("\t cols statistic: ", vss_c)
            print("Test statistic after laundary: ", tess)
            print("\t num: ", len(self.test_docs))
            print("\t rows statistic: ", tess_r)
            print("\t cols statistic: ", tess_c)

    def _positional_mapping(self, docs, labels, rows, cols):
        """
        docs in format:
        [[file_name, text, word_id, [x_left, y_top, x_right, y_bottom], [left, top, right, bottom], max_row_words, max_col_words] ]
        return grid_tables, gird_labels, dict bboxes {file_name:[]}, file_names
        """
        grid_tables = []
        gird_labels = []
        ps_indices_x = []  # positional sampling indices
        ps_indices_y = []  # positional sampling indices
        bboxes = {}
        label_mapids = []
        bbox_mapids = []  # [{}, ] bbox identifier, each id with one or multiple bbox/bboxes
        file_names = []
        for doc in docs:
            items = []
            cols_e = 2 * cols  # use @cols_e larger than required @cols as buffer
            grid_table = np.zeros([rows, cols_e], dtype=np.int32)
            grid_label = np.zeros([rows, cols_e], dtype=np.int8)
            ps_x = np.zeros([rows, cols_e], dtype=np.int32)
            ps_y = np.zeros([rows, cols_e], dtype=np.int32)
            bbox = [[] for c in range(cols_e) for r in range(rows)]
            bbox_id, bbox_mapid = 0, {}  # one word in one or many positions in a bbox is mapped in bbox_mapid
            label_mapid = [[] for _ in range(self.num_classes)]  # each class is connected to several bboxes (words)
            drawing_board = np.zeros([rows, cols_e], dtype=str)
            for item in doc:
                file_name = item[0]
                text = item[1]
                word_id = item[2]
                x_left, y_top, x_right, y_bottom = item[3][:]
                left, top, right, bottom = item[4][:]

                dict_id = self.word_to_index[text]
                entity_id, class_id = self._dress_class(file_name, word_id, labels)

                bbox_id += 1
                #                 if self.fill_bbox: # TBD: overlap avoidance
                #                     top = int(rows * y_top / image_h)
                #                     bottom = int(rows * y_bottom / image_h)
                #                     left = int(cols * x_left / image_w)
                #                     right = int(cols * x_right / image_w)
                #                     grid_table[top:bottom, left:right] = dict_id
                #                     grid_label[top:bottom, left:right] = class_id
                #
                #                     label_mapid[class_id].append(bbox_id)
                #                     for row in range(top, bottom):
                #                         for col in range(left, right):
                #                             bbox_mapid[row*cols+col] = bbox_id
                #
                #                     for y in range(top, bottom):
                #                         for x in range(left, right):
                #                             bbox[y][x] = [x_left, y_top, x_right-x_left, y_bottom-y_top]
                label_mapid[class_id].append(bbox_id)

                # v_c = (y_top - top + (y_bottom-y_top)/2) / (bottom-top)
                # h_c = (x_left - left + (x_right-x_left)/2) / (right-left)
                # v_c = (y_top + (y_bottom-y_top)/2) / bottom
                # h_c = (x_left + (x_right-x_left)/2) / right
                # v_c = (y_top-top) / (bottom-top)
                # h_c = (x_left-left) / (right-left)
                # v_c = (y_top) / (bottom)
                # h_c = (x_left) / (right)
                box_y = y_top + (y_bottom - y_top) / 2
                box_x = x_left  # h_l is used for image feature map positional sampling
                v_c = (y_top - top + (y_bottom - y_top) / 2) / (bottom - top)
                h_c = (x_left - left + (x_right - x_left) / 2) / (right - left)  # h_c is used for sorting items
                row = int(rows * v_c)
                col = int(cols * h_c)
                items.append([row, col, [box_y, box_x], [v_c, h_c], file_name, dict_id, class_id, entity_id, bbox_id,
                              [x_left, y_top, x_right - x_left, y_bottom - y_top]])

            items.sort(key=lambda x: (x[0], x[3], x[5]))  # sort according to row > h_c > bbox_id
            for item in items:
                row, col, [box_y, box_x], [v_c, h_c], file_name, dict_id, class_id, entity_id, bbox_id, box = item
                entity_class_id = entity_id * self.num_classes + class_id

                while col < cols and grid_table[row, col] != 0:
                    col += 1

                    # self.pm_strategy 0: skip if overlap
                # self.pm_strategy 1: shift to find slot if overlap
                # self.pm_strategy 2: expand grid table if overlap
                if self.pm_strategy == 0:
                    if col == cols:
                        print('overlap in {} row {} r{}c{}!'.
                              format(file_name, row, rows, cols))
                        # print(grid_table[row,:])
                        # print('overlap in {} <{}> row {} r{}c{}!'.
                        #      format(file_name, self.index_to_word[dict_id], row, rows, cols))
                    else:
                        grid_table[row, col] = dict_id
                        grid_label[row, col] = entity_class_id
                        bbox_mapid[row * cols + col] = bbox_id
                        bbox[row * cols + col] = box
                elif self.pm_strategy == 1 or self.pm_strategy == 2:
                    ptr = 0
                    if col == cols:  # shift to find slot to drop the current item
                        col -= 1
                        while ptr < cols and grid_table[row, ptr] != 0:
                            ptr += 1
                        if ptr == cols:
                            grid_table[row, :-1] = grid_table[row, 1:]
                        else:
                            grid_table[row, ptr:-1] = grid_table[row, ptr + 1:]

                    if self.pm_strategy == 2:
                        while col < cols_e and grid_table[row, col] != 0:
                            col += 1
                        if col > cols:  # update maximum cols in current grid
                            print(grid_table[row, :col])
                            print('overlap in {} <{}> row {} r{}c{}!'.
                                  format(file_name, self.index_to_word[dict_id], row, rows, cols))
                            cols = col
                        if col == cols_e:
                            print('overlap!')

                    grid_table[row, col] = dict_id
                    grid_label[row, col] = entity_class_id
                    ps_x[row, col] = box_x
                    ps_y[row, col] = box_y
                    bbox_mapid[row * cols + col] = bbox_id
                    bbox[row * cols + col] = box

            cols = self._fit_shape(cols)
            grid_table = grid_table[..., :cols]
            grid_label = grid_label[..., :cols]
            ps_x = np.array(ps_x[..., :cols])
            ps_y = np.array(ps_y[..., :cols])

            if DEBUG:
                self.grid_visualization(file_name, grid_table, grid_label)

            grid_tables.append(np.expand_dims(grid_table, -1))
            gird_labels.append(grid_label)
            ps_indices_x.append(ps_x)
            ps_indices_y.append(ps_y)
            bboxes[file_name] = bbox
            label_mapids.append(label_mapid)
            bbox_mapids.append(bbox_mapid)
            file_names.append(file_name)

        return grid_tables, gird_labels, bboxes, label_mapids, bbox_mapids, file_names, cols, ps_indices_x, ps_indices_y

    def _positional_sampling(self, path, file_names, ps_indices_x, ps_indices_y, updated_cols):
        images, ps_1d_indices = [], []

        ## load image and generate corresponding @ps_1dindices
        max_h, max_w = 0, updated_cols
        for i in range(len(file_names)):
            file_name = file_names[i]
            file_path = join(path, file_name)  # TBD: ensure image is upright
            ps_1d_x = np.array(ps_indices_x[i], dtype=np.float32).reshape([-1])
            ps_1d_y = np.array(ps_indices_y[i], dtype=np.float32).reshape([-1])

            image = cv2.imread(file_path)
            if image is not None:
                h, w, _ = image.shape  # [h,w,c]
                factor = max_w / w

                h = int(h * factor)
                ps_1d_x *= factor  # TBD: implement more accurate mapping method rather than nearest neighbor, since the .4 or .6 leads to two different sampling results
                ps_1d_y *= factor

                ps_1d = np.int32(np.floor(ps_1d_x) + np.floor(ps_1d_y) * max_w)
                max_items = max_w * h - 1
                for i in range(len(ps_1d)):
                    if ps_1d[i] > max_items - 1:
                        ps_1d[i] = max_items - 1

                image = cv2.resize(image, (max_w, h))
                image = (image - 127.5) / 255
            else:
                # print('Warning: {} image not found!'.format(file_path))
                print('{} ignored due to image file not found.'.format(file_path))
                image, ps_1d = None, None
                break

            if image is not None and ps_1d is not None:  # ignore data with no images
                ps_1d_indices.append(ps_1d)
                images.append(image)
                h, w, c = image.shape
                if h > max_h:
                    max_h = h
            else:
                pass
                # print('{} ignored due to image file not found.'.format(file_path))

        ## pad image to the same shape
        for i, image in enumerate(images):
            pad_img = np.zeros([max_h, max_w, 3], dtype=image.dtype)
            pad_img[:image.shape[0], :, :] = image
            images[i] = pad_img

        return images, ps_1d_indices

    def load_data(self, data_files, update_dict=False):
        """
        label_dressed in format:
        {file_id: {class: [{'key_id':[], 'value_id':[], 'key_text':'', 'value_text':''}, ] } }
        load doc words with location and class returned in format: 
        [[file_name, text, word_id, [x_left, y_top, x_right, y_bottom], [left, top, right, bottom], max_row_words, max_col_words] ]
        """
        label_dressed = {}
        doc_dressed = []
        if not data_files:
            print("no data file found.")
        for file in data_files:
            with open(file, encoding='utf-8') as f:
                data = json.load(f)
                file_id = data['global_attributes']['file_id']

                label = self._collect_label(file_id, data['fields'])
                # ignore corrupted data
                if not label:
                    continue
                label_dressed.update(label)

                data = self._collect_data(file_id, data['text_boxes'], update_dict)
                for i in data:
                    doc_dressed.append(i)

        return doc_dressed, label_dressed

    def _cal_rows_cols(self, docs, extra_augmentation=False, dropout=False):
        max_row = self.encoding_factor
        max_col = self.encoding_factor
        for doc in docs:
            for line in doc:
                _, _, _, _, _, max_row_words, max_col_words = line
                if max_row_words > max_row:
                    max_row = max_row_words
                if max_col_words > max_col:
                    max_col = max_col_words

        pre_rows = self._fit_shape(max_row)  # (max_row//self.encoding_factor+1) * self.encoding_factor
        pre_cols = self._fit_shape(max_col)  # (max_col//self.encoding_factor+1) * self.encoding_factor

        rows, cols = 0, 0
        if extra_augmentation:
            pad_row = int(random.gauss(0, self.da_extra_rows * self.encoding_factor))  # abs(random.gauss(0, u))
            pad_col = int(random.gauss(0, self.da_extra_cols * self.encoding_factor))  # random.randint(0, u)

            if self.augment_strategy == 1:  # strategy 1: augment data by increasing grid shape sizes
                pad_row = abs(pad_row)
                pad_col = abs(pad_col)
                rows = self._fit_shape(max_row + pad_row)  # apply upper boundary to avoid OOM
                cols = self._fit_shape(max_col + pad_col)  # apply upper boundary to avoid OOM
            elif self.augment_strategy == 2 or self.augment_strategy == 3:  # strategy 2: augment by increasing or decreasing the target gird shape size
                rows = self._fit_shape(max(self.rows_target + pad_row, max_row))  # protect grid shape
                cols = self._fit_shape(max(self.cols_target + pad_col, max_col))  # protect grid shape
            else:
                raise Exception('unknown augment strategy')
            rows = min(rows, self.rows_ulimit)  # apply upper boundary to avoid OOM
            cols = min(cols, self.cols_ulimit)  # apply upper boundary to avoid OOM
        else:
            rows = pre_rows
            cols = pre_cols
        return rows, cols, pre_rows, pre_cols

    def _fit_shape(self, shape):  # modify shape size to fit the encoding factor
        while shape % self.encoding_factor:
            shape += 1
        return shape

    def _expand_shape(self, shape):  # expand shape size with step 2
        return self._fit_shape(shape + 1)

    def _collect_data(self, file_name, content, update_dict):
        """
        dress and preserve only interested data.
        """
        content_dressed = []
        left, top, right, bottom, buffer = 9999, 9999, 0, 0, 2
        for line in content:
            bbox = line['bbox']  # handle data corrupt
            if len(bbox) == 0:
                continue
            if line['text'] in self.special_dict:  # ignore potential overlap causing characters
                continue

            x_left, y_top, x_right, y_bottom = self._dress_bbox(bbox)
            # TBD: the real image size is better for calculating the relative x/y/w/h
            if x_left < left: left = x_left - buffer
            if y_top < top: top = y_top - buffer
            if x_right > right: right = x_right + buffer
            if y_bottom > bottom: bottom = y_bottom + buffer

            word_id = line['id']
            dressed_texts = self._dress_text(line['text'], update_dict)

            num_block = len(dressed_texts)
            for i, dressed_text in enumerate(dressed_texts):  # handling tokenized text, separate bbox
                new_left = int(x_left + (x_right - x_left) / num_block * (i))
                new_right = int(x_left + (x_right - x_left) / num_block * (i + 1))
                content_dressed.append([file_name, dressed_text, word_id, [new_left, y_top, new_right, y_bottom]])

        # initial calculation of maximum number of words in rows/cols in terms of image size
        num_words_row = [0 for _ in range(bottom)]  # number of words in each row
        num_words_col = [0 for _ in range(right)]  # number of words in each column
        for line in content_dressed:
            _, _, _, [x_left, y_top, x_right, y_bottom] = line
            for y in range(y_top, y_bottom):
                num_words_row[y] += 1
            for x in range(x_left, x_right):
                num_words_col[x] += 1
        max_row_words = self._fit_shape(max(num_words_row))
        max_col_words = 0  # self._fit_shape(max(num_words_col))

        # further expansion of maximum number of words in rows/cols in terms of grid shape
        max_rows = max(self.encoding_factor, max_row_words)
        max_cols = max(self.encoding_factor, max_col_words)
        DONE = False
        while not DONE:
            DONE = True
            grid_table = np.zeros([max_rows, max_cols], dtype=np.int32)
            for line in content_dressed:
                _, _, _, [x_left, y_top, x_right, y_bottom] = line
                row = int(max_rows * (y_top - top + (y_bottom - y_top) / 2) / (bottom - top))
                col = int(max_cols * (x_left - left + (x_right - x_left) / 2) / (right - left))
                # row = int(max_rows * (y_top + (y_bottom-y_top)/2) / (bottom))
                # col = int(max_cols * (x_left + (x_right-x_left)/2) / (right))
                # row = int(max_rows * (y_top-top) / (bottom-top))
                # col = int(max_cols * (x_left-left) / (right-left))
                # row = int(max_rows * (y_top) / (bottom))
                # col = int(max_cols * (x_left) / (right))
                # row = int(max_rows * (y_top + (y_bottom-y_top)/2) / bottom)
                # col = int(max_cols * (x_left + (x_right-x_left)/2) / right)

                while col < max_cols and grid_table[row, col] != 0:  # shift to find slot to drop the current item
                    col += 1
                if col == max_cols:  # shift to find slot to drop the current item
                    col -= 1
                    ptr = 0
                    while ptr < max_cols and grid_table[row, ptr] != 0:
                        ptr += 1
                    if ptr == max_cols:  # overlap cannot be solved in current row, then expand the grid
                        max_cols = self._expand_shape(max_cols)
                        DONE = False
                        break

                    grid_table[row, ptr:-1] = grid_table[row, ptr + 1:]

                if DONE:
                    if row > max_rows or col > max_cols:
                        print('wrong')
                    grid_table[row, col] = 1

        max_rows = self._fit_shape(max_rows)
        max_cols = self._fit_shape(max_cols)

        # print('{} collected in shape: {},{}'.format(file_name, max_rows, max_cols))

        # segment grid into two parts if number of cols is larger than self.cols_target
        data = []
        if self.segment_grid and max_cols > self.cols_segment:
            content_dressed_left = []
            content_dressed_right = []
            cnt = defaultdict(int)  # counter for number of words in a specific row
            cnt_l, cnt_r = defaultdict(int), defaultdict(int)  # update max_cols if larger than self.cols_segment
            left_boundary = max_cols - self.cols_segment
            right_boundary = self.cols_segment
            for i, line in enumerate(content_dressed):
                file_name, dressed_text, word_id, [x_left, y_top, x_right, y_bottom] = line

                row = int(max_rows * (y_top + (y_bottom - y_top) / 2) / bottom)
                cnt[row] += 1
                if cnt[row] <= left_boundary:
                    cnt_l[row] += 1
                    content_dressed_left.append([file_name, dressed_text, word_id, [x_left, y_top, x_right, y_bottom], \
                                                 [left, top, right, bottom], max_rows, self.cols_segment])
                elif left_boundary < cnt[row] <= right_boundary:
                    cnt_l[row] += 1
                    cnt_r[row] += 1
                    content_dressed_left.append([file_name, dressed_text, word_id, [x_left, y_top, x_right, y_bottom], \
                                                 [left, top, right, bottom], max_rows, self.cols_segment])
                    content_dressed_right.append([file_name, dressed_text, word_id, [x_left, y_top, x_right, y_bottom], \
                                                  [left, top, right, bottom], max_rows, max(max(cnt_r.values()), self.cols_segment)])
                else:
                    cnt_r[row] += 1
                    content_dressed_right.append([file_name, dressed_text, word_id, [x_left, y_top, x_right, y_bottom], \
                                                  [left, top, right, bottom], max_rows, max(max(cnt_r.values()), self.cols_segment)])
            # print(sorted(cnt.items(), key=lambda x:x[1], reverse=True))
            # print(sorted(cnt_l.items(), key=lambda x:x[1], reverse=True))
            # print(sorted(cnt_r.items(), key=lambda x:x[1], reverse=True))
            if max(cnt_l.values()) < 2 * self.cols_segment:
                data.append(content_dressed_left)
            if max(cnt_r.values()) < 2 * self.cols_segment:  # avoid OOM, which tends to happen in the right side
                data.append(content_dressed_right)
        else:
            for i, line in enumerate(content_dressed):  # append height/width/numofwords to the list
                file_name, dressed_text, word_id, [x_left, y_top, x_right, y_bottom] = line
                content_dressed[i] = [file_name, dressed_text, word_id, [x_left, y_top, x_right, y_bottom], \
                                      [left, top, right, bottom], max_rows, max_cols]
            data.append(content_dressed)
        return data

    def _collect_label(self, file_id, content):
        """
        dress and preserve only interested data.
        label_dressed in format:
        {file_id: {class: [{'key_id':[], 'value_id':[], 'key_text':'', 'value_text':''}, ] } }
        """
        label_dressed = dict()
        label_dressed[file_id] = {cls: [] for cls in self.classes[1:]}
        for line in content:
            cls = line['field_name']
            if cls in self.classes:
                # identity = line.get('identity', 0)
                label_dressed[file_id][cls].append({'key_id': [], 'value_id': [], 'key_text': '', 'value_text': ''})
                label_dressed[file_id][cls][-1]['key_id'] = line.get('key_id', [])
                label_dressed[file_id][cls][-1]['value_id'] = line['value_id']  # value_id
                label_dressed[file_id][cls][-1]['key_text'] = line.get('key_text', [])
                label_dressed[file_id][cls][-1]['value_text'] = line['value_text']  # value_text

        # handle corrupted data
        for cls in label_dressed[file_id]:
            for idx, label in enumerate(label_dressed[file_id][cls]):
                if len(label) == 0:  # no relevant class in sample @file_id
                    continue
                if (len(label['key_text']) > 0 and len(label['key_id']) == 0) or \
                        (len(label['value_text']) > 0 and len(label['value_id']) == 0):
                    return None

        return label_dressed

    def _dress_class(self, file_name, word_id, labels):
        """
        label_dressed in format:
        {file_id: {class: [{'key_id':[], 'value_id':[], 'key_text':'', 'value_text':''}, ] } }
        """
        if file_name in labels:
            for cls, cls_labels in labels[file_name].items():
                for idx, cls_label in enumerate(cls_labels):
                    for key, values in cls_label.items():
                        if (key == 'key_id' or key == 'value_id') and word_id in values:
                            if key == 'key_id':
                                if self.data_mode == 0:
                                    return idx, self.classes.index(cls) * 2 - 1  # odd
                                elif self.data_mode == 1:
                                    return idx, self.classes.index(cls)
                                else:  # ignore key_id when self.data_mode is not 0 or 1
                                    return 0, 0
                            elif key == 'value_id':
                                if self.data_mode == 0:
                                    return idx, self.classes.index(cls) * 2  # even
                                else:  # when self.data_mode is 1 or 2
                                    return idx, self.classes.index(cls)
            return 0, 0  # 0 is of class type 'DontCare'
        print("No matched labels found for {}".format(file_name))

    def _dress_text(self, text, update_dict):
        """
        three cases covered: 
        alphabetic string, numeric string, special character
        """
        string = text if self.text_case else text.lower()
        for i, c in enumerate(string):
            if is_number(c):
                string = string[:i] + '0' + string[i + 1:]

        strings = [string]
        if self.tokenize:
            strings = self.tokenizer.tokenize(strings[0])
            # print(string, '-->', strings)

        for idx, string in enumerate(strings):
            if string.isalpha():
                if string in self.special_dict:
                    string = self.special_dict[string]
                # TBD: convert a word to its most similar word in a known vocabulary
            elif is_number(string):
                pass
            elif len(string) == 1:  # special character
                pass
            else:
                # TBD: seperate string as parts for alpha and number combinated strings
                # string = re.findall('[a-z]+', string)
                pass

            if string not in self.dictionary.keys():
                if update_dict:
                    self.dictionary[string] = 0
                else:
                    # print('unknown text: ' + string)
                    string = '[UNK]'  # TBD: take special care to unmet words\
            self.dictionary[string] += 1

            strings[idx] = string
        return strings

    def _dress_bbox(self, bbox):
        positions = np.array(bbox).reshape([-1])
        x_left = max(0, min(positions[0::2]))
        x_right = max(positions[0::2])
        y_top = max(0, min(positions[1::2]))
        y_bottom = max(positions[1::2])
        w = x_right - x_left
        h = y_bottom - y_top
        return int(x_left), int(y_top), int(x_right), int(y_bottom)

    def _get_filenames(self, data_path):
        files = []
        for dirpath, dirnames, filenames in walk(data_path):
            for filename in filenames:
                file = join(dirpath, filename)
                if file.endswith('csv') or file.endswith('json'):
                    files.append(file)
        return files

    def grid_visualization(self, file_name, grid, label):
        import cv2
        height, width = np.shape(grid)
        grid_box_h, grid_box_w = 20, 40
        palette = np.zeros([height * grid_box_h, width * grid_box_w, 3], np.uint8)
        font = cv2.FONT_HERSHEY_SIMPLEX
        gt_color = [[255, 250, 240], [152, 245, 255], [127, 255, 212], [100, 149, 237],
                    [192, 255, 62], [175, 238, 238], [255, 130, 171], [240, 128, 128], [255, 105, 180]]
        cv2.putText(palette, file_name + "({},{})".format(height, width), (grid_box_h, grid_box_w), font, 0.6, [255, 0, 0])
        for h in range(height):
            cv2.line(palette, (0, h * grid_box_h), (width * grid_box_w, h * grid_box_h), (100, 100, 100))
            for w in range(width):
                if grid[h, w]:
                    org = (int((w + 1) * grid_box_w * 0.7), int((h + 1) * grid_box_h * 0.9))
                    color = gt_color[label[h, w]]
                    cv2.putText(palette, self.index_to_word[grid[h, w]], org, font, 0.4, color)

        img = cv2.imread(self.doc_path + '/' + file_name)
        if img is not None:
            shape = list(img.shape)
            max_len = 768
            factor = max_len / max(shape)
            shape[0], shape[1] = [int(s * factor) for s in shape[:2]]
            img = cv2.resize(img, (shape[1], shape[0]))
            cv2.imshow("img", img)
        cv2.imshow("grid", palette)
        cv2.waitKey(0)