ChinmayBH's picture
created app.py
bd889e8 verified
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"
)