import json import os from pathlib import Path from typing import List, Dict, Any import google.generativeai as genai from PIL import Image import PyPDF2 import pytesseract from doctr.io import DocumentFile from doctr.models import ocr_predictor # Optional: Gradio for a lightweight UI import gradio as gr # 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 number of packages", "totalUnits": "Unit type (UNIT/MBF/MSFT/etc)", "accountName": "Customer/Account name", "inventories": { "items": [ { "quantityShipped": "Quantity as number, no of packages", "inventoryUnits": "Unit type from document (MBF, FBM, SF, UNIT etc.)", "productName": "Full product description", "productCode": "Product code/SKU", "product": { "category": "Product category (OSB/Lumber/etc)", "unit": "Unit type from document (MBF, FBM, SF, UNIT etc.)", "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 the same unit as document" }, "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(e.g., "2x6x14" means height=6, width=14, length=2) 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 8. Unit types must be (PCS/PKG/MBF/MSFT/etc) Return ONLY valid JSON, no markdown formatting or explanations.""" # Temporary: print available models #for model in genai.list_models(): # print(model) 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: return [] except Exception as e: print(f"Error converting PDF to images: {e}") return [] def extract_text_from_image(img_path: str) -> str: """Extract text using DocTR for better structure""" try: doc = DocumentFile.from_images(img_path) result = ocr_model(doc) export = result.export() lines = [] # Collect line-wise text preserving order for page in export['pages']: for block in page['blocks']: for line in block['lines']: line_text = " ".join([w['value'] for w in line['words']]) lines.append(line_text) return "\n".join(lines) except Exception as e: print(f"Error extracting text from image {img_path}: {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': 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', '.gif', '.bmp']: 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" # ===== Add OCR here ===== text = pytesseract.image_to_string(img) processed_data["combined_text"] += f"\n--- {file_name} (Image) ---\n{text}\n" elif file_ext in ['.txt']: 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: genai.configure(api_key=api_key) #model = genai.GenerativeModel('gemini-1.5-flash') model = genai.GenerativeModel('models/gemini-2.5-flash') # recommended # ya phir agar Cloud AI API hai to 'text-bison-001' print("available models : ", genai.list_models()) # Prepare content 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() 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] 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: print("⚠️ Please provide at least one document.") return if not api_key: print("⚠️ Please provide your Gemini API key.") return # Step 1: Process files print("📄 Processing files...") processed_data = process_files(files) # Step 2: Extract with Gemini print("🤖 Extracting data with Gemini AI...") result = extract_with_gemini(processed_data, api_key) if result.get("success"): json_output = json.dumps(result["data"], indent=2) print(" Extraction Successful!") print(json_output) # ===== Save to output.json ===== output_file = "output.json" with open(output_file, "w", encoding="utf-8") as f: f.write(json_output) print(f"JSON saved to {output_file}") return json_output else: print(f" Extraction Failed: {result.get('error', 'Unknown error')}") print("Raw Response:", result.get('raw_response', 'No response')) return None # --------------------------- # Lightweight web UI wrapper # --------------------------- # This UI layer calls the exact same processing functions above. # It does not modify extraction logic, only provides a user-friendly front end. def _gradio_wrapper(uploaded_files): """ uploaded_files: list of temporary file dicts that Gradio provides. Returns: status_message, json_text, preview_text """ if not uploaded_files: return ("No files uploaded.", "{}", "") # Map Gradio file objects to file paths that process_documents expects file_paths = [] for f in uploaded_files: # Gradio supplies a dict-like object with 'name' pointing to the temp path # Accept either direct path or dict with 'name' if isinstance(f, str) and os.path.exists(f): file_paths.append(f) else: # f may be a tempfile-like object or dict try: temp_path = f.name # file-like object if os.path.exists(temp_path): file_paths.append(temp_path) else: # attempt to copy bytes to a local temp file content = None if hasattr(f, "read"): content = f.read() elif isinstance(f, dict) and "name" in f: file_paths.append(f["name"]) continue if content: # create a temp file tmp_dir = Path("gradio_tmp") tmp_dir.mkdir(exist_ok=True) dest = tmp_dir / Path(f.name).name with open(dest, "wb") as out: out.write(content) file_paths.append(str(dest)) except Exception: # last-resort: try to interpret as path string try: if isinstance(f, dict) and "name" in f and os.path.exists(f["name"]): file_paths.append(f["name"]) except Exception: pass if not file_paths: return ("Uploaded files could not be located.", "{}", "") status_msg = "Processing..." # Call the existing processing pipeline (no changes) json_result = process_documents(file_paths, GEMINI_API_KEY) if json_result: # process_documents returns JSON string on success pretty = json_result try: parsed = json.loads(pretty) preview = "" # build a compact preview: show PO and first product name if available po = parsed.get("poNumber") inv = parsed.get("inventories", {}).get("items", []) first_prod = inv[0].get("productName") if inv else None preview = f"PO: {po}\nFirst product: {first_prod}" except Exception: preview = pretty[:100] + "..." return ("Extraction completed.", pretty, preview) else: return ("Extraction failed. Check console for details.", "{}", "") def build_ui(): """Create a simple web UI that uses the same processing code above.""" with gr.Blocks() as ui: gr.Markdown("## Document Extractor — Upload files to extract structured shipping data") gr.Markdown(""" ### 💡 Tips: - Upload multiple files for batch processing - For images: ensure text is clear and well-lit - For PDFs: both text-based and scanned PDFs work - The AI will analyze visual content even if text extraction fails """) with gr.Row(): with gr.Column(scale=2): file_input = gr.File( label="Select documents (PDF, image, text)", file_count="multiple", file_types=[".pdf", ".jpg", ".jpeg", ".png", ".gif", ".bmp", ".txt", ".csv", ".doc", ".docx"] ) run_btn = gr.Button("Extract", variant="primary") with gr.Column(scale=3): status = gr.Textbox(label="Status", lines=2) output_json = gr.Code(label="Extracted JSON", language="json", lines=20) preview = gr.Textbox(label="Quick preview", lines=4) run_btn.click(fn=_gradio_wrapper, inputs=[file_input], outputs=[status, output_json, preview]) return ui if __name__ == "__main__": # Keep the original hardcoded call unchanged for CLI usage files_to_process = ["sample1.pdf"] # Replace with your PDF paths # Run CLI extraction (preserves original behavior) process_documents(files_to_process, GEMINI_API_KEY) # Launch the UI (optional). Comment out the next lines if you don't want the web UI. demo = build_ui() demo.launch(server_name="0.0.0.0", server_port=7860, share=False)