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