change class name from column to table column

app.py

@@ -1,10 +1,7 @@
 import gradio as gr
-from huggingface_hub import hf_hub_download
-from PIL import Image
-import torch
 from transformers import AutoImageProcessor, AutoModelForObjectDetection
+import torch
 
-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,11 +17,19 @@ 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
+    predicted_class_logits = outputs.logits[0].cpu().numpy()  # Class logits for the first image
+    predicted_classes = predicted_class_logits.argmax(-1)  # Get class predictions
     class_names = model.config.id2label  # Get the class name mapping
-    print(class_names)
-    # Return the bounding boxes for display
-    return {"boxes": predicted_boxes.tolist(), "classes": predicted_classes.tolist()}
+    
+    # Filter predictions to only include columns based on class name
+    column_boxes = []
+    for idx, class_id in enumerate(predicted_classes):
+        class_name = class_names[class_id]
+        if "table column" in class_name.lower():  # Check if the class name contains 'column'
+            column_boxes.append(predicted_boxes[idx])
+
+    # Return the bounding boxes for columns
+    return {"boxes": column_boxes, "classes": ["table column"] * len(column_boxes)}
 
 # Set up the Gradio interface
 interface = gr.Interface(