Sarath2002/Form_Understanding_using_LayoutLMV3

1from datasets import load_dataset2import numpy as np3from transformers import LayoutLMv3Processor, LayoutLMv3ForTokenClassification4from datasets import load_dataset5from PIL import Image, ImageDraw, ImageFont6import torch7 8 9 10tokenizer = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base")11model = LayoutLMv3ForTokenClassification.from_pretrained(r"models")12"""device = torch.device("cuda")13model.cuda()14"""15labels = ['O', 'B-HEADER', 'I-HEADER', 'B-QUESTION', 'I-QUESTION', 'B-ANSWER', 'I-ANSWER']16id2label = {v: k for v, k in enumerate(labels)}17label2color = {18    "question": "blue",19    "answer": "green",20    "header": "orange",21    "other": "violet",22}23 24 25def unnormalize_box(bbox, width, height):26    return [27        width * (bbox[0] / 1000),28        height * (bbox[1] / 1000),29        width * (bbox[2] / 1000),30        height * (bbox[3] / 1000),31    ]32 33 34def iob_to_label(label):35    label = label[2:]36    if not label:37        return "other"38    return label39 40 41def processor(image):42    image = image.convert("RGB")43    width, height = image.size44    45 46    # encode47    encoding = tokenizer(48        image, truncation=True, return_offsets_mapping=True, return_tensors="pt"49    )50    offset_mapping = encoding.pop("offset_mapping")51    52    encoding = encoding.to('cuda')53 54    # forward pass55    outputs = model(**encoding)56    57    # get predictions58    predictions = outputs.logits.argmax(-1).squeeze().tolist()59    token_boxes = encoding.bbox.squeeze().tolist()60    61 62    # only keep non-subword predictions63    is_subword = np.array(offset_mapping.squeeze().tolist())[:, 0] != 064    true_predictions = [65        id2label[pred] for idx, pred in enumerate(predictions) if not is_subword[idx]66    ]67    true_boxes = [68        unnormalize_box(box, width, height)69        for idx, box in enumerate(token_boxes)70        if not is_subword[idx]71    ]72    73    74    75    draw = ImageDraw.Draw(image)76    font = ImageFont.load_default()77    for prediction, box in zip(true_predictions, true_boxes):78        predicted_label = iob_to_label(prediction).lower()79        draw.rectangle(box, outline=label2color[predicted_label])80        draw.text(81            (box[0] + 10, box[1] - 10),82            text=predicted_label,83            fill=label2color[predicted_label],84            font=font,85        )86 87    return image