app.py

import streamlit as st
import os
import json
import fitz
from io import BytesIO
from PIL import Image
import pandas as pd
import zipfile
import tempfile

def extract_text_images(
        pdf_path: str, output_folder: str,
        minimum_font_size: int,
        extraction_type: str = 'both'
        ) -> dict:
    """
    Extracts text and/or images from a PDF and organizes them by pages.
    """
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    extraction_data = []
    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:
                    for line in block["lines"]:
                        for span in line["spans"]:
                            font_size = span["size"]
                            top = span["bbox"][1]

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

                elements.append({
                    'type': 'text',
                    'font_size': line[0]['size'],
                    'page': page_number + 1,
                    'content': line_text,
                    'x0': line[0]['bbox'][0],
                    'top': top,
                })

        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']))

        page_content = []
        for element in elements:
            if element['type'] == 'text':
                if page_content and page_content[-1]['type'] == 'text':
                    page_content[-1]['content'] += " " + element['content']
                else:
                    page_content.append({
                        'type': 'text',
                        'content': element['content']
                    })
            elif element['type'] == 'image':
                page_content.append({
                    'type': 'image',
                    'path': element['path']
                })

        extraction_data.append({
            'page': page_number + 1,
            'content': page_content
        })

    pdf_document.close()

    return extraction_data

def convert_to_xlsx(data: dict) -> BytesIO:
    """
    Converts the extracted data to an XLSX file.
    """
    rows = []

    for item in data:
        page_number = item['page']
        content_list = item['content']
        
        for content in content_list:
            if content['type'] == 'text':
                rows.append({
                    'Page': page_number,
                    'Content': content['content']
                })
            elif content['type'] == 'image':
                rows.append({
                    'Page': page_number,
                    'Content': f"[Image: {content['path']}]"
                })

    df = pd.DataFrame(rows)

    output = BytesIO()
    with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
        df.to_excel(writer, index=False, sheet_name='Extraction')

    output.seek(0)
    return output

def create_zip_with_json_and_images(output_folder, extraction_data):
    """
    Creates a ZIP file containing both images and JSON data.
    """
    zip_buffer = BytesIO()
    with zipfile.ZipFile(zip_buffer, "w") as zip_file:
        # Add JSON file
        json_data = json.dumps(extraction_data, ensure_ascii=False, indent=4).encode('utf-8')
        zip_file.writestr("extraction_data.json", json_data)
        
        # Add images
        for item in extraction_data:
            for content in item['content']:
                if content['type'] == 'image':
                    image_path = content['path']
                    image_name = os.path.basename(image_path)
                    zip_file.write(image_path, image_name)

    zip_buffer.seek(0)
    return zip_buffer

def main():
    st.markdown("<h1 style='text-align: center; color: blue;'>PDF DATA SNACHER:PAGEWISE</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 and organizes them by pages. 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)


    pdf_file = st.file_uploader("Upload PDF", type="pdf")

    if pdf_file is not None:
        num_pages_to_preview = st.sidebar.slider(
            "Select number of pages to preview:", 
            min_value=1, max_value=5, value=1
        )

        pdf_document = fitz.open(stream=pdf_file.read(), filetype="pdf")
        for page_num in range(min(num_pages_to_preview, pdf_document.page_count)):
            page = pdf_document.load_page(page_num)
            pix = page.get_pixmap()
            image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            st.sidebar.image(image, caption=f"Page {page_num + 1} Preview", use_column_width=True)

    st.info("You can select **only text** or **only images** or **text and images both** to extract form pdf")
    extraction_type = st.selectbox(
        "Choose extraction type:",
        ("text", "images", "both")
    )

    st.info("Minimum font size is the size below which size, the text will get ignored for extraction")
    minimum_font_size = st.number_input(
        "Minimum font size to extract:",
        min_value=1, value=2
    )

    output_folder = st.text_input("Output folder path:")

    if st.button("Start Extraction"):
        if pdf_file is not None and output_folder:
            with tempfile.TemporaryDirectory() as temp_dir:
                temp_pdf_path = os.path.join(temp_dir, pdf_file.name)
                with open(temp_pdf_path, "wb") as f:
                    f.write(pdf_file.getvalue())

                extraction_data = extract_text_images(
                    temp_pdf_path,
                    temp_dir,
                    minimum_font_size,
                    extraction_type
                )

                st.json(extraction_data)

                if extraction_type == 'images' or extraction_type == 'both':
                    zip_data = create_zip_with_json_and_images(temp_dir, extraction_data)
                    st.download_button(
                        label="Download ZIP",
                        data=zip_data,
                        file_name='extraction_data.zip',
                        mime='application/zip'
                    )

                xlsx_data = convert_to_xlsx(extraction_data)

                col1, col2 = st.columns(2)
                with col1:
                    st.download_button(
                        label="Download JSON",
                        data=json.dumps(extraction_data, ensure_ascii=False, indent=4).encode('utf-8'),
                        file_name='extraction_data.json',
                        mime='application/json'
                    )
                with col2:
                    st.download_button(
                        label="Download XLSX",
                        data=xlsx_data,
                        file_name='extraction_data.xlsx',
                        mime='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
                    )
        else:
            st.error("Please upload a PDF file and provide an output folder path.")

    st.markdown(
        """
        <style>
        .footer {
            position: fixed;
            bottom: 0;
            left: 0;
            width: 100%;
            background-color: #F0F0F0;
            font-family:cursive;
            text-align: right;
            padding: 5px 0;
            font-size:20px;
            font-weight: bold;
            color: #FF0000;
        }
        </style>
        <div class="footer">
            CREATED BY: CHINMAY BHALERAO
        </div>
        """,
        unsafe_allow_html=True
    )

if __name__ == "__main__":
    main()