File size: 1,475 Bytes
6df3c38 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import numpy as np
from transformers import AutoModelForTokenClassification, AutoProcessor
def normalize_box(bbox, width, height):
return [
int(bbox[0]*(1000/width)),
int(bbox[1]*(1000/height)),
int(bbox[2]*(1000/width)),
int(bbox[3]*(1000/height)),
]
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 unnormalize_box(bbox, width, height):
return [
width * (bbox[0] / 1000),
height * (bbox[1] / 1000),
width * (bbox[2] / 1000),
height * (bbox[3] / 1000),
]
def adjacent(w1, w2):
if w1['label'] == w2['label'] and abs(w1['id'] - w2['id']) == 1:
return True
return False
def random_color():
return np.random.randint(0, 255, 3)
def image_label_2_color(annotation):
if 'output' in annotation.keys():
image_labels = set([span['label'] for span in annotation['output']])
label2color = {f'{label}': (random_color()[0], random_color()[
1], random_color()[2]) for label in image_labels}
return label2color
else:
raise ValueError('please use "output" as annotation key')
def load_model(model_path):
model = AutoModelForTokenClassification.from_pretrained(model_path)
return model
def load_processor():
processor = AutoProcessor.from_pretrained(
"microsoft/layoutlmv3-base", apply_ocr=False)
return processor
|