Upload LayoutLMv2Main_cord2_gOcr_folder.py

LayoutLMv2Main_cord2_gOcr_folder.py

@@ -0,0 +1,234 @@
+# -*- 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)