|
import numpy as np |
|
from preprocess import normalize_box |
|
import copy |
|
|
|
def classifyTokens(model, input_ids, attention_mask, bbox, offset_mapping): |
|
outputs = model(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask) |
|
|
|
predictions = outputs.logits.argmax(-1).squeeze().tolist() |
|
return predictions |
|
|
|
def compare_boxes(b1,b2): |
|
b1 = np.array([c for c in b1]) |
|
b2 = np.array([c for c in b2]) |
|
equal = np.array_equal(b1,b2) |
|
return equal |
|
|
|
def mergable(w1,w2): |
|
if w1['label'] == w2['label']: |
|
threshold = 7 |
|
if abs(w1['box'][1] - w2['box'][1]) < threshold or abs(w1['box'][-1] - w2['box'][-1]) < threshold: |
|
return True |
|
return False |
|
return False |
|
|
|
def convert_data(data, tokenizer, img_size): |
|
def normalize_bbox(bbox, size): |
|
return [ |
|
int(1000 * bbox[0] / size[0]), |
|
int(1000 * bbox[1] / size[1]), |
|
int(1000 * bbox[2] / size[0]), |
|
int(1000 * bbox[3] / size[1]), |
|
] |
|
|
|
|
|
def simplify_bbox(bbox): |
|
return [ |
|
min(bbox[0::2]), |
|
min(bbox[1::2]), |
|
max(bbox[2::2]), |
|
max(bbox[3::2]), |
|
] |
|
|
|
def merge_bbox(bbox_list): |
|
x0, y0, x1, y1 = list(zip(*bbox_list)) |
|
return [min(x0), min(y0), max(x1), max(y1)] |
|
|
|
tokenized_doc = {"input_ids": [], "bbox": [], "labels": [], "attention_mask":[]} |
|
entities = [] |
|
id2label = {} |
|
entity_id_to_index_map = {} |
|
empty_entity = set() |
|
for line in data: |
|
if len(line["text"]) == 0: |
|
empty_entity.add(line["id"]) |
|
continue |
|
id2label[line["id"]] = line["label"] |
|
tokenized_inputs = tokenizer( |
|
line["text"], |
|
add_special_tokens=False, |
|
return_offsets_mapping=True, |
|
return_attention_mask=True, |
|
) |
|
text_length = 0 |
|
ocr_length = 0 |
|
bbox = [] |
|
for token_id, offset in zip(tokenized_inputs["input_ids"], tokenized_inputs["offset_mapping"]): |
|
if token_id == 6: |
|
bbox.append(None) |
|
continue |
|
text_length += offset[1] - offset[0] |
|
tmp_box = [] |
|
while ocr_length < text_length: |
|
ocr_word = line["words"].pop(0) |
|
ocr_length += len( |
|
tokenizer._tokenizer.normalizer.normalize_str(ocr_word["text"].strip()) |
|
) |
|
tmp_box.append(simplify_bbox(ocr_word["box"])) |
|
if len(tmp_box) == 0: |
|
tmp_box = last_box |
|
bbox.append(normalize_bbox(merge_bbox(tmp_box), img_size)) |
|
last_box = tmp_box |
|
bbox = [ |
|
[bbox[i + 1][0], bbox[i + 1][1], bbox[i + 1][0], bbox[i + 1][1]] if b is None else b |
|
for i, b in enumerate(bbox) |
|
] |
|
if line["label"] == "other": |
|
label = ["O"] * len(bbox) |
|
else: |
|
label = [f"I-{line['label'].upper()}"] * len(bbox) |
|
label[0] = f"B-{line['label'].upper()}" |
|
tokenized_inputs.update({"bbox": bbox, "labels": label}) |
|
if label[0] != "O": |
|
entity_id_to_index_map[line["id"]] = len(entities) |
|
entities.append( |
|
{ |
|
"start": len(tokenized_doc["input_ids"]), |
|
"end": len(tokenized_doc["input_ids"]) + len(tokenized_inputs["input_ids"]), |
|
"label": line["label"].upper(), |
|
} |
|
) |
|
for i in tokenized_doc: |
|
tokenized_doc[i] = tokenized_doc[i] + tokenized_inputs[i] |
|
|
|
chunk_size = 512 |
|
output = {} |
|
for chunk_id, index in enumerate(range(0, len(tokenized_doc["input_ids"]), chunk_size)): |
|
item = {} |
|
entities_in_this_span = [] |
|
for k in tokenized_doc: |
|
item[k] = tokenized_doc[k][index : index + chunk_size] |
|
global_to_local_map = {} |
|
for entity_id, entity in enumerate(entities): |
|
if ( |
|
index <= entity["start"] < index + chunk_size |
|
and index <= entity["end"] < index + chunk_size |
|
): |
|
entity["start"] = entity["start"] - index |
|
entity["end"] = entity["end"] - index |
|
global_to_local_map[entity_id] = len(entities_in_this_span) |
|
entities_in_this_span.append(entity) |
|
item.update( |
|
{ |
|
"entities": entities_in_this_span |
|
} |
|
) |
|
for key in item.keys(): |
|
output[key] = output.get(key, []) + item[key] |
|
return output |
|
|
|
def dfs(i, merged, width, height, visited, df_words): |
|
v_threshold = int(.01 * height) |
|
h_threshold = int(.08 * width) |
|
visited.add(i) |
|
merged.append(df_words[i]) |
|
|
|
for j in range(len(df_words)): |
|
if j not in visited: |
|
w1 = df_words[i]['words'][0] |
|
w2 = df_words[j]['words'][0] |
|
|
|
|
|
if (abs(w1['box'][1] - w2['box'][1]) < v_threshold or abs(w1['box'][-1] - w2['box'][-1]) < v_threshold) \ |
|
and (df_words[i]['label'] == df_words[j]['label']) \ |
|
and (abs(w1['box'][0] - w2['box'][0]) < h_threshold or abs(w1['box'][-2] - w2['box'][-2]) < h_threshold): |
|
dfs(j,merged, width, height, visited, df_words) |
|
return merged |
|
|
|
def createEntities(model, predictions, input_ids, ocr_df, tokenizer, img_size, bbox): |
|
width, height = img_size |
|
words = [] |
|
for index,row in ocr_df.iterrows(): |
|
word = {} |
|
origin_box = [row['left'],row['top'],row['left']+row['width'],row['top']+row['height']] |
|
word['word_text'] = row['text'] |
|
word['word_box'] = origin_box |
|
word['normalized_box'] = normalize_box(word['word_box'], width, height) |
|
words.append(word) |
|
|
|
raw_input_ids = input_ids[0].tolist() |
|
token_boxes = bbox.squeeze().tolist() |
|
special_tokens = [tokenizer.cls_token_id, tokenizer.sep_token_id, tokenizer.pad_token_id] |
|
|
|
input_ids = [id for id in raw_input_ids if id not in special_tokens] |
|
predictions = [model.config.id2label[prediction] for i,prediction in enumerate(predictions) if not (raw_input_ids[i] in special_tokens)] |
|
actual_boxes = [box for i,box in enumerate(token_boxes) if not (raw_input_ids[i] in special_tokens )] |
|
|
|
assert(len(actual_boxes) == len(predictions)) |
|
|
|
for word in words: |
|
word_labels = [] |
|
token_labels = [] |
|
word_tagging = None |
|
for i,box in enumerate(actual_boxes,start=0): |
|
if compare_boxes(word['normalized_box'],box): |
|
if predictions[i] != 'O': |
|
word_labels.append(predictions[i][2:]) |
|
else: |
|
word_labels.append('O') |
|
token_labels.append(predictions[i]) |
|
if word_labels != []: |
|
word_tagging = word_labels[0] if word_labels[0] != 'O' else word_labels[-1] |
|
else: |
|
word_tagging = 'O' |
|
word['word_labels'] = token_labels |
|
word['word_tagging'] = word_tagging |
|
|
|
filtered_words = [{'id':i,'text':word['word_text'], |
|
'label':word['word_tagging'], |
|
'box':word['word_box'], |
|
'words':[{'box':word['word_box'],'text':word['word_text']}]} for i,word in enumerate(words) if word['word_tagging'] != 'O'] |
|
|
|
merged_taggings = [] |
|
df_words = filtered_words.copy() |
|
visited = set() |
|
for i in range(len(df_words)): |
|
if i not in visited: |
|
merged_taggings.append(dfs(i,[], width, height, visited, df_words)) |
|
merged_words = [] |
|
for i,merged_tagging in enumerate(merged_taggings): |
|
if ((len(merged_tagging) > 1)) or (merged_tagging[0]['label'] == 'ANSWER'): |
|
new_word = {} |
|
merging_word = " ".join([word['text'] for word in merged_tagging]) |
|
merging_box = [merged_tagging[0]['box'][0]-5,merged_tagging[0]['box'][1]-10,merged_tagging[-1]['box'][2]+5,merged_tagging[-1]['box'][3]+10] |
|
new_word['text'] = merging_word |
|
new_word['box'] = merging_box |
|
new_word['label'] = merged_tagging[0]['label'] |
|
new_word['id'] = filtered_words[-1]['id']+i+1 |
|
new_word['words'] = [{'box':word['box'],'text':word['text']} for word in merged_tagging] |
|
|
|
merged_words.append(new_word) |
|
|
|
filtered_words.extend(merged_words) |
|
predictions = [word['label'] for word in filtered_words] |
|
actual_boxes = [word['box'] for word in filtered_words] |
|
unique_taggings = set(predictions) |
|
|
|
output = convert_data(copy.deepcopy(merged_words), tokenizer, img_size) |
|
return output, merged_words |
|
|
|
|
