import gradio as gr
import numpy as np
import fitz  # PyMuPDF
import spaces
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(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[int(y1 * scale_factor):int(y2 * scale_factor), int(x1 * scale_factor):int(x2 * scale_factor)]
        for (x1, y1, x2, y2) in boxes
    ]
    return cropped_images

@spaces.GPU
def process_pdf(pdf_file):
    # Open the PDF file
    doc = fitz.open(pdf_file)
    all_cropped_images = []

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

    # Calculate the scaling factor
    scale_factor = high_dpi / low_dpi

    # Pre-cache all page pixmaps at low DPI
    low_res_pixmaps = [page.get_pixmap(dpi=low_dpi) for page in doc]
    
    # Loop through each page
    for page_num, low_res_pix in enumerate(low_res_pixmaps):
        low_res_img = np.frombuffer(low_res_pix.samples, dtype=np.uint8).reshape(low_res_pix.height, low_res_pix.width, 3)
        
        # Get bounding boxes from low DPI image
        boxes = infer_image_and_get_boxes(low_res_img)
        
        if boxes:
            # Load high DPI image for cropping only if boxes are found
            high_res_pix = doc[page_num].get_pixmap(dpi=high_dpi)
            high_res_img = np.frombuffer(high_res_pix.samples, dtype=np.uint8).reshape(high_res_pix.height, high_res_pix.width, 3)
            
            # 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()