Spaces:
Sleeping
Sleeping
| import json | |
| import gradio as gr | |
| from pdfminer.high_level import extract_pages | |
| from pdfminer.layout import LTTextBoxHorizontal, LTFigure, LTImage | |
| import os | |
| import io | |
| from PIL import Image | |
| import pandas as pd | |
| import tabula | |
| import camelot | |
| from PyPDF2 import PdfReader | |
| def parse_pdf(pdf_file, output_format, progress=gr.Progress()): | |
| """ | |
| Parses a PDF file, extracts text, tables, and images, and formats the output. | |
| Args: | |
| pdf_file: Path to the uploaded PDF file. | |
| output_format: Desired output format ("JSON", "Markdown", or "HTML"). | |
| progress: Gradio Progress object for displaying progress. | |
| Returns: | |
| tuple: Extracted text and download data in the specified format. | |
| Returns an empty string and None if there is an error. | |
| """ | |
| try: | |
| with open(pdf_file, 'rb') as file: | |
| pages = list(extract_pages(file)) # Convert generator to list | |
| text = "" | |
| tables = [] | |
| images = [] | |
| # Iterate through pages and extract text and images | |
| for i, page in enumerate(pages): | |
| progress(i / len(pages)) # Update progress bar | |
| for element in page: | |
| if isinstance(element, LTTextBoxHorizontal): | |
| text += element.get_text() | |
| elif isinstance(element, (LTFigure, LTImage)): | |
| try: | |
| if hasattr(element, 'stream'): | |
| image_data = element.stream.read() | |
| image = Image.open(io.BytesIO(image_data)) | |
| image_filename = f"extracted_image_{len(images)}.png" | |
| image.save(image_filename) | |
| images.append({"filename": image_filename}) | |
| else: | |
| for child in element: | |
| if isinstance(child, LTImage): | |
| image_data = child.stream.read() | |
| image = Image.open(io.BytesIO(image_data)) | |
| image_filename = f"extracted_image_{len(images)}.png" | |
| image.save(image_filename) | |
| images.append({"filename": image_filename}) | |
| except Exception as e: | |
| print(f"Error extracting image: {e}") | |
| # Enhanced table extraction (tabula-py preferred, fallback to camelot) | |
| try: | |
| tables = tabula.read_pdf(pdf_file, pages='all', multiple_tables=True) | |
| except Exception as e: | |
| print(f"tabula-py failed: {e}. Trying camelot...") | |
| try: | |
| camelot_tables = camelot.read_pdf(pdf_file) | |
| for table in camelot_tables: | |
| tables.append(table.df) | |
| except Exception as e: | |
| print(f"camelot also failed: {e}. No tables extracted.") | |
| # Format extracted data based on user selection | |
| if output_format == "JSON": | |
| json_data = { | |
| "text": text, | |
| "tables": [table.to_dict() for table in tables], | |
| "images": images | |
| } | |
| download_data = json.dumps(json_data, indent=4) # Add indentation for readability | |
| elif output_format == "Markdown": | |
| markdown_text = f"# Extracted Text\n\n{text}\n\n# Tables\n" | |
| for i, table in enumerate(tables): | |
| markdown_text += f"## Table {i+1}\n" | |
| markdown_text += table.to_markdown(index=False) + "\n\n" | |
| # Image embedding in Markdown (using relative paths) | |
| markdown_text += "\n\n# Images\n\n" | |
| for image in images: | |
| image_path = os.path.join(os.getcwd(), image["filename"]) | |
| markdown_text += f'\n' | |
| download_data = markdown_text | |
| elif output_format == "HTML": | |
| html_text = f"<p>{text}</p>\n\n<h2>Tables</h2>\n" | |
| for i, table in enumerate(tables): | |
| html_text += f"<h2>Table {i+1}</h2>\n" | |
| html_text += table.to_html() + "<br>" | |
| # Image embedding in HTML (using relative paths) | |
| html_text += "\n\n<h2>Images</h2>\n\n" | |
| for image in images: | |
| image_path = os.path.join(os.getcwd(), image["filename"]) | |
| html_text += f'<img src="{image_path}" alt="Image"><br>\n' | |
| download_data = html_text.encode("utf-8") # Encode for HTML download | |
| return text, download_data | |
| except Exception as main_e: | |
| print(f"A main error occurred: {main_e}") | |
| return "", None # Return empty string and None in case of error | |
| iface = gr.Interface( | |
| fn=parse_pdf, | |
| inputs=["file", gr.Dropdown(["JSON", "Markdown", "HTML"])], # Remove gr.Progress() from inputs | |
| outputs=[ | |
| gr.Text(label="Output Text"), | |
| gr.File(label="Download Output") | |
| ], | |
| title="PDF Parser", | |
| description="Parse a PDF and choose the output format." | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch(share=False) |