import gradio as gr
import numpy as np
from pdf2image import convert_from_path
from PIL import Image
from ultralytics import YOLOv10

# Load the trained model
model = YOLOv10("best.pt")

# Define the class indices for figures and tables
figure_class_index = 3  # class index for figures
table_class_index = 4   # class index for tables

# Function to perform inference on an image and return bounding boxes for figures and tables
def infer_image_and_get_boxes(image, confidence_threshold=0.6):
    results = model.predict(np.array(image))
    boxes = [
        (int(box.xyxy[0][0]), int(box.xyxy[0][1]), int(box.xyxy[0][2]), int(box.xyxy[0][3]))
        for result in results for box in result.boxes
        if int(box.cls[0]) in {figure_class_index, table_class_index} and box.conf[0] > confidence_threshold
    ]
    return boxes

# Function to crop images from the boxes
def crop_images_from_boxes(image, boxes, scale_factor):
    cropped_images = [
        image.crop((int(x1 * scale_factor), int(y1 * scale_factor), int(x2 * scale_factor), int(y2 * scale_factor)))
        for (x1, y1, x2, y2) in boxes
    ]
    return cropped_images

@spaces.GPU
def process_pdf(pdf_file):
    all_cropped_images = []

    # Set the DPI for inference and high resolution for cropping
    low_dpi = 50
    high_dpi = 300

    # Convert PDF pages to images at low DPI
    low_res_images = convert_from_path(pdf_file.name, dpi=low_dpi)
    
    # Calculate the scaling factor
    scale_factor = high_dpi / low_dpi

    for page_num, low_res_img in enumerate(low_res_images):
        # Get bounding boxes from low DPI image
        boxes = infer_image_and_get_boxes(low_res_img)
        
        if boxes:
            # Convert the specific page to high DPI only if boxes are found
            high_res_img = convert_from_path(pdf_file.name, dpi=high_dpi, first_page=page_num+1, last_page=page_num+1)[0]
            
            # Crop images at high DPI
            cropped_imgs = crop_images_from_boxes(high_res_img, boxes, scale_factor)
            all_cropped_images.extend(cropped_imgs)

    return all_cropped_images

# Create Gradio interface
iface = gr.Interface(
    fn=process_pdf,
    inputs=gr.File(label="Upload a PDF"),
    outputs=gr.Gallery(label="Cropped Figures and Tables from PDF Pages"),
    title="Fast document layout analysis based on YOLOv10",
    description="Upload a PDF file to get cropped figures and tables from each page."
)

# Launch the app
iface.launch()