Update app.py

app.py

@@ -1,9 +1,8 @@
 import streamlit as st
 import fitz  # PyMuPDF
-import pandas as pd
 import os
 
-# Streamlit UI
+# Configure Streamlit page
 st.set_page_config(page_title="PDF to Structured Data", layout="centered")
 
 # Custom Styling
@@ -21,15 +20,23 @@ st.markdown(
     unsafe_allow_html=True,
 )
 
+# App title
 st.markdown("<h1 style='text-align: center;'>📄 PDF to Structured Data</h1>", unsafe_allow_html=True)
-st.markdown("<p style='text-align: center; color: gray;'>powered by Google DeepMind Gemini 2.0 Flash</p>", unsafe_allow_html=True)
+st.markdown("<p style='text-align: center; color: gray;'>powered by AI</p>", unsafe_allow_html=True)
 
 # File uploader
-uploaded_file = st.file_uploader("Drop your PDF here or click to browse", type=["pdf"], help="Maximum file size: 100MB")
+uploaded_file = st.file_uploader(
+    "Drop your PDF here or click to browse", 
+    type=["pdf"], 
+    help="Maximum file size: 100MB"
+)
 
-# Text input for structure description
+# User input for data structure description
 st.markdown("<h4 style='color: #4A90E2;'>Describe the structure and type of data you want to extract from the PDF.</h4>", unsafe_allow_html=True)
-data_description = st.text_area("Example: Extract all invoice details including invoice number, date, items, prices, and total amount...", "")
+data_description = st.text_area(
+    "Example: Extract all invoice details including invoice number, date, items, prices, and total amount...",
+    ""
+)
 
 # Function to extract structured data
 def extract_text_and_structure(pdf_path):
@@ -37,18 +44,33 @@ def extract_text_and_structure(pdf_path):
     structured_data = {"title": "", "headings": [], "paragraphs": []}
 
     for page in doc:
-        text = page.get_text("text")
         blocks = page.get_text("blocks")
 
         for block in blocks:
             content = block[4].strip()
             if not content:
-                continue
-
-            font_size = page.get_text("dict")["blocks"][blocks.index(block)]["lines"][0]["spans"][0]["size"]
-            if font_size > 14:
-                structured_data["headings"].append(content)
-            else:
+                continue  # Skip empty blocks
+
+            # Get font size safely
+            try:
+                text_dict = page.get_text("dict")
+                block_index = blocks.index(block)
+
+                block_data = text_dict["blocks"][block_index] if "blocks" in text_dict and block_index < len(text_dict["blocks"]) else {}
+
+                if "lines" in block_data and block_data["lines"]:
+                    if "spans" in block_data["lines"][0] and block_data["lines"][0]["spans"]:
+                        font_size = block_data["lines"][0]["spans"][0]["size"]
+                        if font_size > 14:
+                            structured_data["headings"].append(content)
+                        else:
+                            structured_data["paragraphs"].append(content)
+                    else:
+                        structured_data["paragraphs"].append(content)
+                else:
+                    structured_data["paragraphs"].append(content)
+
+            except Exception as e:
                 structured_data["paragraphs"].append(content)
 
     if structured_data["headings"]:
@@ -57,25 +79,31 @@ def extract_text_and_structure(pdf_path):
     return structured_data
 
 # Extract Data Button
-extract_button = st.button("Extract Data")
-
-if extract_button and uploaded_file:
-    with open("temp.pdf", "wb") as f:
-        f.write(uploaded_file.getbuffer())
-
-    st.success("PDF Uploaded Successfully!")
-    
-    extracted_data = extract_text_and_structure("temp.pdf")
-
-    st.subheader("Title")
-    st.write(extracted_data["title"])
-
-    st.subheader("Headings")
-    for heading in extracted_data["headings"]:
-        st.write(f"- {heading}")
-
-    st.subheader("Paragraphs")
-    for para in extracted_data["paragraphs"]:
-        st.write(para)
-
-    os.remove("temp.pdf")
+if st.button("Extract Data", use_container_width=True):
+    if uploaded_file is not None:
+        with st.spinner("Processing your PDF..."):
+            temp_path = "temp.pdf"
+            with open(temp_path, "wb") as f:
+                f.write(uploaded_file.getbuffer())
+
+            extracted_data = extract_text_and_structure(temp_path)
+            os.remove(temp_path)
+
+            # Display extracted data
+            st.success("✅ Extraction Complete!")
+
+            if extracted_data["title"]:
+                st.subheader("📌 Title")
+                st.write(extracted_data["title"])
+
+            if extracted_data["headings"]:
+                st.subheader("📑 Headings")
+                for heading in extracted_data["headings"]:
+                    st.write(f"- {heading}")
+
+            if extracted_data["paragraphs"]:
+                st.subheader("📖 Paragraphs")
+                for para in extracted_data["paragraphs"]:
+                    st.write(para)
+    else:
+        st.warning("⚠️ Please upload a PDF file before extracting data.")