File size: 1,716 Bytes
cb14db8 5c56c76 25186f8 5c56c76 cb14db8 5c56c76 fa8646f 5c56c76 fa8646f 5c56c76 fa8646f 5c56c76 25186f8 a6fc7d1 5c56c76 25186f8 5c56c76 25186f8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
import gradio as gr
from huggingface_hub import hf_hub_download
from PIL import Image, ImageDraw
import torch
from transformers import AutoImageProcessor, AutoModelForObjectDetection
# 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_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
draw.rectangle([x0 * width, y0 * height, x1 * width, y1 * height], outline="red", width=3)
# Return the image with bounding boxes drawn
return image
# 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
)
# Launch the Gradio app
interface.launch()
|