Find whats the name,id and classes inside the model

app.py

@@ -20,22 +20,27 @@ def predict(image):
     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
-    
-    # Filter predictions to only include columns based on class name
-    column_boxes = []
+
+    # Collect the class IDs and labels along with the bounding boxes
+    result = []
     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])
+        result.append({
+            "class_id": int(class_id),
+            "class_name": class_name,
+            "bounding_box": predicted_boxes[idx].tolist()  # Convert to list for JSON serialization
+        })
 
-    # Return the bounding boxes for columns
-    return {"boxes": column_boxes, "classes": ["table column"] * len(column_boxes)}
+    # Return the bounding boxes and classes
+    return result
 
 # 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="json",  # Outputting a JSON with the boxes and classes
+    outputs="json",  # Outputting a JSON with the class labels, IDs, and bounding boxes
+    title="Table Structure Recognition",  # Add title for clarity
+    description="Upload an image and see the detected table columns and their corresponding class IDs.",
 )
 
 # Launch the Gradio app