# -*- coding: utf-8 -*- """inference with LayoutLMv2ForTokenClassification .ipynb Automatically generated by Colaboratory. Original file is located at https://colab.research.google.com/drive/1nhfx6XRncq2XsOBREZGJI7tRByt_TJIa ## Inference with LayoutLMv2ForTokenClassification + Gradio demo In this notebook, we are going to perform inference with `LayoutLMv2ForTokenClassification` on new document images, when no label information is accessible. At the end, we will also make a cool [Gradio](https://gradio.app/) demo, that turns our inference code into a cool web interface. ## Install libraries Let's first install the required libraries: * HuggingFace Transformers + Detectron2 (for the model) * HuggingFace Datasets (for getting the data) * PyTesseract (for OCR) """ # !pip install -q transformers # !pip install -q gradio # !pip install 'git+https://github.com/facebookresearch/detectron2.git' # !pip install -q datasets # !sudo apt install tesseract-ocr # !pip install -q pytesseract # pip install torchvision # import gradio as gr import os import time import numpy as np from transformers import LayoutLMv2Processor, LayoutLMv2ForTokenClassification from datasets import load_dataset import torch from transformers import LayoutLMv2ForTokenClassification from PIL import Image, ImageDraw, ImageFont import json from GoogleVisionService import GoogleVisionService from getTextHelper import cord_label_to_color, get_word_boxes_google, get_word_boxes_tesseract, getImg, getImgAndPath, normalize_bbox, unnormalize_box from datasets import load_dataset import pytesseract import cv2 class labelCounter: lbl_i = 0 lbl = None device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') apply_ocr = False # apply_ocr=True processor = LayoutLMv2Processor.from_pretrained( "microsoft/layoutlmv2-base-uncased", apply_ocr=apply_ocr) workingPath = "/Users/eliaweiss/Documents/doc2txt/lineCv/" docNumberList = [6, 7, 10, 21, 25, 29, 48, 50] # load the fine-tuned model from the hub model = LayoutLMv2ForTokenClassification.from_pretrained( "doc2txt/layoutlmv2-finetuned-cord") model.to(device) # datasets = load_dataset("MarkusDressel/cord") """Let's create a list containing all unique labels, as well as dictionaries mapping integers to their label names and vice versa. This will be useful to convert the model's predictions to actual label names.""" # labels = datasets['train'].features['ner_tags'].feature.names labels = ['I-menu.cnt', 'I-menu.discountprice', 'I-menu.etc', 'I-menu.itemsubtotal', 'I-menu.nm', 'I-menu.num', 'I-menu.price', 'I-menu.sub_cnt', 'I-menu.sub_etc', 'I-menu.sub_nm', 'I-menu.sub_price', 'I-menu.sub_unitprice', 'I-menu.unitprice', 'I-menu.vatyn', 'I-sub_total.discount_price', 'I-sub_total.etc', 'I-sub_total.othersvc_price', 'I-sub_total.service_price', 'I-sub_total.subtotal_price', 'I-sub_total.tax_price', 'I-total.cashprice', 'I-total.changeprice', 'I-total.creditcardprice', 'I-total.emoneyprice', 'I-total.menuqty_cnt', 'I-total.menutype_cnt', 'I-total.total_etc', 'I-total.total_price', 'I-void_menu.nm', 'I-void_menu.price'] # print(labels) id2label = {v: k for v, k in enumerate(labels)} label2id = {k: v for v, k in enumerate(labels)} """## Inference """ # font = ImageFont.load_default() font = ImageFont.truetype( "/Users/eliaweiss/work/ocrPlus/ocrPlus/DejaVuSans.ttf", 50) # print(example.keys()) # for nn in docNumberList: def getNextLabel(labels, true_predictions): labelCounter.lbl = None i = 0 while labelCounter.lbl not in true_predictions: labelCounter.lbl_i += 1 if labelCounter.lbl_i >= len(labels) - 1: labelCounter.lbl_i = 0 labelCounter.lbl = labels[labelCounter.lbl_i] i += 1 if i >= len(labels): break return labelCounter.lbl def iob_to_label(label): label = label[2:] if not label: return 'other' return label def drawLabels(image, true_predictions, true_boxes): image_tmp = image.copy() draw = ImageDraw.Draw(image_tmp) color = cord_label_to_color(labelCounter.lbl) draw.text((10, 10), text=labelCounter.lbl, fill=color, font=font) for prediction, box in zip(true_predictions, true_boxes): predicted_label = iob_to_label(prediction).lower() color = cord_label_to_color(prediction) if not labelCounter.lbl in prediction: continue # color = label2color[predicted_label] if predicted_label in label2color else 'black' draw.rectangle(box, outline=color, width=5) return image_tmp # folder = "/Users/eliaweiss/.cache/huggingface/datasets/downloads/extracted/87634c2ab68012df3def8353986bcb092170ef7341c69e1a9cd97be52e513079/CORD/test/image/" # folder = "/Users/eliaweiss/Documents/doc2txt/en_invoice_printed" folder = "/Users/eliaweiss/ai/ICDAR-2019-SROIE/data/img" for img_name in os.listdir(folder): labelCounter.lbl_i = 0 start_time = time.time() img_path = os.path.join(folder, img_name) image = Image.open(img_path) # image = Image.open(example['image_path']) # pathOcr = workingPath + docNumber+".json" # with open(pathOcr, encoding="utf-8") as f: # gOcrJson = json.load(f) gOcr = GoogleVisionService(img_path) gOcrJson = gOcr.googleOcr() image = image.convert("RGB") width, height = image.size """We prepare it for the model using `LayoutLMv2Processor`.""" # Extract words and bounding boxes words = [] boxes = [] words, boxes = get_word_boxes_google(gOcrJson) boxes = [normalize_bbox(box, width, height) for box in boxes] # # Use pytesseract to perform OCR on the image # cv_image = cv2.imread(image_path) # gray = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY) # # Get word-level bounding boxes using pytesseract # data = pytesseract.image_to_data(gray, output_type=pytesseract.Output.DICT) # words, boxes = get_word_boxes_tesseract(data) # boxes = [normalize_bbox(box, width, height) for box in boxes] if not apply_ocr: encoding = processor(image, words, boxes=boxes, return_offsets_mapping=True, return_tensors="pt") else: encoding = processor( image, return_offsets_mapping=True, return_tensors="pt") offset_mapping = encoding.pop('offset_mapping') print(encoding.keys()) """Next, let's move everything to the GPU, if it's available.""" for k, v in encoding.items(): encoding[k] = v.to(device) # forward pass outputs = model(**encoding) # print(outputs.logits.shape) print("Time: " + str(time.time() - start_time)) """Let's create the true predictions as well as the true boxes. With "true", I mean only taking into account tokens that are at the start of a given word. We can use the `offset_mapping` returned by the processor to determine which tokens are a subword.""" predictions = outputs.logits.argmax(-1).squeeze().tolist() token_boxes = encoding.bbox.squeeze().tolist() import numpy as np is_subword = np.array(offset_mapping.squeeze().tolist())[:, 0] != 0 true_predictions = [id2label[pred] for idx, pred in enumerate(predictions) if not is_subword[idx]] true_boxes = [unnormalize_box(box, width, height) for idx, box in enumerate( token_boxes) if not is_subword[idx]] # print(true_predictions) # print(true_boxes) """Let's visualize the result!""" labelCounter.lbl = "I-total.total_price" # getNextLabel(labels, true_predictions) image_tmp = drawLabels(image, true_predictions, true_boxes) # display the image with cv2 # Display the image image_np = np.array(image_tmp) # cv2.imwrite('output_image_v2.jpg', image_np) cv2.imshow('Window Name ', image_np) while (1): k = cv2.waitKey(0) & 0xFF if k == 255: continue if k == 126: # shift + ` break print("k", k) if k == 9: # tab - change direction labelCounter.lbl = getNextLabel(labels, true_predictions) image_tmp = drawLabels(image, true_predictions, true_boxes) # display the image with cv2 # Display the image image_np = np.array(image_tmp) # cv2.imwrite('output_image_v2.jpg', image_np) cv2.imshow('Window Name ', image_np)