601 lines
24 KiB
Python
601 lines
24 KiB
Python
|
import os
|
||
|
|
from typing import Any
|
||
|
|
import numpy as np
|
||
|
|
import pandas as pd
|
||
|
|
import imagesize
|
||
|
|
import itertools
|
||
|
|
from PIL import Image
|
||
|
|
import argparse
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from utils.utils import read_ocr_result_from_txt, read_json, post_process_basic_ocr
|
||
|
|
from utils.run_ocr import load_ocr_engine, process_img
|
||
|
|
from lightning_modules.utils import sliding_windows
|
||
|
|
|
||
|
|
|
||
|
|
class KVUProcess:
|
||
|
|
def __init__(self, tokenizer_layoutxlm, feature_extractor, backbone_type, class_names, slice_interval, window_size, run_ocr, max_seq_length=512, mode=0):
|
||
|
|
self.tokenizer_layoutxlm = tokenizer_layoutxlm
|
||
|
|
self.feature_extractor = feature_extractor
|
||
|
|
|
||
|
|
self.max_seq_length = max_seq_length
|
||
|
|
self.backbone_type = backbone_type
|
||
|
|
self.class_names = class_names
|
||
|
|
|
||
|
|
self.slice_interval = slice_interval
|
||
|
|
self.window_size = window_size
|
||
|
|
self.run_ocr = run_ocr
|
||
|
|
self.mode = mode
|
||
|
|
|
||
|
|
self.pad_token_id_layoutxlm = tokenizer_layoutxlm.convert_tokens_to_ids(tokenizer_layoutxlm._pad_token)
|
||
|
|
self.cls_token_id_layoutxlm = tokenizer_layoutxlm.convert_tokens_to_ids(tokenizer_layoutxlm._cls_token)
|
||
|
|
self.sep_token_id_layoutxlm = tokenizer_layoutxlm.convert_tokens_to_ids(tokenizer_layoutxlm._sep_token)
|
||
|
|
self.unk_token_id_layoutxlm = tokenizer_layoutxlm.convert_tokens_to_ids(self.tokenizer_layoutxlm._unk_token)
|
||
|
|
|
||
|
|
|
||
|
|
self.class_idx_dic = dict(
|
||
|
|
[(class_name, idx) for idx, class_name in enumerate(self.class_names)]
|
||
|
|
)
|
||
|
|
self.ocr_engine = None
|
||
|
|
if self.run_ocr == 1:
|
||
|
|
self.ocr_engine = load_ocr_engine()
|
||
|
|
|
||
|
|
def __call__(self, img_path: str, ocr_path: str) -> list:
|
||
|
|
if (self.run_ocr == 1) or (not os.path.exists(ocr_path)):
|
||
|
|
process_img(img_path, "tmp.txt", self.ocr_engine, export_img=False)
|
||
|
|
ocr_path = "tmp.txt"
|
||
|
|
lbboxes, lwords = read_ocr_result_from_txt(ocr_path)
|
||
|
|
lwords = post_process_basic_ocr(lwords)
|
||
|
|
bbox_windows = sliding_windows(lbboxes, self.window_size, self.slice_interval)
|
||
|
|
word_windows = sliding_windows(lwords, self.window_size, self.slice_interval)
|
||
|
|
assert len(bbox_windows) == len(word_windows), f"Shape of lbboxes and lwords after sliding window is not the same {len(bbox_windows)} # {len(word_windows)}"
|
||
|
|
|
||
|
|
width, height = imagesize.get(img_path)
|
||
|
|
images = [Image.open(img_path).convert("RGB")]
|
||
|
|
image_features = torch.from_numpy(self.feature_extractor(images)['pixel_values'][0].copy())
|
||
|
|
|
||
|
|
|
||
|
|
if self.mode == 0:
|
||
|
|
output = self.preprocess(lbboxes, lwords,
|
||
|
|
{'image': image_features, 'width': width, 'height': height, 'img_path': img_path},
|
||
|
|
max_seq_length=self.max_seq_length)
|
||
|
|
elif self.mode == 1:
|
||
|
|
output = {}
|
||
|
|
windows = []
|
||
|
|
for i in range(len(bbox_windows)):
|
||
|
|
_words = word_windows[i]
|
||
|
|
_bboxes = bbox_windows[i]
|
||
|
|
windows.append(
|
||
|
|
self.preprocess(
|
||
|
|
_bboxes, _words,
|
||
|
|
{'image': image_features, 'width': width, 'height': height, 'img_path': img_path},
|
||
|
|
max_seq_length=self.max_seq_length)
|
||
|
|
)
|
||
|
|
|
||
|
|
output['windows'] = windows
|
||
|
|
elif self.mode == 2:
|
||
|
|
output = {}
|
||
|
|
windows = []
|
||
|
|
output['doduments'] = self.preprocess(lbboxes, lwords,
|
||
|
|
{'image': image_features, 'width': width, 'height': height, 'img_path': img_path},
|
||
|
|
max_seq_length=2048)
|
||
|
|
for i in range(len(bbox_windows)):
|
||
|
|
_words = word_windows[i]
|
||
|
|
_bboxes = bbox_windows[i]
|
||
|
|
windows.append(
|
||
|
|
self.preprocess(
|
||
|
|
_bboxes, _words,
|
||
|
|
{'image': image_features, 'width': width, 'height': height, 'img_path': img_path},
|
||
|
|
max_seq_length=self.max_seq_length)
|
||
|
|
)
|
||
|
|
|
||
|
|
output['windows'] = windows
|
||
|
|
else:
|
||
|
|
raise ValueError(
|
||
|
|
f"Not supported mode: {self.mode }"
|
||
|
|
)
|
||
|
|
return output
|
||
|
|
|
||
|
|
|
||
|
|
def preprocess(self, bounding_boxes, words, feature_maps, max_seq_length):
|
||
|
|
input_ids_layoutxlm = np.ones(max_seq_length, dtype=int) * self.pad_token_id_layoutxlm
|
||
|
|
|
||
|
|
attention_mask_layoutxlm = np.zeros(max_seq_length, dtype=int)
|
||
|
|
|
||
|
|
bbox = np.zeros((max_seq_length, 8), dtype=np.float32)
|
||
|
|
|
||
|
|
are_box_first_tokens = np.zeros(max_seq_length, dtype=np.bool_)
|
||
|
|
|
||
|
|
list_layoutxlm_tokens = []
|
||
|
|
|
||
|
|
list_bbs = []
|
||
|
|
list_words = []
|
||
|
|
lwords = [''] * max_seq_length
|
||
|
|
|
||
|
|
box_to_token_indices = []
|
||
|
|
cum_token_idx = 0
|
||
|
|
|
||
|
|
cls_bbs = [0.0] * 8
|
||
|
|
len_overlap_tokens = 0
|
||
|
|
len_non_overlap_tokens = 0
|
||
|
|
len_valid_tokens = 0
|
||
|
|
|
||
|
|
for word_idx, (bounding_box, word) in enumerate(zip(bounding_boxes, words)):
|
||
|
|
bb = [[bounding_box[0], bounding_box[1]], [bounding_box[2], bounding_box[1]], [bounding_box[2], bounding_box[3]], [bounding_box[0], bounding_box[3]]]
|
||
|
|
layoutxlm_tokens = self.tokenizer_layoutxlm.convert_tokens_to_ids(self.tokenizer_layoutxlm.tokenize(word))
|
||
|
|
|
||
|
|
this_box_token_indices = []
|
||
|
|
|
||
|
|
len_valid_tokens += len(layoutxlm_tokens)
|
||
|
|
if word_idx < self.slice_interval:
|
||
|
|
len_non_overlap_tokens += len(layoutxlm_tokens)
|
||
|
|
|
||
|
|
if len(layoutxlm_tokens) == 0:
|
||
|
|
layoutxlm_tokens.append(self.unk_token_id)
|
||
|
|
|
||
|
|
|
||
|
|
if len(list_layoutxlm_tokens) + len(layoutxlm_tokens) > max_seq_length - 2:
|
||
|
|
break
|
||
|
|
|
||
|
|
|
||
|
|
list_layoutxlm_tokens += layoutxlm_tokens
|
||
|
|
|
||
|
|
# min, max clipping
|
||
|
|
for coord_idx in range(4):
|
||
|
|
bb[coord_idx][0] = max(0.0, min(bb[coord_idx][0], feature_maps['width']))
|
||
|
|
bb[coord_idx][1] = max(0.0, min(bb[coord_idx][1], feature_maps['height']))
|
||
|
|
|
||
|
|
bb = list(itertools.chain(*bb))
|
||
|
|
bbs = [bb for _ in range(len(layoutxlm_tokens))]
|
||
|
|
texts = [word for _ in range(len(layoutxlm_tokens))]
|
||
|
|
|
||
|
|
for _ in layoutxlm_tokens:
|
||
|
|
cum_token_idx += 1
|
||
|
|
this_box_token_indices.append(cum_token_idx)
|
||
|
|
|
||
|
|
list_bbs.extend(bbs)
|
||
|
|
list_words.extend(texts) ###
|
||
|
|
box_to_token_indices.append(this_box_token_indices)
|
||
|
|
|
||
|
|
sep_bbs = [feature_maps['width'], feature_maps['height']] * 4
|
||
|
|
|
||
|
|
# For [CLS] and [SEP]
|
||
|
|
list_layoutxlm_tokens = (
|
||
|
|
[self.cls_token_id_layoutxlm]
|
||
|
|
+ list_layoutxlm_tokens[: max_seq_length - 2]
|
||
|
|
+ [self.sep_token_id_layoutxlm]
|
||
|
|
)
|
||
|
|
|
||
|
|
if len(list_bbs) == 0:
|
||
|
|
# When len(json_obj["words"]) == 0 (no OCR result)
|
||
|
|
list_bbs = [cls_bbs] + [sep_bbs]
|
||
|
|
else: # len(list_bbs) > 0
|
||
|
|
list_bbs = [cls_bbs] + list_bbs[: max_seq_length - 2] + [sep_bbs]
|
||
|
|
# list_words = ['CLS'] + list_words[: max_seq_length - 2] + ['SEP']
|
||
|
|
list_words = [self.tokenizer_layoutxlm._cls_token] + list_words[: max_seq_length - 2] + [self.tokenizer_layoutxlm._sep_token]
|
||
|
|
|
||
|
|
len_list_layoutxlm_tokens = len(list_layoutxlm_tokens)
|
||
|
|
input_ids_layoutxlm[:len_list_layoutxlm_tokens] = list_layoutxlm_tokens
|
||
|
|
attention_mask_layoutxlm[:len_list_layoutxlm_tokens] = 1
|
||
|
|
|
||
|
|
|
||
|
|
bbox[:len_list_layoutxlm_tokens, :] = list_bbs
|
||
|
|
lwords[:len_list_layoutxlm_tokens] = list_words ###
|
||
|
|
|
||
|
|
# Normalize bbox -> 0 ~ 1
|
||
|
|
bbox[:, [0, 2, 4, 6]] = bbox[:, [0, 2, 4, 6]] / feature_maps['width']
|
||
|
|
bbox[:, [1, 3, 5, 7]] = bbox[:, [1, 3, 5, 7]] / feature_maps['height']
|
||
|
|
|
||
|
|
if self.backbone_type in ("layoutlm", "layoutxlm", "xlm-roberta"):
|
||
|
|
bbox = bbox[:, [0, 1, 4, 5]]
|
||
|
|
bbox = bbox * 1000
|
||
|
|
bbox = bbox.astype(int)
|
||
|
|
else:
|
||
|
|
assert False
|
||
|
|
|
||
|
|
st_indices = [
|
||
|
|
indices[0]
|
||
|
|
for indices in box_to_token_indices
|
||
|
|
if indices[0] < max_seq_length
|
||
|
|
]
|
||
|
|
are_box_first_tokens[st_indices] = True
|
||
|
|
|
||
|
|
assert len_list_layoutxlm_tokens == len_valid_tokens + 2
|
||
|
|
len_overlap_tokens = len_valid_tokens - len_non_overlap_tokens
|
||
|
|
|
||
|
|
ntokens = max_seq_length if max_seq_length == 512 else len_valid_tokens + 2
|
||
|
|
|
||
|
|
input_ids_layoutxlm = input_ids_layoutxlm[:ntokens]
|
||
|
|
attention_mask_layoutxlm = attention_mask_layoutxlm[:ntokens]
|
||
|
|
bbox = bbox[:ntokens]
|
||
|
|
are_box_first_tokens = are_box_first_tokens[:ntokens]
|
||
|
|
|
||
|
|
|
||
|
|
input_ids_layoutxlm = torch.from_numpy(input_ids_layoutxlm)
|
||
|
|
attention_mask_layoutxlm = torch.from_numpy(attention_mask_layoutxlm)
|
||
|
|
bbox = torch.from_numpy(bbox)
|
||
|
|
are_box_first_tokens = torch.from_numpy(are_box_first_tokens)
|
||
|
|
|
||
|
|
len_valid_tokens = torch.tensor(len_valid_tokens)
|
||
|
|
len_overlap_tokens = torch.tensor(len_overlap_tokens)
|
||
|
|
return_dict = {
|
||
|
|
"img_path": feature_maps['img_path'],
|
||
|
|
"words": list_words,
|
||
|
|
"len_overlap_tokens": len_overlap_tokens,
|
||
|
|
'len_valid_tokens': len_valid_tokens,
|
||
|
|
"image": feature_maps['image'],
|
||
|
|
"input_ids_layoutxlm": input_ids_layoutxlm,
|
||
|
|
"attention_mask_layoutxlm": attention_mask_layoutxlm,
|
||
|
|
"are_box_first_tokens": are_box_first_tokens,
|
||
|
|
"bbox": bbox,
|
||
|
|
}
|
||
|
|
return return_dict
|
||
|
|
|
||
|
|
def load_ground_truth(self, json_file):
|
||
|
|
json_obj = read_json(json_file)
|
||
|
|
width = json_obj["meta"]["imageSize"]["width"]
|
||
|
|
height = json_obj["meta"]["imageSize"]["height"]
|
||
|
|
|
||
|
|
input_ids = np.ones(self.max_seq_length, dtype=int) * self.pad_token_id_layoutxlm
|
||
|
|
bbox = np.zeros((self.max_seq_length, 8), dtype=np.float32)
|
||
|
|
attention_mask = np.zeros(self.max_seq_length, dtype=int)
|
||
|
|
|
||
|
|
itc_labels = np.zeros(self.max_seq_length, dtype=int)
|
||
|
|
are_box_first_tokens = np.zeros(self.max_seq_length, dtype=np.bool_)
|
||
|
|
|
||
|
|
# stc_labels stores the index of the previous token.
|
||
|
|
# A stored index of max_seq_length (512) indicates that
|
||
|
|
# this token is the initial token of a word box.
|
||
|
|
stc_labels = np.ones(self.max_seq_length, dtype=np.int64) * self.max_seq_length
|
||
|
|
el_labels = np.ones(self.max_seq_length, dtype=int) * self.max_seq_length
|
||
|
|
el_labels_from_key = np.ones(self.max_seq_length, dtype=int) * self.max_seq_length
|
||
|
|
|
||
|
|
|
||
|
|
list_tokens = []
|
||
|
|
list_bbs = []
|
||
|
|
list_words = []
|
||
|
|
box2token_span_map = []
|
||
|
|
lwords = [''] * self.max_seq_length
|
||
|
|
|
||
|
|
box_to_token_indices = []
|
||
|
|
cum_token_idx = 0
|
||
|
|
|
||
|
|
cls_bbs = [0.0] * 8
|
||
|
|
|
||
|
|
for word_idx, word in enumerate(json_obj["words"]):
|
||
|
|
this_box_token_indices = []
|
||
|
|
|
||
|
|
tokens = word["layoutxlm_tokens"]
|
||
|
|
bb = word["boundingBox"]
|
||
|
|
text = word["text"]
|
||
|
|
|
||
|
|
if len(tokens) == 0:
|
||
|
|
tokens.append(self.unk_token_id)
|
||
|
|
|
||
|
|
if len(list_tokens) + len(tokens) > self.max_seq_length - 2:
|
||
|
|
break
|
||
|
|
|
||
|
|
box2token_span_map.append(
|
||
|
|
[len(list_tokens) + 1, len(list_tokens) + len(tokens) + 1]
|
||
|
|
) # including st_idx
|
||
|
|
list_tokens += tokens
|
||
|
|
|
||
|
|
# min, max clipping
|
||
|
|
for coord_idx in range(4):
|
||
|
|
bb[coord_idx][0] = max(0.0, min(bb[coord_idx][0], width))
|
||
|
|
bb[coord_idx][1] = max(0.0, min(bb[coord_idx][1], height))
|
||
|
|
|
||
|
|
bb = list(itertools.chain(*bb))
|
||
|
|
bbs = [bb for _ in range(len(tokens))]
|
||
|
|
texts = [text for _ in range(len(tokens))]
|
||
|
|
|
||
|
|
for _ in tokens:
|
||
|
|
cum_token_idx += 1
|
||
|
|
this_box_token_indices.append(cum_token_idx)
|
||
|
|
|
||
|
|
list_bbs.extend(bbs)
|
||
|
|
list_words.extend(texts) ####
|
||
|
|
box_to_token_indices.append(this_box_token_indices)
|
||
|
|
|
||
|
|
sep_bbs = [width, height] * 4
|
||
|
|
|
||
|
|
# For [CLS] and [SEP]
|
||
|
|
list_tokens = (
|
||
|
|
[self.cls_token_id_layoutxlm]
|
||
|
|
+ list_tokens[: self.max_seq_length - 2]
|
||
|
|
+ [self.sep_token_id_layoutxlm]
|
||
|
|
)
|
||
|
|
if len(list_bbs) == 0:
|
||
|
|
# When len(json_obj["words"]) == 0 (no OCR result)
|
||
|
|
list_bbs = [cls_bbs] + [sep_bbs]
|
||
|
|
else: # len(list_bbs) > 0
|
||
|
|
list_bbs = [cls_bbs] + list_bbs[: self.max_seq_length - 2] + [sep_bbs]
|
||
|
|
# list_words = ['CLS'] + list_words[: self.max_seq_length - 2] + ['SEP'] ###
|
||
|
|
list_words = [self.tokenizer_layoutxlm._cls_token] + list_words[: self.max_seq_length - 2] + [self.tokenizer_layoutxlm._sep_token]
|
||
|
|
|
||
|
|
|
||
|
|
len_list_tokens = len(list_tokens)
|
||
|
|
input_ids[:len_list_tokens] = list_tokens
|
||
|
|
attention_mask[:len_list_tokens] = 1
|
||
|
|
|
||
|
|
bbox[:len_list_tokens, :] = list_bbs
|
||
|
|
lwords[:len_list_tokens] = list_words
|
||
|
|
|
||
|
|
# Normalize bbox -> 0 ~ 1
|
||
|
|
bbox[:, [0, 2, 4, 6]] = bbox[:, [0, 2, 4, 6]] / width
|
||
|
|
bbox[:, [1, 3, 5, 7]] = bbox[:, [1, 3, 5, 7]] / height
|
||
|
|
|
||
|
|
if self.backbone_type in ("layoutlm", "layoutxlm"):
|
||
|
|
bbox = bbox[:, [0, 1, 4, 5]]
|
||
|
|
bbox = bbox * 1000
|
||
|
|
bbox = bbox.astype(int)
|
||
|
|
else:
|
||
|
|
assert False
|
||
|
|
|
||
|
|
st_indices = [
|
||
|
|
indices[0]
|
||
|
|
for indices in box_to_token_indices
|
||
|
|
if indices[0] < self.max_seq_length
|
||
|
|
]
|
||
|
|
are_box_first_tokens[st_indices] = True
|
||
|
|
|
||
|
|
# Label
|
||
|
|
classes_dic = json_obj["parse"]["class"]
|
||
|
|
for class_name in self.class_names:
|
||
|
|
if class_name == "others":
|
||
|
|
continue
|
||
|
|
if class_name not in classes_dic:
|
||
|
|
continue
|
||
|
|
|
||
|
|
for word_list in classes_dic[class_name]:
|
||
|
|
is_first, last_word_idx = True, -1
|
||
|
|
for word_idx in word_list:
|
||
|
|
if word_idx >= len(box_to_token_indices):
|
||
|
|
break
|
||
|
|
box2token_list = box_to_token_indices[word_idx]
|
||
|
|
for converted_word_idx in box2token_list:
|
||
|
|
if converted_word_idx >= self.max_seq_length:
|
||
|
|
break # out of idx
|
||
|
|
|
||
|
|
if is_first:
|
||
|
|
itc_labels[converted_word_idx] = self.class_idx_dic[
|
||
|
|
class_name
|
||
|
|
]
|
||
|
|
is_first, last_word_idx = False, converted_word_idx
|
||
|
|
else:
|
||
|
|
stc_labels[converted_word_idx] = last_word_idx
|
||
|
|
last_word_idx = converted_word_idx
|
||
|
|
|
||
|
|
# Label
|
||
|
|
relations = json_obj["parse"]["relations"]
|
||
|
|
for relation in relations:
|
||
|
|
if relation[0] >= len(box2token_span_map) or relation[1] >= len(
|
||
|
|
box2token_span_map
|
||
|
|
):
|
||
|
|
continue
|
||
|
|
if (
|
||
|
|
box2token_span_map[relation[0]][0] >= self.max_seq_length
|
||
|
|
or box2token_span_map[relation[1]][0] >= self.max_seq_length
|
||
|
|
):
|
||
|
|
continue
|
||
|
|
|
||
|
|
word_from = box2token_span_map[relation[0]][0]
|
||
|
|
word_to = box2token_span_map[relation[1]][0]
|
||
|
|
# el_labels[word_to] = word_from
|
||
|
|
|
||
|
|
if el_labels[word_to] != 512 and el_labels_from_key[word_to] != 512:
|
||
|
|
continue
|
||
|
|
|
||
|
|
# if self.second_relations == 1:
|
||
|
|
# if itc_labels[word_from] == 4 and (itc_labels[word_to] in (2, 3)):
|
||
|
|
# el_labels[word_to] = word_from # pair of (header, key) or (header-value)
|
||
|
|
# else:
|
||
|
|
#### 1st relation => ['key, 'value']
|
||
|
|
#### 2st relation => ['header', 'key'or'value']
|
||
|
|
if itc_labels[word_from] == 2 and itc_labels[word_to] == 3:
|
||
|
|
el_labels_from_key[word_to] = word_from # pair of (key-value)
|
||
|
|
if itc_labels[word_from] == 4 and (itc_labels[word_to] in (2, 3)):
|
||
|
|
el_labels[word_to] = word_from # pair of (header, key) or (header-value)
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
input_ids = torch.from_numpy(input_ids)
|
||
|
|
bbox = torch.from_numpy(bbox)
|
||
|
|
attention_mask = torch.from_numpy(attention_mask)
|
||
|
|
|
||
|
|
itc_labels = torch.from_numpy(itc_labels)
|
||
|
|
are_box_first_tokens = torch.from_numpy(are_box_first_tokens)
|
||
|
|
stc_labels = torch.from_numpy(stc_labels)
|
||
|
|
el_labels = torch.from_numpy(el_labels)
|
||
|
|
el_labels_from_key = torch.from_numpy(el_labels_from_key)
|
||
|
|
|
||
|
|
return_dict = {
|
||
|
|
# "image": feature_maps,
|
||
|
|
"input_ids": input_ids,
|
||
|
|
"bbox": bbox,
|
||
|
|
"words": lwords,
|
||
|
|
"attention_mask": attention_mask,
|
||
|
|
"itc_labels": itc_labels,
|
||
|
|
"are_box_first_tokens": are_box_first_tokens,
|
||
|
|
"stc_labels": stc_labels,
|
||
|
|
"el_labels": el_labels,
|
||
|
|
"el_labels_from_key": el_labels_from_key
|
||
|
|
}
|
||
|
|
|
||
|
|
return return_dict
|
||
|
|
|
||
|
|
|
||
|
|
class DocumentKVUProcess(KVUProcess):
|
||
|
|
def __init__(self, tokenizer_layoutxlm, feature_extractor, backbone_type, class_names, max_window_count, slice_interval, window_size, run_ocr, max_seq_length=512, mode=0):
|
||
|
|
super().__init__(tokenizer_layoutxlm, feature_extractor, backbone_type, class_names, slice_interval, window_size, run_ocr, max_seq_length, mode)
|
||
|
|
self.max_window_count = max_window_count
|
||
|
|
self.pad_token_id = self.pad_token_id_layoutxlm
|
||
|
|
self.cls_token_id = self.cls_token_id_layoutxlm
|
||
|
|
self.sep_token_id = self.sep_token_id_layoutxlm
|
||
|
|
self.unk_token_id = self.unk_token_id_layoutxlm
|
||
|
|
self.tokenizer = self.tokenizer_layoutxlm
|
||
|
|
|
||
|
|
def __call__(self, img_path: str, ocr_path: str) -> list:
|
||
|
|
if (self.run_ocr == 1) and (not os.path.exists(ocr_path)):
|
||
|
|
process_img(img_path, "tmp.txt", self.ocr_engine, export_img=False)
|
||
|
|
ocr_path = "tmp.txt"
|
||
|
|
lbboxes, lwords = read_ocr_result_from_txt(ocr_path)
|
||
|
|
lwords = post_process_basic_ocr(lwords)
|
||
|
|
|
||
|
|
width, height = imagesize.get(img_path)
|
||
|
|
images = [Image.open(img_path).convert("RGB")]
|
||
|
|
image_features = torch.from_numpy(self.feature_extractor(images)['pixel_values'][0].copy())
|
||
|
|
output = self.preprocess(lbboxes, lwords,
|
||
|
|
{'image': image_features, 'width': width, 'height': height, 'img_path': img_path})
|
||
|
|
return output
|
||
|
|
|
||
|
|
def preprocess(self, bounding_boxes, words, feature_maps):
|
||
|
|
n_words = len(words)
|
||
|
|
output_dicts = {'windows': [], 'documents': []}
|
||
|
|
n_empty_windows = 0
|
||
|
|
|
||
|
|
for i in range(self.max_window_count):
|
||
|
|
input_ids = np.ones(self.max_seq_length, dtype=int) * self.pad_token_id
|
||
|
|
bbox = np.zeros((self.max_seq_length, 8), dtype=np.float32)
|
||
|
|
attention_mask = np.zeros(self.max_seq_length, dtype=int)
|
||
|
|
are_box_first_tokens = np.zeros(self.max_seq_length, dtype=np.bool_)
|
||
|
|
|
||
|
|
if n_words == 0:
|
||
|
|
n_empty_windows += 1
|
||
|
|
output_dicts['windows'].append({
|
||
|
|
"image": feature_maps['image'],
|
||
|
|
"input_ids": torch.from_numpy(input_ids),
|
||
|
|
"bbox": torch.from_numpy(bbox),
|
||
|
|
"words": [],
|
||
|
|
"attention_mask": torch.from_numpy(attention_mask),
|
||
|
|
"are_box_first_tokens": torch.from_numpy(are_box_first_tokens),
|
||
|
|
})
|
||
|
|
continue
|
||
|
|
|
||
|
|
start_word_idx = i * self.window_size
|
||
|
|
stop_word_idx = min(n_words, (i+1)*self.window_size)
|
||
|
|
|
||
|
|
if start_word_idx >= stop_word_idx:
|
||
|
|
n_empty_windows += 1
|
||
|
|
output_dicts['windows'].append(output_dicts['windows'][-1])
|
||
|
|
continue
|
||
|
|
|
||
|
|
list_tokens = []
|
||
|
|
list_bbs = []
|
||
|
|
list_words = []
|
||
|
|
lwords = [''] * self.max_seq_length
|
||
|
|
|
||
|
|
box_to_token_indices = []
|
||
|
|
cum_token_idx = 0
|
||
|
|
|
||
|
|
cls_bbs = [0.0] * 8
|
||
|
|
|
||
|
|
for _, (bounding_box, word) in enumerate(zip(bounding_boxes[start_word_idx:stop_word_idx], words[start_word_idx:stop_word_idx])):
|
||
|
|
bb = [[bounding_box[0], bounding_box[1]], [bounding_box[2], bounding_box[1]], [bounding_box[2], bounding_box[3]], [bounding_box[0], bounding_box[3]]]
|
||
|
|
tokens = self.tokenizer.convert_tokens_to_ids(self.tokenizer.tokenize(word))
|
||
|
|
|
||
|
|
this_box_token_indices = []
|
||
|
|
|
||
|
|
if len(tokens) == 0:
|
||
|
|
tokens.append(self.unk_token_id)
|
||
|
|
|
||
|
|
if len(list_tokens) + len(tokens) > self.max_seq_length - 2:
|
||
|
|
break
|
||
|
|
|
||
|
|
list_tokens += tokens
|
||
|
|
|
||
|
|
# min, max clipping
|
||
|
|
for coord_idx in range(4):
|
||
|
|
bb[coord_idx][0] = max(0.0, min(bb[coord_idx][0], feature_maps['width']))
|
||
|
|
bb[coord_idx][1] = max(0.0, min(bb[coord_idx][1], feature_maps['height']))
|
||
|
|
|
||
|
|
bb = list(itertools.chain(*bb))
|
||
|
|
bbs = [bb for _ in range(len(tokens))]
|
||
|
|
texts = [word for _ in range(len(tokens))]
|
||
|
|
|
||
|
|
for _ in tokens:
|
||
|
|
cum_token_idx += 1
|
||
|
|
this_box_token_indices.append(cum_token_idx)
|
||
|
|
|
||
|
|
list_bbs.extend(bbs)
|
||
|
|
list_words.extend(texts) ###
|
||
|
|
box_to_token_indices.append(this_box_token_indices)
|
||
|
|
|
||
|
|
sep_bbs = [feature_maps['width'], feature_maps['height']] * 4
|
||
|
|
|
||
|
|
# For [CLS] and [SEP]
|
||
|
|
list_tokens = (
|
||
|
|
[self.cls_token_id]
|
||
|
|
+ list_tokens[: self.max_seq_length - 2]
|
||
|
|
+ [self.sep_token_id]
|
||
|
|
)
|
||
|
|
if len(list_bbs) == 0:
|
||
|
|
# When len(json_obj["words"]) == 0 (no OCR result)
|
||
|
|
list_bbs = [cls_bbs] + [sep_bbs]
|
||
|
|
else: # len(list_bbs) > 0
|
||
|
|
list_bbs = [cls_bbs] + list_bbs[: self.max_seq_length - 2] + [sep_bbs]
|
||
|
|
if len(list_words) < 510:
|
||
|
|
list_words.extend(['</p>' for _ in range(510 - len(list_words))])
|
||
|
|
list_words = [self.tokenizer._cls_token] + list_words[: self.max_seq_length - 2] + [self.tokenizer._sep_token]
|
||
|
|
|
||
|
|
len_list_tokens = len(list_tokens)
|
||
|
|
input_ids[:len_list_tokens] = list_tokens
|
||
|
|
attention_mask[:len_list_tokens] = 1
|
||
|
|
|
||
|
|
bbox[:len_list_tokens, :] = list_bbs
|
||
|
|
lwords[:len_list_tokens] = list_words ###
|
||
|
|
|
||
|
|
# Normalize bbox -> 0 ~ 1
|
||
|
|
bbox[:, [0, 2, 4, 6]] = bbox[:, [0, 2, 4, 6]] / feature_maps['width']
|
||
|
|
bbox[:, [1, 3, 5, 7]] = bbox[:, [1, 3, 5, 7]] / feature_maps['height']
|
||
|
|
|
||
|
|
bbox = bbox[:, [0, 1, 4, 5]]
|
||
|
|
bbox = bbox * 1000
|
||
|
|
bbox = bbox.astype(int)
|
||
|
|
|
||
|
|
st_indices = [
|
||
|
|
indices[0]
|
||
|
|
for indices in box_to_token_indices
|
||
|
|
if indices[0] < self.max_seq_length
|
||
|
|
]
|
||
|
|
are_box_first_tokens[st_indices] = True
|
||
|
|
|
||
|
|
input_ids = torch.from_numpy(input_ids)
|
||
|
|
bbox = torch.from_numpy(bbox)
|
||
|
|
attention_mask = torch.from_numpy(attention_mask)
|
||
|
|
are_box_first_tokens = torch.from_numpy(are_box_first_tokens)
|
||
|
|
|
||
|
|
return_dict = {
|
||
|
|
"image": feature_maps['image'],
|
||
|
|
"input_ids": input_ids,
|
||
|
|
"bbox": bbox,
|
||
|
|
"words": list_words,
|
||
|
|
"attention_mask": attention_mask,
|
||
|
|
"are_box_first_tokens": are_box_first_tokens,
|
||
|
|
}
|
||
|
|
output_dicts["windows"].append(return_dict)
|
||
|
|
|
||
|
|
attention_mask = torch.cat([o['attention_mask'] for o in output_dicts["windows"]])
|
||
|
|
are_box_first_tokens = torch.cat([o['are_box_first_tokens'] for o in output_dicts["windows"]])
|
||
|
|
if n_empty_windows > 0:
|
||
|
|
attention_mask[self.max_seq_length * (self.max_window_count - n_empty_windows):] = torch.from_numpy(np.zeros(self.max_seq_length * n_empty_windows, dtype=int))
|
||
|
|
are_box_first_tokens[self.max_seq_length * (self.max_window_count - n_empty_windows):] = torch.from_numpy(np.zeros(self.max_seq_length * n_empty_windows, dtype=np.bool_))
|
||
|
|
bbox = torch.cat([o['bbox'] for o in output_dicts["windows"]])
|
||
|
|
words = []
|
||
|
|
for o in output_dicts['windows']:
|
||
|
|
words.extend(o['words'])
|
||
|
|
|
||
|
|
return_dict = {
|
||
|
|
"attention_mask": attention_mask,
|
||
|
|
"bbox": bbox,
|
||
|
|
"are_box_first_tokens": are_box_first_tokens,
|
||
|
|
"n_empty_windows": n_empty_windows,
|
||
|
|
"words": words
|
||
|
|
}
|
||
|
|
output_dicts['documents'] = return_dict
|
||
|
|
|
||
|
|
return output_dicts
|
||
|
|
|
||
|
|
|
||
|
|
|