File size: 5,640 Bytes
b7d2529 a971389 b7d2529 a971389 b7d2529 a971389 b7d2529 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
import gradio as gr
import numpy as np
from PIL import Image
import torch
import pandas as pd
from transformers import AutoImageProcessor, AutoModelForObjectDetection, AutoProcessor, Pix2StructForConditionalGeneration
import torch
from io import StringIO
device="cpu"
MAX_PATCHES = 1024
MAX_NEW_TOKENS = 1024
TABLE_THRESHOLD = 0.9
TABLE_PADDING = 5
# Detection related
table_detr_processor = AutoImageProcessor.from_pretrained("microsoft/table-transformer-detection")
table_detr_model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-detection", revision="no_timm")
table_detr_model.to(device)
table_detr_model.eval()
no_table_found = Image.open("app_assets/no_table_found.png")
# Recognition related
table_recog_processor = AutoProcessor.from_pretrained("KennethTM/pix2struct-base-table2html")
table_recog_model = Pix2StructForConditionalGeneration.from_pretrained("KennethTM/pix2struct-base-table2html")
table_recog_model.to(device)
table_recog_model.eval()
def table_detection(image, threshold=TABLE_THRESHOLD):
inputs = table_detr_processor(images=image, return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}
with torch.inference_mode():
outputs = table_detr_model(**inputs)
target_sizes = torch.tensor([image.size[::-1]])
results = table_detr_processor.post_process_object_detection(outputs, threshold=threshold, target_sizes=target_sizes)
table_boxes = [i for i in results[0]["boxes"]]
tables = []
if len(table_boxes) == 0:
tables.append(no_table_found)
else:
padding = TABLE_PADDING
for box in table_boxes:
box = [int(i) for i in box]
box[0] = max(0, box[0]-padding)
box[1] = max(0, box[1]-padding)
box[2] = min(image.width, box[2]+padding)
box[3] = min(image.height, box[3]+padding)
tables.append(image.crop(box))
return tables
def table_recognition(image, max_new_tokens = MAX_NEW_TOKENS):
encoding = table_recog_processor(image, return_tensors="pt", max_patches=MAX_PATCHES)
with torch.inference_mode():
flattened_patches = encoding.pop("flattened_patches").to(device)
attention_mask = encoding.pop("attention_mask").to(device)
predictions = table_recog_model.generate(flattened_patches=flattened_patches, attention_mask=attention_mask, max_new_tokens=max_new_tokens)
predictions_decoded = table_recog_processor.tokenizer.batch_decode(predictions, skip_special_tokens=True)
table_html = predictions_decoded[0]
return table_html
def table_recognition_outputs(image):
# Table to HTML
table_html = table_recognition(image)
# Write HTML to files
with open("table.html", "w") as file:
file.write(table_html)
df = pd.read_html(StringIO(table_html))[0]
df.to_csv("table.csv", index=False)
return [table_html,
gr.DownloadButton("Download HTML", value="table.html", visible=True),
gr.DownloadButton("Download CSV", value="table.csv", visible=True)]
demo_detection = [
"app_assets/example_one_table.jpg",
"app_assets/example_two_tables.jpg",
]
demo_recognition = [
"app_assets/example_recog_1.jpg",
"app_assets/example_recog_2.jpg",
]
with gr.Blocks() as demo:
with gr.Tab("Recognition"):
gr.Markdown("# Table recognition")
gr.Markdown("This model ([KennethTM/pix2struct-base-table2html](https://huggingface.co/KennethTM/pix2struct-base-table2html)) converts an image of a table to HTML format and is finetuned from [Pix2Struct base model](https://huggingface.co/google/pix2struct-base).")
gr.Markdown("The model expects an image containing only a table. If the table is embedded in a document, first use the detection model in the 'Detection' tab.")
gr.Markdown("*Note that recognition model inference is slow on CPU (a few minutes), please be patient*")
with gr.Row():
with gr.Column():
input_table = gr.Image(type="pil", label="Table", show_label=True, scale=1)
with gr.Column():
output_html = gr.HTML(label="Table (HTML format)", show_label=False)
with gr.Row():
download_html = gr.DownloadButton(visible=False)
download_csv = gr.DownloadButton(visible=False)
with gr.Row():
examples = gr.Examples(demo_recognition, input_table, cache_examples=False, label="Example tables (MMTab dataset)")
input_table.change(fn=table_recognition_outputs, inputs=input_table, outputs=[output_html, download_html, download_csv])
with gr.Tab("Detection"):
gr.Markdown("# Table detection")
gr.Markdown("This model detect tables in a document image with [Microsoft's Table Transformer model](https://huggingface.co/microsoft/table-transformer-detection).")
gr.Markdown("Use the detection to find tables, download the results and use as input for table recognition in the 'Recognition' tab.")
with gr.Row():
with gr.Column():
input_image = gr.Image(type="pil", label="Document", show_label=True, scale=1)
with gr.Column():
output_gallery = gr.Gallery(type="pil", label="Tables", show_label=True, scale=1, format="png")
with gr.Row():
examples = gr.Examples(demo_detection, input_image, cache_examples=False, label="Example documents (PubTabNet dataset)")
input_image.change(fn=table_detection, inputs=input_image, outputs=output_gallery)
demo.launch()
|