Upload app (2).py

app (2).py

@@ -0,0 +1,185 @@
+import gradio as gr
+import json
+from pathlib import Path
+from typing import List, Dict, Any
+from PIL import Image
+import PyPDF2
+import pytesseract
+import google.generativeai as genai
+import tempfile
+
+
+# ==================== Configure Gemini API ====================
+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, Do NOT convert units
+4. Calculate MBF/SF when possible from dimensions and piece count
+5. If a field is not found, use null
+6. For multiple products, create separate items
+7. Extract custom fields like Mill, Vendor
+Return ONLY valid JSON, no markdown formatting or explanations."""
+
+# ==================== Utility functions ====================
+def extract_text_from_pdf(pdf_file) -> str:
+    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]:
+    try:
+        from pdf2image import convert_from_path
+        images = convert_from_path(pdf_file)
+        return images
+    except Exception as e:
+        print(f"Error converting PDF to images: {e}")
+        return []
+
+def extract_text_from_image(img: Image.Image) -> str:
+    try:
+        text = pytesseract.image_to_string(img)
+        return text
+    except Exception as e:
+        print(f"Error extracting text from image: {e}")
+        return ""
+
+def process_files(files: List[str]) -> Dict[str, Any]:
+    processed_data = {
+        "files": [],
+        "combined_text": "",
+        "images": []
+    }
+    
+    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':
+                text = extract_text_from_pdf(file_path)
+                file_data["content"] = text
+                processed_data["combined_text"] += f"\n--- {file_name} ---\n{text}\n"
+                images = convert_pdf_to_images(file_path)
+                processed_data["images"].extend(images)
+                
+            elif file_ext in ['.jpg', '.jpeg', '.png', '.bmp', '.gif']:
+                img = Image.open(file_path)
+                processed_data["images"].append(img)
+                text = extract_text_from_image(img)
+                processed_data["combined_text"] += f"\n--- {file_name} ---\n{text}\n"
+                file_data["content"] = f"Image file: {file_name}"
+                
+            elif file_ext in ['.txt']:
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    text = f.read()
+                processed_data["combined_text"] += f"\n--- {file_name} ---\n{text}\n"
+                file_data["content"] = text
+            
+            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]) -> Dict[str, Any]:
+    try:
+        model = genai.GenerativeModel('models/gemini-2.5-flash')
+        content = [EXTRACTION_PROMPT]
+        if processed_data["combined_text"]:
+            content.append(f"\nDocument Text:\n{processed_data['combined_text']}")
+        for img in processed_data["images"][:5]:
+            content.append(img)
+        response = model.generate_content(content)
+        response_text = response.text.strip()
+        # Clean Markdown
+        for mark in ["```json", "```"]:
+            response_text = response_text.replace(mark, "")
+        extracted_data = json.loads(response_text)
+        return extracted_data
+    except Exception as e:
+        return {"error": str(e)}
+
+# ==================== Gradio function ====================
+def gradio_extraction(uploaded_files):
+    file_paths = []
+
+    for file in uploaded_files:
+        src_path = Path(file.name)
+        file_name = src_path.name
+        tmp_path = Path(tempfile.gettempdir()) / file_name
+
+        with open(src_path, "rb") as src, open(tmp_path, "wb") as dst:
+            dst.write(src.read())
+
+        file_paths.append(str(tmp_path))
+    
+    processed_data = process_files(file_paths)
+    extracted_data = extract_with_gemini(processed_data)
+    
+    with open("output.json", "w", encoding="utf-8") as f:
+        json.dump(extracted_data, f, indent=2)
+    
+    return json.dumps(extracted_data, indent=2), "output.json"
+
+
+# ==================== Gradio Interface ====================
+iface = gr.Interface(
+    fn=gradio_extraction,
+    inputs = gr.File(file_types=[".pdf", ".jpg", ".jpeg", ".png", ".bmp", ".txt"], file_count="multiple"),
+    outputs=[
+        gr.Textbox(label="Extracted JSON",lines=15, max_lines=30), 
+        gr.File(label="Download JSON")
+    ],
+    title="Shipping Document Text Extractor",
+    description="Upload PDFs or images of shipping/logistics documents and get structured JSON output.",
+    theme=gr.themes.Base(primary_hue="blue")
+)
+
+iface.launch()