Create app.py

app.py

@@ -0,0 +1,137 @@
+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, 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](https://huggingface.co/datasets/SpursgoZmy/MMTab))")
+
+        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](https://huggingface.co/datasets/apoidea/pubtabnet-html))")
+
+        input_image.change(fn=table_detection, inputs=input_image, outputs=output_gallery)
+
+demo.launch()