import gradio as gr
from PIL import Image, ImageDraw, ImageFont
import random
import pandas as pd
import numpy as np
from datasets import concatenate_datasets
from operator import itemgetter
import collections
# download datasets
from datasets import load_dataset
dataset_small = load_dataset("pierreguillou/DocLayNet-small")
dataset_base = load_dataset("pierreguillou/DocLayNet-base")
id2label = {idx:label for idx,label in enumerate(dataset_small["train"].features["categories"].feature.names)}
label2id = {label:idx for idx,label in id2label.items()}
labels = [label for idx, label in id2label.items()]
# need to change the coordinates format
def convert_box(box):
x, y, w, h = tuple(box) # the row comes in (left, top, width, height) format
actual_box = [x, y, x+w, y+h] # we turn it into (left, top, left+widght, top+height) to get the actual box
return actual_box
# get back original size
def original_box(box, original_width, original_height, coco_width, coco_height):
return [
int(original_width * (box[0] / coco_width)),
int(original_height * (box[1] / coco_height)),
int(original_width * (box[2] / coco_width)),
int(original_height* (box[3] / coco_height)),
]
# function to sort bounding boxes
def get_sorted_boxes(bboxes):
# sort by y from page top to bottom
bboxes = sorted(bboxes, key=itemgetter(1), reverse=False)
y_list = [bbox[1] for bbox in bboxes]
# sort by x from page left to right when boxes with same y
if len(list(set(y_list))) != len(y_list):
y_list_duplicates_indexes = dict()
y_list_duplicates = [item for item, count in collections.Counter(y_list).items() if count > 1]
for item in y_list_duplicates:
y_list_duplicates_indexes[item] = [i for i, e in enumerate(y_list) if e == item]
bbox_list_y_duplicates = sorted(np.array(bboxes)[y_list_duplicates_indexes[item]].tolist(), key=itemgetter(0), reverse=False)
np_array_bboxes = np.array(bboxes)
np_array_bboxes[y_list_duplicates_indexes[item]] = np.array(bbox_list_y_duplicates)
bboxes = np_array_bboxes.tolist()
return bboxes
# categories colors
label2color = {
'Caption': 'brown',
'Footnote': 'orange',
'Formula': 'gray',
'List-item': 'yellow',
'Page-footer': 'red',
'Page-header': 'red',
'Picture': 'violet',
'Section-header': 'orange',
'Table': 'green',
'Text': 'blue',
'Title': 'pink'
}
# image witout content
examples_dir = 'samples/'
images_wo_content = examples_dir + "wo_content.png"
df_paragraphs_wo_content, df_lines_wo_content = pd.DataFrame(), pd.DataFrame()
df_paragraphs_wo_content["paragraphs"] = [0]
df_paragraphs_wo_content["categories"] = ["no content"]
df_paragraphs_wo_content["texts"] = ["no content"]
df_paragraphs_wo_content["bounding boxes"] = ["no content"]
df_lines_wo_content["lines"] = [0]
df_lines_wo_content["categories"] = ["no content"]
df_lines_wo_content["texts"] = ["no content"]
df_lines_wo_content["bounding boxes"] = ["no content"]
# lists
font = ImageFont.load_default()
dataset_names = ["small", "base"]
splits = ["all", "train", "validation", "test"]
domains = ["all", "Financial Reports", "Manuals", "Scientific Articles", "Laws & Regulations", "Patents", "Government Tenders"]
domains_names = [domain_name.lower().replace(" ", "_").replace("&", "and") for domain_name in domains]
categories = labels + ["all"]
# function to get a rendom image and all data from DocLayNet
def generate_annotated_image(dataset_name, split, domain, category):
# error message
msg_error = ""
# get dataset
if dataset_name == "small": example = dataset_small
else: example = dataset_base
# get split
if split == "all":
example = concatenate_datasets([example["train"], example["validation"], example["test"]])
else:
example = example[split]
# get domain
domain_name = domains_names[domains.index(domain)]
if domain_name != "all":
example = example.filter(lambda example: example["doc_category"] == domain_name)
if len(example) == 0:
msg_error = f'There is no image with at least one labeled bounding box that matches your settings (dataset: "DocLayNet {dataset_name}" / domain: "{domain}" / split: "{split}").'
example = dict()
# get category
idx_list = list()
if category != "all":
for idx, categories_list in enumerate(example["categories"]):
if int(label2id[category]) in categories_list:
idx_list.append(idx)
if len(idx_list) > 0:
example = example.select(idx_list)
else:
msg_error = f'There is no image with at least one labeled bounding box that matches your settings (dataset: "DocLayNet {dataset_name}" / split: "{split}" / domain: "{domain}" / category: "{category}").'
example = dict()
if len(msg_error) > 0:
# save image files
Image.open(images_wo_content).save("wo_content.png")
# save csv files
df_paragraphs_wo_content.to_csv("paragraphs_wo_content.csv", encoding="utf-8", index=False)
df_lines_wo_content.to_csv("lines_wo_content.csv", encoding="utf-8", index=False)
return msg_error, "wo_content.png", images_wo_content, images_wo_content, "wo_content.png", "wo_content.png", df_paragraphs_wo_content, df_lines_wo_content, gr.File.update(value="paragraphs_wo_content.csv", visible=False), gr.File.update(value="lines_wo_content.csv", visible=False)
else:
# get random image & PDF data
index = random.randint(0, len(example))
image = example[index]["image"] # original image
coco_width, coco_height = example[index]["coco_width"], example[index]["coco_height"]
original_width, original_height = example[index]["original_width"], example[index]["original_height"]
original_filename = example[index]["original_filename"]
page_no = example[index]["page_no"]
num_pages = example[index]["num_pages"]
# resize image to original
image = image.resize((original_width, original_height))
# get image of PDF without bounding boxes
img_file = original_filename.replace(".pdf", ".png")
image.save(img_file)
# get corresponding annotations
texts = example[index]["texts"]
bboxes_block = example[index]["bboxes_block"]
bboxes_line = example[index]["bboxes_line"]
categories = example[index]["categories"]
domain = example[index]["doc_category"]
# convert boxes to original
original_bboxes_block = [original_box(convert_box(box), original_width, original_height, coco_width, coco_height) for box in bboxes_block]
original_bboxes_line = [original_box(convert_box(box), original_width, original_height, coco_width, coco_height) for box in bboxes_line]
original_bboxes = [original_bboxes_block, original_bboxes_line]
##### block boxes #####
# get list of unique block boxes
original_blocks = dict()
original_bboxes_block_list = list()
original_bbox_block_prec = list()
for count_block, original_bbox_block in enumerate(original_bboxes_block):
if original_bbox_block != original_bbox_block_prec:
original_bbox_block_indexes = [i for i, original_bbox in enumerate(original_bboxes_block) if original_bbox == original_bbox_block]
original_blocks[count_block] = original_bbox_block_indexes
original_bboxes_block_list.append(original_bbox_block)
original_bbox_block_prec = original_bbox_block
# get list of categories and texts by unique block boxes
category_block_list, text_block_list = list(), list()
for original_bbox_block in original_bboxes_block_list:
count_block = original_bboxes_block.index(original_bbox_block)
original_bbox_block_indexes = original_blocks[count_block]
category_block = categories[original_bbox_block_indexes[0]]
category_block_list.append(category_block)
if id2label[category_block] == "Text" or id2label[category_block] == "Caption" or id2label[category_block] == "Footnote":
text_block = ' '.join(np.array(texts)[original_bbox_block_indexes].tolist())
elif id2label[category_block] == "Section-header" or id2label[category_block] == "Title" or id2label[category_block] == "Picture" or id2label[category_block] == "Formula" or id2label[category_block] == "List-item" or id2label[category_block] == "Table" or id2label[category_block] == "Page-header" or id2label[category_block] == "Page-footer":
text_block = '\n'.join(np.array(texts)[original_bbox_block_indexes].tolist())
text_block_list.append(text_block)
# sort data from y = 0 to end of page (and after, x=0 to end of page when necessary)
sorted_original_bboxes_block_list = get_sorted_boxes(original_bboxes_block_list)
sorted_original_bboxes_block_list_indexes = [original_bboxes_block_list.index(item) for item in sorted_original_bboxes_block_list]
sorted_category_block_list = np.array(category_block_list)[sorted_original_bboxes_block_list_indexes].tolist()
sorted_text_block_list = np.array(text_block_list)[sorted_original_bboxes_block_list_indexes].tolist()
##### line boxes ####
# sort data from y = 0 to end of page (and after, x=0 to end of page when necessary)
original_bboxes_line_list = original_bboxes_line
category_line_list = categories
text_line_list = texts
sorted_original_bboxes_line_list = get_sorted_boxes(original_bboxes_line_list)
sorted_original_bboxes_line_list_indexes = [original_bboxes_line_list.index(item) for item in sorted_original_bboxes_line_list]
sorted_category_line_list = np.array(category_line_list)[sorted_original_bboxes_line_list_indexes].tolist()
sorted_text_line_list = np.array(text_line_list)[sorted_original_bboxes_line_list_indexes].tolist()
# setup images & PDF data
columns = 2
images = [image.copy(), image.copy()]
num_imgs = len(images)
imgs, df_paragraphs, df_lines = dict(), pd.DataFrame(), pd.DataFrame()
for i, img in enumerate(images):
draw = ImageDraw.Draw(img)
for box, label_idx, text in zip(original_bboxes[i], categories, texts):
label = id2label[label_idx]
color = label2color[label]
draw.rectangle(box, outline=color)
text = text.encode('latin-1', 'replace').decode('latin-1') # https://stackoverflow.com/questions/56761449/unicodeencodeerror-latin-1-codec-cant-encode-character-u2013-writing-to
draw.text((box[0] + 10, box[1] - 10), text=label, fill=color, font=font)
if i == 0:
imgs["paragraphs"] = img
# save
img_paragraphs = "img_paragraphs_" + original_filename.replace(".pdf", ".png")
img.save(img_paragraphs)
df_paragraphs["paragraphs"] = list(range(len(sorted_original_bboxes_block_list)))
df_paragraphs["categories"] = [id2label[label_idx] for label_idx in sorted_category_block_list]
df_paragraphs["texts"] = sorted_text_block_list
df_paragraphs["bounding boxes"] = [str(bbox) for bbox in sorted_original_bboxes_block_list]
# save
csv_paragraphs = "csv_paragraphs_" + original_filename.replace(".pdf", ".csv")
df_paragraphs.to_csv(csv_paragraphs, encoding="utf-8", index=False)
else:
imgs["lines"] = img
# save
img_lines = "img_lines_" + original_filename.replace(".pdf", ".png")
img.save(img_lines)
df_lines["lines"] = list(range(len(sorted_original_bboxes_line_list)))
df_lines["categories"] = [id2label[label_idx] for label_idx in sorted_category_line_list]
df_lines["texts"] = sorted_text_line_list
df_lines["bounding boxes"] = [str(bbox) for bbox in sorted_original_bboxes_line_list]
# save
csv_lines = "csv_lines_" + original_filename.replace(".pdf", ".csv")
df_lines.to_csv(csv_lines, encoding="utf-8", index=False)
msg = f'The page {page_no} of the PDF "{original_filename}" (domain: "{domain}") matches your settings.'
return msg, img_file, imgs["paragraphs"], imgs["lines"], img_paragraphs, img_lines, df_paragraphs, df_lines, gr.File.update(value=csv_paragraphs, visible=True), gr.File.update(value=csv_lines, visible=True)
# gradio APP
with gr.Blocks(title="DocLayNet image viewer", css=".gradio-container") as demo:
gr.HTML("""
DocLayNet image viewer
(01/29/2023) This APP is an image viewer of the DocLayNet dataset and a data extraction tool.
Make your settings and the output will show 2 images of a randomly selected PDF with labeled bounding boxes, one of paragraphs and the other of lines, and their corresponding tables of texts with their labels.
For example, if you select the domain "laws_and_regulations" and the category "Caption", you will get a random PDF that corresponds to these settings (ie, it will have at least one bounding box labeled with "Caption" in the PDF).
WARNING: if the app crashes or runs without providing a result, refresh the page (DocLayNet image viewer) and run a search again. If the same problem occurs again, prefer the DocLayNet small. Thanks.