Reverted

app.py

@@ -1,9 +1,10 @@
 import gradio as gr
 from huggingface_hub import hf_hub_download
-from PIL import Image, ImageDraw
+from PIL import Image
 import torch
 from transformers import AutoImageProcessor, AutoModelForObjectDetection
 
+gr.load("models/microsoft/table-transformer-structure-recognition").launch()
 # Load the processor and model for table structure recognition
 processor = AutoImageProcessor.from_pretrained("microsoft/table-transformer-structure-recognition")
 model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-structure-recognition")
@@ -20,34 +21,16 @@ def predict(image):
     # Extract bounding boxes and class labels
     predicted_boxes = outputs.pred_boxes[0].cpu().numpy()  # First image
     predicted_classes = outputs.logits.argmax(-1).cpu().numpy()  # Class predictions
-    
-    # Create a drawing context for the image
-    draw = ImageDraw.Draw(image)
-    width, height = image.size
-    
-    # Loop over all detected boxes and draw them on the image
-    for box in predicted_boxes:
-        # Box coordinates are normalized, so multiply by image dimensions
-        x0, y0, x1, y1 = box
-        
-        # Ensure that y0 < y1 and x0 < x1
-        if x1 < x0:
-            x0, x1 = x1, x0
-        if y1 < y0:
-            y0, y1 = y1, y0
-        
-        # Draw the rectangle
-        draw.rectangle([x0 * width, y0 * height, x1 * width, y1 * height], outline="red", width=3)
-    
-    # Return the image with bounding boxes drawn
-    return image
+
+    # Return the bounding boxes for display
+    return {"boxes": predicted_boxes.tolist(), "classes": predicted_classes.tolist()}
 
 # Set up the Gradio interface
 interface = gr.Interface(
     fn=predict,  # The function that gets called when an image is uploaded
     inputs=gr.Image(type="pil"),  # Image input (as PIL image)
-    outputs=gr.Image(type="pil"),  # Outputting the image with boxes drawn
+    outputs="json",  # Outputting a JSON with the boxes and classes
 )
 
 # Launch the Gradio app
-interface.launch()
+interface.launch()