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()