Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import numpy as np | |
| import fitz # PyMuPDF | |
| from ultralytics import YOLOv10 | |
| import spaces | |
| # Load the trained model | |
| model = YOLOv10("best.pt") | |
| # Define the class indices for figures and tables | |
| figure_class_index = 4 # class index for figures | |
| table_class_index = 3 # 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]), int(box.cls[0])) | |
| 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): | |
| figures = [] | |
| tables = [] | |
| for (x1, y1, x2, y2, cls) in boxes: | |
| cropped_img = image[int(y1 * scale_factor):int(y2 * scale_factor), int(x1 * scale_factor):int(x2 * scale_factor)] | |
| if cls == figure_class_index: | |
| figures.append(cropped_img) | |
| elif cls == table_class_index: | |
| tables.append(cropped_img) | |
| return figures, tables | |
| def process_pdf(pdf_file): | |
| # Open the PDF file | |
| doc = fitz.open(pdf_file) | |
| all_figures = [] | |
| all_tables = [] | |
| # 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 | |
| figures, tables = crop_images_from_boxes(high_res_img, boxes, scale_factor) | |
| all_figures.extend(figures) | |
| all_tables.extend(tables) | |
| return all_figures, all_tables | |
| # Create Gradio interface | |
| with gr.Blocks() as app: | |
| gr.Markdown( | |
| """ | |
| # PDF Figures and Tables Extraction | |
| Upload a PDF file to extract figures and tables using YOLOv10. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| file_input = gr.File(label="Upload a PDF") | |
| with gr.Column(): | |
| extract_button = gr.Button("Extract") | |
| with gr.Row(): | |
| with gr.Column(): | |
| figures_gallery = gr.Gallery(label="Figures from PDF", object_fit='scale-down') | |
| with gr.Column(): | |
| tables_gallery = gr.Gallery(label="Tables from PDF", object_fit='scale-down') | |
| extract_button.click(process_pdf, inputs=file_input, outputs=[figures_gallery, tables_gallery]) | |
| app.launch() | |