| import gradio as gr | |
| from transformers import AutoImageProcessor, AutoModelForObjectDetection | |
| import torch | |
| # 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") | |
| # Define the inference function | |
| def predict(image): | |
| # Preprocess the input image | |
| inputs = processor(images=image, return_tensors="pt") | |
| # Perform object detection using the model | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| # Extract bounding boxes and class labels | |
| predicted_boxes = outputs.pred_boxes[0].cpu().numpy() # First 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 = [] | |
| 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( | |
| 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 | |
| ) | |
| # Launch the Gradio app | |
| interface.launch() | |