Update app.py

app.py

@@ -2,7 +2,7 @@ import matplotlib.pyplot as plt
 import matplotlib.patches as patches
 from matplotlib.patches import Patch
 import io
-from PIL import Image
+from PIL import Image, ImageDraw
 
 from transformers import TableTransformerImageProcessor, AutoModelForObjectDetection
 import torch
@@ -13,6 +13,10 @@ import gradio as gr
 processor = TableTransformerImageProcessor(max_size=800)
 model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-detection", revision="no_timm")
 
+# load table structure recognition model
+structure_processor = TableTransformerImageProcessor(max_size=1000)
+structure_model = AutoModelForObjectDetection.from_pretrained("microsoft/table-structure-recognition-v1.1-all")
+
 
 # for output bounding box post-processing
 def box_cxcywh_to_xyxy(x):
@@ -103,7 +107,7 @@ def visualize_detected_tables(img, det_tables):
     return fig
 
 
-def detect_table(image):
+def detect_and_crop_table(image):
     # prepare image for the model
     pixel_values = processor(image, return_tensors="pt").pixel_values
 
@@ -117,8 +121,41 @@ def detect_table(image):
     detected_tables = outputs_to_objects(outputs, image.size, id2label)
 
     # visualize
-    fig = visualize_detected_tables(image, detected_tables)
-    image = fig2img(fig)
+    # fig = visualize_detected_tables(image, detected_tables)
+    # image = fig2img(fig)
+
+    # crop first detected table out of image
+    cropped_table = image.crop(objects[0]["bbox"])
+
+    return cropped_table
+
+
+def recognize_table(image):
+    # prepare image for the model
+    pixel_values = structure_processor(images=cropped_table, return_tensors="pt").pixel_values
+
+    # forward pass
+    with torch.no_grad():
+        outputs = structure_model(pixel_values)
+
+    # postprocess to get individual elements
+    id2label = structure_modelmodel.config.id2label
+    id2label[len(structure_modelmodel.config.id2label)] = "no object"
+    detected_tables = outputs_to_objects(outputs, image.size, id2label)
+
+    # visualize cells on cropped table
+    draw = ImageDraw.Draw(image)
+
+    for cell in cells:
+        draw.rectangle(cell["bbox"], outline="red")
+        
+    return image
+
+
+def process_pdf(image):
+    cropped_table = detect_and_crop_table(image)
+
+    image = recognize_table(cropped_table)
 
     return image
     
@@ -127,7 +164,7 @@ title = "Demo: table detection with Table Transformer"
 description = "Demo for the Table Transformer (TATR)."
 examples =[['image.png']]
 
-app = gr.Interface(fn=detect_table, 
+app = gr.Interface(fn=process_pdf, 
                      inputs=gr.Image(type="pil"), 
                      outputs=gr.Image(type="pil", label="Detected table"),
                      title=title,