..."
if input_ids_dict[str(bboxes_list[-1])][0] == (tokenizer.convert_tokens_to_ids('')):
del input_ids_dict[str(bboxes_list[-1])]
bboxes_list = bboxes_list[:-1]
# get texts by line
input_ids_list = input_ids_dict.values()
texts_list = [tokenizer.decode(input_ids) for input_ids in input_ids_list]
# display DataFrame
df = pd.DataFrame({"texts": texts_list, "input_ids": input_ids_list, "bboxes": bboxes_list})
return image, df, num_tokens, page_no, num_pages
# In[ ]:
# display chunk of PDF image and its data
def display_chunk_paragraphs_inference(index_chunk=None):
# get image and image data
image, df, num_tokens, page_no, num_pages = get_encoded_chunk_inference(index_chunk=index_chunk)
# get data from dataframe
input_ids = df["input_ids"]
texts = df["texts"]
bboxes = df["bboxes"]
print(f'Chunk ({num_tokens} tokens) of the PDF (page: {page_no+1} / {num_pages})\n')
# display image with bounding boxes
print(">> PDF image with bounding boxes of paragraphs\n")
draw = ImageDraw.Draw(image)
labels = list()
for box, text in zip(bboxes, texts):
color = "red"
draw.rectangle(box, outline=color)
# resize image to original
width, height = image.size
image = image.resize((int(0.5*width), int(0.5*height)))
# convert to cv and display
img = np.array(image, dtype='uint8') # PIL to cv2
cv2_imshow(img)
cv2.waitKey(0)
# display image dataframe
print("\n>> Dataframe of annotated paragraphs\n")
cols = ["texts", "bboxes"]
df = df[cols]
display(df)
# #### APP function
# In[ ]:
# APP outputs
def app_outputs(uploaded_pdf):
filename, msg, images = pdf_to_images(uploaded_pdf)
num_images = len(images)
if not msg.startswith("Error with the PDF"):
# Extraction of image data (text and bounding boxes)
dataset, texts_lines, texts_pars, texts_lines_par, row_indexes, par_boxes, line_boxes, lines_par_boxes = extraction_data_from_image(images)
# prepare our data in the format of the model
encoded_dataset = dataset.map(prepare_inference_features_paragraph, batched=True, batch_size=64, remove_columns=dataset.column_names)
custom_encoded_dataset = CustomDataset(encoded_dataset, tokenizer)
# Get predictions (token level)
outputs, images_ids_list, chunk_ids, input_ids, bboxes = predictions_token_level(images, custom_encoded_dataset)
# Get predictions (paragraph level)
probs_bbox, bboxes_list_dict, input_ids_dict_dict, probs_dict_dict, df = predictions_paragraph_level(dataset, outputs, images_ids_list, chunk_ids, input_ids, bboxes)
# Get labeled images with lines bounding boxes
images = get_labeled_images(dataset, images_ids_list, bboxes_list_dict, probs_dict_dict)
img_files = list()
# get image of PDF without bounding boxes
for i in range(num_images):
if filename != "files/blank.png": img_file = f"img_{i}_" + filename.replace(".pdf", ".png")
else: img_file = filename.replace(".pdf", ".png")
img_file = img_file.replace("/", "_")
images[i].save(img_file)
img_files.append(img_file)
if num_images < max_imgboxes:
img_files += [image_blank]*(max_imgboxes - num_images)
images += [Image.open(image_blank)]*(max_imgboxes - num_images)
for count in range(max_imgboxes - num_images):
df[num_images + count] = pd.DataFrame()
else:
img_files = img_files[:max_imgboxes]
images = images[:max_imgboxes]
df = dict(itertools.islice(df.items(), max_imgboxes))
# save
csv_files = list()
for i in range(max_imgboxes):
csv_file = f"csv_{i}_" + filename.replace(".pdf", ".csv")
csv_file = csv_file.replace("/", "_")
csv_files.append(gr.File.update(value=csv_file, visible=True))
df[i].to_csv(csv_file, encoding="utf-8", index=False)
else:
img_files, images, csv_files = [""]*max_imgboxes, [""]*max_imgboxes, [""]*max_imgboxes
img_files[0], img_files[1] = image_blank, image_blank
images[0], images[1] = Image.open(image_blank), Image.open(image_blank)
csv_file = "csv_wo_content.csv"
csv_files[0], csv_files[1] = gr.File.update(value=csv_file, visible=True), gr.File.update(value=csv_file, visible=True)
df, df_empty = dict(), pd.DataFrame()
df[0], df[1] = df_empty.to_csv(csv_file, encoding="utf-8", index=False), df_empty.to_csv(csv_file, encoding="utf-8", index=False)
return msg, img_files[0], img_files[1], images[0], images[1], csv_files[0], csv_files[1], df[0], df[1]
# ## Model & tokenizer
# In[ ]:
from transformers import LayoutLMv2ForTokenClassification # LayoutXLMTokenizerFast,
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# model
# tokenizer = LayoutXLMTokenizerFast.from_pretrained(model_id)
model = LayoutLMv2ForTokenClassification.from_pretrained(model_id);
model.to(device);
# feature extractor
from transformers import LayoutLMv2FeatureExtractor
feature_extractor = LayoutLMv2FeatureExtractor(apply_ocr=False)
# tokenizer
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
# In[ ]:
# get labels
id2label = model.config.id2label
label2id = model.config.label2id
num_labels = len(id2label)
# ## Gradio APP
# In[ ]:
with gr.Blocks(title="Inference APP for Document Understanding at paragraph level (v2 - LayoutXLM base)", css=".gradio-container") as demo:
gr.HTML("""
Inference APP for Document Understanding at paragraph level (v2 - LayoutXLM base)
LayoutXLM: Multimodal Pre-training for Multilingual Visually-rich Document Understanding is a Document Understanding model that uses both layout and text in order to detect labels of bounding boxes. Combined with the model XML-RoBERTa base, this finetuned model has the capacity to understand any language. Finetuned on the dataset DocLayNet base, it can classifly any bounding box (and its OCR text) to 11 labels (Caption, Footnote, Formula, List-item, Page-footer, Page-header, Picture, Section-header, Table, Text, Title).
It relies on an external OCR engine to get words and bounding boxes from the document image. Thus, let's run in this APP an OCR engine (PyTesseract) to get the bounding boxes, then run Layout XLM base (already fine-tuned on the dataset DocLayNet base at paragraph level) on the individual tokens and then, visualize the result at paragraph level!
It allows to get all pages of any PDF (of any language) with bounding boxes labeled at paragraph level and the associated dataframes with labeled data (bounding boxes, texts, labels) :-)
More information about the DocLayNet datasets, the finetuning of the model and this APP in the following blog posts: