layoutlmv3FineTuning/Layoutlm_inference/utils.py

import numpy as np
from transformers import AutoModelForTokenClassification, AutoProcessor

from dotenv import load_dotenv
import os

# Load .env file
load_dotenv()

# Access variables
dummy_key = os.getenv("dummy_key")
# secret_key = os.getenv("SECRET_KEY")
# debug_mode = os.getenv("DEBUG")

# print(f"Database URL: {database_url}")
# print(f"Secret Key: {secret_key}")
# print(f"Debug Mode: {debug_mode}")




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,use_auth_token=dummy_key)
    return model

def load_processor(model_name_or_path):
    processor = AutoProcessor.from_pretrained(
        model_name_or_path, apply_ocr=False,use_auth_token=dummy_key)
    return processor