created app.py

app.py

@@ -0,0 +1,298 @@
+import os
+import json
+import fitz
+import streamlit as st
+from PIL import Image
+import io
+import pandas as pd
+import pickle
+import zipfile
+import tempfile
+
+def extract_text_images(
+        pdf_path: str, output_folder: str,
+        minimum_font_size: int,
+        mode: str = 'headerwise',
+        header_font_sizes: list[float] = None,
+        tolerance: float = 0.01,
+        extraction_type: str = 'both',
+        headers_are_capital: bool = False
+        ) -> str:
+    if not os.path.exists(output_folder):
+        os.makedirs(output_folder)
+
+    extraction_data = []
+    current_header = None
+    current_header_content = []
+
+    def add_current_header_content() -> None:
+        nonlocal current_header, current_header_content
+        if current_header:
+            extraction_data.append({
+                'header': current_header,
+                'content': current_header_content
+            })
+            current_header_content = []
+        current_header = None
+
+    def is_header_font_size(font_size: float) -> bool:
+        return any(
+            abs(font_size - header_font_size) <= tolerance
+            for header_font_size in header_font_sizes
+        )
+
+    def is_bold(font: str) -> bool:
+        return 'bold' in font.lower()
+
+    pdf_document = fitz.open(pdf_path)
+
+    for page_number in range(pdf_document.page_count):
+        page = pdf_document.load_page(page_number)
+        elements = []
+
+        if extraction_type in ('text', 'both'):
+            text_blocks = page.get_text("dict")["blocks"]
+            lines = {}
+
+            for block in text_blocks:
+                if block["type"] == 0:  # Text block
+                    for line in block["lines"]:
+                        for span in line["spans"]:
+                            font_size = span["size"]
+                            top = span["bbox"][1]
+                            font = span["font"]
+
+                            if font_size < minimum_font_size:
+                                continue
+
+                            if top not in lines:
+                                lines[top] = []
+                            lines[top].append(span)
+
+            for top in sorted(lines.keys()):
+                line = lines[top]
+                line_text = " ".join([span['text'] for span in line])
+                line_font_size = line[0]['size']
+                font = line[0]['font']
+
+                if headers_are_capital:
+                    line_text_is_header = line_text.isupper()
+                else:
+                    line_text_is_header = True
+
+                elements.append({
+                    'type': 'text',
+                    'font_size': line_font_size,
+                    'page': page_number + 1,
+                    'content': line_text,
+                    'x0': line[0]['bbox'][0],
+                    'top': top,
+                    'font': font,
+                    'is_header': line_text_is_header
+                })
+
+        if extraction_type in ('images', 'both'):
+            image_list = page.get_images(full=True)
+
+            for img_index, img in enumerate(image_list):
+                xref = img[0]
+                base_image = pdf_document.extract_image(xref)
+                image_bytes = base_image["image"]
+                image_filename = os.path.join(
+                    output_folder,
+                    f"page_{page_number + 1}_img_{img_index + 1}.png"
+                )
+
+                with open(image_filename, "wb") as img_file:
+                    img_file.write(image_bytes)
+
+                img_rect = page.get_image_bbox(img)
+                elements.append({
+                    'type': 'image',
+                    'page': page_number + 1,
+                    'path': image_filename,
+                    'x0': img_rect.x0,
+                    'top': img_rect.y0
+                })
+
+        elements.sort(key=lambda e: (e['top'], e['x0']))
+
+        if mode == 'headerwise':
+            for element in elements:
+                if element['type'] == 'text' and element['is_header'] and is_header_font_size(element['font_size']) and is_bold(element['font']):
+                    add_current_header_content()
+                    current_header = element['content']
+                elif element['type'] == 'text':
+                    if current_header_content and current_header_content[-1]['type'] == 'text':
+                        current_header_content[-1]['content'] += " " + element['content']
+                    else:
+                        current_header_content.append({
+                            'type': 'text',
+                            'content': element['content']
+                        })
+                elif element['type'] == 'image':
+                    current_header_content.append({
+                        'type': 'image',
+                        'path': element['path']
+                    })
+
+        if mode == 'headerwise':
+            add_current_header_content()
+
+    pdf_document.close()
+
+    json_output_path = os.path.join(output_folder, 'extraction_data.json')
+    with open(json_output_path, 'w', encoding='utf-8') as json_file:
+        json.dump(extraction_data, json_file, ensure_ascii=False, indent=4)
+
+    # Save to XLSX
+    df = pd.json_normalize(extraction_data, sep='_')
+    xlsx_output_path = os.path.join(output_folder, 'extraction_data.xlsx')
+    df.to_excel(xlsx_output_path, index=False)
+
+    # Save to Pickle
+    pickle_output_path = os.path.join(output_folder, 'extraction_data.pkl')
+    with open(pickle_output_path, 'wb') as pickle_file:
+        pickle.dump(extraction_data, pickle_file)
+
+    # Create ZIP file
+    zip_output_path = os.path.join(output_folder, 'extraction_data.zip')
+    with zipfile.ZipFile(zip_output_path, 'w') as zipf:
+        zipf.write(json_output_path, os.path.basename(json_output_path))
+        zipf.write(xlsx_output_path, os.path.basename(xlsx_output_path))
+        zipf.write(pickle_output_path, os.path.basename(pickle_output_path))
+        if extraction_type in ('images', 'both'):
+            for root, _, files in os.walk(output_folder):
+                for file in files:
+                    if file.endswith('.png'):
+                        zipf.write(os.path.join(root, file), file)
+
+    return json_output_path, xlsx_output_path, pickle_output_path, zip_output_path
+
+def render_pdf_page_as_image(pdf_path: str, page_number: int, zoom: float = 2.0) -> io.BytesIO:
+    # Render PDF page as an image
+    pdf_document = fitz.open(pdf_path)
+    page = pdf_document.load_page(page_number - 1)  # Page number is zero-indexed in fitz
+    pix = page.get_pixmap(matrix=fitz.Matrix(zoom, zoom))
+    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+    img_bytes = io.BytesIO()
+    img.save(img_bytes, format="PNG")
+    img_bytes.seek(0)
+    pdf_document.close()
+    return img_bytes
+
+# Streamlit UI
+
+st.markdown("<h1 style='text-align: center; color: blue;'>PDF DATA SNATCHER: HEADERWISE</h1>", unsafe_allow_html=True)
+st.markdown("<h3 style='text-align: center;color: brown;'>Extract valuable text and images from PDFs effortlessly and Convert PDFs into editable text and high-quality images </h3>", unsafe_allow_html=True)
+
+with st.expander("Click here for more information"):
+    st.write("""
+        **This application allows you to extract text and images from PDF files Headerswise. You can choose to extract only text, only images, or both, and the extracted data can be downloaded in JSON or XLSX format. Additionally, if you choose to extract images, you can download a ZIP file containing both the images and the JSON data.**
+        - **What is different about this app?**    
+        - 1. The sequence of text and images will get maintained as per its order in pdf file
+        - 2. You have options to extract entities from pdf
+        - 3. You can download data in JSON or XLSX format
+        - **PDF Preview:** You can preview a few pages of the uploaded PDF in the sidebar.
+        - **Extraction Type:** Choose whether to extract text, images, or both.
+        - **Minimum Font Size:** Set a threshold for the font size; text below this size will be ignored during extraction.
+        - **Output:** Download the extracted data as a JSON file, an Excel file, or a ZIP file (if images are included).
+        - *AUTHOR : CHINMAY BHALERAO*
+        """)
+
+    st.sidebar.markdown(
+        """
+        <div style="background-color: lightgray; padding: 2px; border-radius: 2px; text-align: center;">
+        <h2 style="color: blue; margin: 0;">PDF PREVIEW</h2>
+        </div>
+        """, unsafe_allow_html=True)
+
+# Upload PDF file
+uploaded_file = st.file_uploader("Upload a PDF file", type=["pdf"])
+
+if uploaded_file is not None:
+    # Create a temporary directory
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Save uploaded file to a temporary path
+        temp_pdf_path = os.path.join(temp_dir, "temp_uploaded_file.pdf")
+        with open(temp_pdf_path, "wb") as f:
+            f.write(uploaded_file.getbuffer())
+
+        # Number of pages to preview
+        num_pages_to_preview = st.sidebar.slider("Number of Pages to Preview", min_value=1, max_value=5, value=1)
+
+        # Generate and display thumbnails for selected number of pages
+        st.sidebar.write("Preview of Uploaded PDF:")
+        for page_number in range(1, num_pages_to_preview + 1):
+            thumbnail_bytes = render_pdf_page_as_image(temp_pdf_path, page_number=page_number)
+            st.sidebar.image(thumbnail_bytes, caption=f"Page {page_number}", width=300)
+
+        # Extraction type
+        st.info("You can select **only text** or **only images** or **text and images both** to extract form pdf")
+
+        extraction_type = st.radio("Extraction Type", options=['text', 'images', 'both'])
+
+        # Headers are capital
+        headers_are_capital = st.checkbox("Are Headers in Capital Letters?", value=False)
+
+        # Minimum font size
+        st.info("Minimum font size is the size below which size, the text will get ignored for extraction")
+
+        minimum_font_size = st.slider("Minimum Font Size", min_value=8, max_value=20, value=10)
+
+        # Header font sizes
+        header_font_sizes_input = st.text_input(
+            "Header Font Sizes (comma-separated, e.g., 10, 12.5, 14.75)", value="16.0")
+        header_font_sizes = [float(size.strip()) for size in header_font_sizes_input.split(',') if size.strip().replace('.', '', 1).isdigit()]
+
+        if st.button("Start Extraction"):
+            # Run extraction
+            json_output_path, xlsx_output_path, pickle_output_path, zip_output_path = extract_text_images(
+                pdf_path=temp_pdf_path,
+                output_folder=temp_dir,
+                minimum_font_size=minimum_font_size,
+                mode='headerwise',  # Pagewise mode has been removed
+                header_font_sizes=header_font_sizes,
+                extraction_type=extraction_type,
+                headers_are_capital=headers_are_capital
+            )
+
+            st.success("Extraction complete!")
+
+            # Display download options
+            with open(json_output_path, 'r', encoding='utf-8') as json_file:
+                extracted_data = json.load(json_file)
+            st.json(extracted_data)  # Display JSON data in Streamlit
+
+            with open(json_output_path, "rb") as file:
+                st.download_button(
+                    label="Download JSON",
+                    data=file,
+                    file_name="extraction_data.json",
+                    mime="application/json"
+                )
+
+            with open(xlsx_output_path, "rb") as file:
+                st.download_button(
+                    label="Download XLSX",
+                    data=file,
+                    file_name="extraction_data.xlsx",
+                    mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
+                )
+
+            with open(pickle_output_path, "rb") as file:
+                st.download_button(
+                    label="Download Pickle",
+                    data=file,
+                    file_name="extraction_data.pkl",
+                    mime="application/octet-stream"
+                )
+
+            if extraction_type in ('images', 'both'):
+                with open(zip_output_path, "rb") as file:
+                    st.download_button(
+                        label="Download ZIP",
+                        data=file,
+                        file_name="extraction_data.zip",
+                        mime="application/zip"
+                    )