Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import json | |
| import os | |
| from pathlib import Path | |
| import base64 | |
| from typing import List, Dict, Any | |
| import google.generativeai as genai | |
| from PIL import Image | |
| import PyPDF2 | |
| import io | |
| # Configure Gemini API | |
| # GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") | |
| GEMINI_API_KEY = "AIzaSyB2b80YwNHs3Yj6RZOTL8wjXk2YhxCluOA" | |
| if GEMINI_API_KEY: | |
| genai.configure(api_key=GEMINI_API_KEY) | |
| EXTRACTION_PROMPT = """You are a shipping document data extraction specialist. Extract structured data from the provided shipping/logistics documents. | |
| Extract the following fields into a JSON format: | |
| { | |
| "poNumber": "Purchase Order Number", | |
| "shipFrom": "Origin/Ship From Location", | |
| "carrierType": "Transportation type (RAIL/TRUCK/etc)", | |
| "originCarrier": "Carrier name (CN/CPRS/etc)", | |
| "railCarNumber": "Rail car identifier", | |
| "totalQuantity": "Total quantity as number", | |
| "totalUnits": "Unit type (UNIT/MBF/MSFT/etc)", | |
| "accountName": "Customer/Account name", | |
| "inventories": { | |
| "items": [ | |
| { | |
| "quantityShipped": "Quantity as number", | |
| "inventoryUnits": "Unit type", | |
| "productName": "Full product description", | |
| "productCode": "Product code/SKU", | |
| "product": { | |
| "category": "Product category (OSB/Lumber/etc)", | |
| "unit": "Unit count as number", | |
| "pcs": "Pieces per unit", | |
| "mbf": "Thousand board feet (if applicable)", | |
| "sf": "Square feet (if applicable)", | |
| "pcsHeight": "Height in inches", | |
| "pcsWidth": "Width in inches", | |
| "pcsLength": "Length in feet" | |
| }, | |
| "customFields": [ | |
| "Mill||Mill Name", | |
| "Vendor||Vendor Name" | |
| ] | |
| } | |
| ] | |
| } | |
| } | |
| IMPORTANT INSTRUCTIONS: | |
| 1. Extract ALL products/items found in the document | |
| 2. Convert text numbers to actual numbers (e.g., "54" → 54) | |
| 3. Parse dimensions carefully (e.g., "2X6X14" means height=2, width=6, length=14) | |
| 4. Calculate MBF/SF when possible from dimensions and piece count | |
| 5. If a field is not found, use null (not empty string) | |
| 6. For multiple products, create separate items in the inventories.items array | |
| 7. Extract custom fields like Mill, Vendor from document metadata | |
| Return ONLY valid JSON, no markdown formatting or explanations.""" | |
| def extract_text_from_pdf(pdf_file) -> str: | |
| """Extract text from PDF file""" | |
| try: | |
| pdf_reader = PyPDF2.PdfReader(pdf_file) | |
| text = "" | |
| for page in pdf_reader.pages: | |
| text += page.extract_text() + "\n" | |
| return text | |
| except Exception as e: | |
| return f"Error extracting PDF text: {str(e)}" | |
| def convert_pdf_to_images(pdf_file) -> List[Image.Image]: | |
| """Convert PDF pages to images""" | |
| try: | |
| from pdf2image import convert_from_path | |
| images = convert_from_path(pdf_file) | |
| return images | |
| except ImportError: | |
| # Fallback if pdf2image not available | |
| return [] | |
| except Exception as e: | |
| print(f"Error converting PDF to images: {e}") | |
| return [] | |
| def process_files(files: List[str]) -> Dict[str, Any]: | |
| """Process uploaded files and extract text/images""" | |
| processed_data = { | |
| "files": [], | |
| "combined_text": "", | |
| "images": [] | |
| } | |
| if not files: | |
| return processed_data | |
| for file_path in files: | |
| file_name = Path(file_path).name | |
| file_ext = Path(file_path).suffix.lower() | |
| file_data = { | |
| "filename": file_name, | |
| "type": file_ext, | |
| "content": "" | |
| } | |
| try: | |
| if file_ext == '.pdf': | |
| # Extract text from PDF | |
| text = extract_text_from_pdf(file_path) | |
| file_data["content"] = text | |
| processed_data["combined_text"] += f"\n--- {file_name} ---\n{text}\n" | |
| # Try to convert PDF to images for visual extraction | |
| images = convert_pdf_to_images(file_path) | |
| processed_data["images"].extend(images) | |
| elif file_ext in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']: | |
| # Load image | |
| img = Image.open(file_path) | |
| processed_data["images"].append(img) | |
| file_data["content"] = f"Image file: {file_name}" | |
| processed_data["combined_text"] += f"\n--- {file_name} (Image) ---\n" | |
| elif file_ext in ['.txt']: | |
| # Read text file | |
| with open(file_path, 'r', encoding='utf-8') as f: | |
| text = f.read() | |
| file_data["content"] = text | |
| processed_data["combined_text"] += f"\n--- {file_name} ---\n{text}\n" | |
| processed_data["files"].append(file_data) | |
| except Exception as e: | |
| file_data["content"] = f"Error processing file: {str(e)}" | |
| processed_data["files"].append(file_data) | |
| return processed_data | |
| def extract_with_gemini(processed_data: Dict[str, Any], api_key: str) -> Dict[str, Any]: | |
| """Extract structured data using Gemini API""" | |
| if not api_key: | |
| return {"error": "Gemini API key not provided"} | |
| try: | |
| # Configure Gemini with provided API key | |
| genai.configure(api_key=api_key) | |
| # Use Gemini Pro Vision if images available, otherwise Pro | |
| if processed_data["images"]: | |
| model = genai.GenerativeModel('gemini-1.5-flash') | |
| # Prepare content with images and text | |
| content = [EXTRACTION_PROMPT] | |
| # Add text content | |
| if processed_data["combined_text"]: | |
| content.append(f"\nDocument Text:\n{processed_data['combined_text']}") | |
| # Add images (limit to first 5 for token management) | |
| for img in processed_data["images"][:5]: | |
| content.append(img) | |
| response = model.generate_content(content) | |
| else: | |
| # Text-only extraction | |
| model = genai.GenerativeModel('gemini-1.5-flash') | |
| prompt = f"{EXTRACTION_PROMPT}\n\nDocument Text:\n{processed_data['combined_text']}" | |
| response = model.generate_content(prompt) | |
| # Parse response | |
| response_text = response.text.strip() | |
| # Remove markdown code blocks if present | |
| if response_text.startswith("```json"): | |
| response_text = response_text[7:] | |
| if response_text.startswith("```"): | |
| response_text = response_text[3:] | |
| if response_text.endswith("```"): | |
| response_text = response_text[:-3] | |
| response_text = response_text.strip() | |
| # Parse JSON | |
| extracted_data = json.loads(response_text) | |
| return { | |
| "success": True, | |
| "data": extracted_data, | |
| "raw_response": response_text | |
| } | |
| except json.JSONDecodeError as e: | |
| return { | |
| "success": False, | |
| "error": f"JSON parsing error: {str(e)}", | |
| "raw_response": response.text if 'response' in locals() else None | |
| } | |
| except Exception as e: | |
| return { | |
| "success": False, | |
| "error": f"Extraction error: {str(e)}" | |
| } | |
| def process_documents(files, api_key): | |
| """Main processing function""" | |
| if not files: | |
| return "⚠️ Please upload at least one document.", None, None | |
| if not api_key: | |
| return "⚠️ Please provide your Gemini API key.", None, None | |
| try: | |
| # Step 1: Process files | |
| status = "📄 Processing files..." | |
| processed_data = process_files(files) | |
| # Step 2: Extract with Gemini | |
| status = "🤖 Extracting data with Gemini AI..." | |
| result = extract_with_gemini(processed_data, api_key) | |
| if result.get("success"): | |
| # Format JSON output | |
| json_output = json.dumps(result["data"], indent=2) | |
| # Create summary | |
| data = result["data"] | |
| summary = f"""✅ **Extraction Successful!** | |
| **Shipment Details:** | |
| - PO Number: {data.get('poNumber', 'N/A')} | |
| - Ship From: {data.get('shipFrom', 'N/A')} | |
| - Carrier: {data.get('originCarrier', 'N/A')} ({data.get('carrierType', 'N/A')}) | |
| - Rail Car: {data.get('railCarNumber', 'N/A')} | |
| - Total Quantity: {data.get('totalQuantity', 'N/A')} {data.get('totalUnits', '')} | |
| - Account: {data.get('accountName', 'N/A')} | |
| **Products Found:** {len(data.get('inventories', {}).get('items', []))} | |
| """ | |
| return summary, json_output, None | |
| else: | |
| error_msg = f"❌ **Extraction Failed**\n\nError: {result.get('error', 'Unknown error')}" | |
| raw = result.get('raw_response', 'No response') | |
| return error_msg, None, raw | |
| except Exception as e: | |
| return f"❌ **Processing Error**\n\n{str(e)}", None, None | |
| # Create Gradio Interface | |
| with gr.Blocks(title="Shipping Document Extractor", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown(""" | |
| # 📦 Shipping Document Data Extractor | |
| Upload shipping documents (PDFs, images, etc.) to automatically extract structured data. | |
| Powered by Google Gemini AI. | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| api_key_input = gr.Textbox( | |
| label="🔑 Gemini API Key", | |
| placeholder="Enter your Gemini API key", | |
| type="password", | |
| value=GEMINI_API_KEY or "" | |
| ) | |
| gr.Markdown("[Get your API key](https://makersuite.google.com/app/apikey)") | |
| file_input = gr.File( | |
| label="📁 Upload Documents", | |
| file_count="multiple", | |
| file_types=[".pdf", ".jpg", ".jpeg", ".png", ".txt"] | |
| ) | |
| process_btn = gr.Button("🚀 Extract Data", variant="primary", size="lg") | |
| gr.Markdown(""" | |
| ### Supported Files: | |
| - PDF documents | |
| - Images (JPG, PNG) | |
| - Text files | |
| """) | |
| with gr.Column(scale=2): | |
| status_output = gr.Markdown(label="Status") | |
| with gr.Tabs(): | |
| with gr.Tab("📊 Extracted JSON"): | |
| json_output = gr.Code( | |
| label="Structured Data (JSON)", | |
| language="json", | |
| lines=20 | |
| ) | |
| download_btn = gr.DownloadButton( | |
| label="💾 Download JSON", | |
| visible=False | |
| ) | |
| with gr.Tab("🔍 Raw Response"): | |
| raw_output = gr.Code( | |
| label="Raw API Response", | |
| language="text", | |
| lines=20 | |
| ) | |
| # Event handlers | |
| def update_download(json_str): | |
| if json_str: | |
| return gr.DownloadButton(visible=True, value=json_str) | |
| return gr.DownloadButton(visible=False) | |
| process_btn.click( | |
| fn=process_documents, | |
| inputs=[file_input, api_key_input], | |
| outputs=[status_output, json_output, raw_output] | |
| ) | |
| json_output.change( | |
| fn=update_download, | |
| inputs=[json_output], | |
| outputs=[download_btn] | |
| ) | |
| gr.Markdown(""" | |
| --- | |
| ### 💡 Tips: | |
| - Upload multiple documents at once for batch processing | |
| - Better quality images/PDFs = better extraction accuracy | |
| - The AI extracts: PO numbers, carrier info, products, quantities, dimensions, and more | |
| """) | |
| # Launch the app | |
| if __name__ == "__main__": | |
| demo.launch() |