import gradio as gr
from huggingface_hub import hf_hub_download
from PIL import Image
import torch
import pytesseract
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")

# Check if GPU is available and use it; otherwise, use CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Define the inference and OCR function
def predict(image):
    # Preprocess the input image
    inputs = processor(images=image, return_tensors="pt").to(device)
    
    # Perform object detection using the model
    with torch.no_grad():
        outputs = model(**inputs)

    # Extract bounding boxes and filter for columns
    predicted_boxes = outputs.pred_boxes[0].cpu().numpy()  # First image
    predicted_classes = outputs.logits.argmax(-1).cpu().numpy()  # Class predictions

    # Prepare OCR results
    ocr_results = []

    image_width, image_height = image.size  # Get original image dimensions

    # Iterate over detected boxes and perform OCR on columns
    for box in predicted_boxes:
        # Unpack the normalized bounding box (x_min, y_min, x_max, y_max)
        x_min, y_min, x_max, y_max = box

        # Calculate width and height (denormalize)
        width = x_max - x_min
        height = y_max - y_min

        # Filter for columns based on aspect ratio (height > width)
        if height / width > 2:  # A threshold for vertical aspect ratio (adjust if needed)
            # Convert normalized coordinates to pixel values
            left = int(x_min * image_width)
            top = int(y_min * image_height)
            right = int(x_max * image_width)
            bottom = int(y_max * image_height)

            # Crop the image to the bounding box area
            cropped_image = image.crop((left, top, right, bottom))

            # Perform OCR on the cropped image
            ocr_text = pytesseract.image_to_string(cropped_image)

            # Append OCR result for this box
            ocr_results.append({
                "box": [left, top, right, bottom],
                "text": ocr_text
            })

    # Return OCR results
    return {"ocr_results": ocr_results}

# 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 OCR results
)

# Launch the Gradio app
interface.launch()