Update app.py

app.py

@@ -1,64 +1,67 @@
+import streamlit as st
+import pdfplumber
 import pandas as pd
-import fitz  # PyMuPDF for PDF processing
-from transformers import AutoProcessor, AutoModelForDocumentQuestionAnswering
-import torch
-from PIL import Image
+from io import BytesIO
 
-# Load Hugging Face OCR model (e.g., LayoutLM)
-processor = AutoProcessor.from_pretrained("microsoft/layoutlmv3-base")
-model = AutoModelForDocumentQuestionAnswering.from_pretrained("microsoft/layoutlmv3-base")
+def is_header(text):
+    """Identify headers using common keywords."""
+    keywords = ['Purchase Order', 'Supplier Order', 'GSTIN', 'Annexure', 'Terms', 'Currency']
+    return any(keyword in text for keyword in keywords)
 
-# Function to extract text from PDF using OCR
-def extract_text_from_pdf(pdf_path):
-    pdf_document = fitz.open(pdf_path)
-    all_text = []
+def extract_cleaned_tables(pdf_file):
+    """Extract tables while skipping headers and arranging them by pages."""
+    tables = []
     
-    # Loop through PDF pages
-    for page_num in range(len(pdf_document)):
-        page = pdf_document.load_page(page_num)
-        pix = page.get_pixmap()  # Convert PDF page to image
-        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-        
-        # Preprocess image for Hugging Face model
-        inputs = processor(images=img, return_tensors="pt")
-        
-        # Model inference
-        with torch.no_grad():
-            outputs = model(**inputs)
-        
-        # Extract and decode text
-        text = processor.decode(outputs.logits.argmax(-1).squeeze().tolist())
-        all_text.append(text)
+    with pdfplumber.open(pdf_file) as pdf:
+        for page_num, page in enumerate(pdf.pages):
+            text = page.extract_text()
+            
+            # Skip pages with header-heavy content
+            if is_header(text):
+                continue
+            
+            page_tables = page.extract_tables()
+            for table in page_tables:
+                if table:
+                    df = pd.DataFrame(table[1:], columns=table[0])
+                    # Fix misalignment issues (if 'Unit' in wrong columns, move it)
+                    if 'Delivery Date' in df.columns and 'Unit' in df.columns:
+                        mask = df['Delivery Date'].str.contains(r'NOS|PCS', na=False)
+                        df.loc[mask, 'Unit'] = df.loc[mask, 'Delivery Date']
+                        df.loc[mask, 'Delivery Date'] = None
+                    
+                    tables.append((f"Page_{page_num+1}", df))
     
-    return all_text
+    return tables
 
-# Parse text to extract relevant information (example logic, modify based on PO structure)
-def parse_po_details(extracted_text):
-    po_data = []
-    
-    for line in extracted_text:
-        # Example: Extract lines related to material, quantity, rate, and value
-        if "Material" in line or "Quantity" in line or "Rate" in line or "Value" in line:
-            po_data.append(line.strip())
-    
-    return po_data
-
-# Function to export data to CSV
-def export_to_csv(po_data, output_csv_path):
-    df = pd.DataFrame(po_data, columns=["PO Details"])
-    df.to_csv(output_csv_path, index=False)
-
-# Path to your PDF file
-pdf_path = 'your_pdf_path_here.pdf'
+# Streamlit App
+st.title("Enhanced PO Extraction Tool")
 
-# Extract text using Hugging Face OCR
-extracted_text = extract_text_from_pdf(pdf_path)
+uploaded_file = st.file_uploader("Upload PO PDF", type=["pdf"])
 
-# Parse the extracted text to get relevant PO details
-po_data = parse_po_details(extracted_text)
+if uploaded_file:
+    try:
+        # Extract and clean tables from the uploaded PDF
+        extracted_tables = extract_cleaned_tables(uploaded_file)
 
-# Export the parsed data to a CSV file
-output_csv_path = 'requirement.csv'
-export_to_csv(po_data, output_csv_path)
+        if extracted_tables:
+            st.success("Tables extracted successfully!")
+            
+            # Create an Excel file with multiple sheets
+            excel_buffer = BytesIO()
+            with pd.ExcelWriter(excel_buffer, engine='openpyxl') as writer:
+                for sheet_name, df in extracted_tables:
+                    df.to_excel(writer, index=False, sheet_name=sheet_name)
+            excel_buffer.seek(0)
 
-print(f"Purchase order details have been exported to {output_csv_path}")
+            # Provide download options
+            st.download_button(
+                label="Download as Excel",
+                data=excel_buffer,
+                file_name="po_data.xlsx",
+                mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
+            )
+        else:
+            st.warning("No valid tables found.")
+    except Exception as e:
+        st.error(f"An error occurred: {e}")