Spaces:
Running on Zero
Running on Zero
| import os | |
| import io | |
| import json | |
| import logging | |
| from typing import Optional | |
| from fastapi import FastAPI, UploadFile, File, HTTPException, status | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import JSONResponse | |
| import gradio as gr | |
| import pandas as pd | |
| from src.config import settings, logger | |
| from src.schemas import DocumentIntelligenceResult, HealthResponse, AnalyzeResponse | |
| from src.ocr_service import ocr_service | |
| from src.document_analyzer import document_analyzer | |
| from contextlib import asynccontextmanager | |
| async def lifespan(app: FastAPI): | |
| """Validate settings and initialize models on startup.""" | |
| settings.validate() | |
| logger.info("AI Document Intelligence service started successfully.") | |
| yield | |
| # Initialize FastAPI App | |
| app = FastAPI( | |
| title="AI Document Intelligence OCR API", | |
| description="Enterprise document analysis combining PaddleOCR and NVIDIA Text LLMs", | |
| version="1.0.0", | |
| lifespan=lifespan, | |
| ) | |
| # Enable CORS for cloud & cross-origin frontends | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| async def health_check(): | |
| """Health check endpoint indicating OCR engine readiness and NVIDIA config.""" | |
| return HealthResponse( | |
| status="ok", | |
| version="1.0.0", | |
| ocr_engine=ocr_service.active_engine, | |
| nvidia_model=settings.NVIDIA_MODEL, | |
| nvidia_configured=bool(settings.NVIDIA_API_KEY), | |
| ) | |
| async def analyze_document(file: UploadFile = File(...)): | |
| """ | |
| Accepts JPG/PNG image upload: | |
| 1. Validates file format and size. | |
| 2. Extracts text using PaddleOCR. | |
| 3. Analyzes structure and entities using NVIDIA LLM. | |
| 4. Returns structured JSON schema. | |
| """ | |
| if not file.content_type or not ( | |
| file.content_type.startswith("image/") | |
| or file.filename.lower().endswith((".jpg", ".jpeg", ".png", ".webp", ".bmp")) | |
| ): | |
| raise HTTPException( | |
| status_code=status.HTTP_400_BAD_REQUEST, | |
| detail="Invalid file format. Please upload a JPEG or PNG image." | |
| ) | |
| # Read image bytes | |
| contents = await file.read() | |
| file_size_mb = len(contents) / (1024 * 1024) | |
| if file_size_mb > settings.MAX_IMAGE_SIZE_MB: | |
| raise HTTPException( | |
| status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, | |
| detail=f"File size ({file_size_mb:.2f}MB) exceeds maximum limit of {settings.MAX_IMAGE_SIZE_MB}MB." | |
| ) | |
| try: | |
| result = document_analyzer.analyze(contents) | |
| return AnalyzeResponse( | |
| success=True, | |
| ocr_text=result.extracted_text, | |
| document_intelligence=result | |
| ) | |
| except ValueError as e: | |
| logger.error("Validation / API configuration error: %s", e) | |
| raise HTTPException( | |
| status_code=status.HTTP_400_BAD_REQUEST, | |
| detail=str(e) | |
| ) | |
| except RuntimeError as e: | |
| logger.error("Processing runtime error: %s", e) | |
| raise HTTPException( | |
| status_code=status.HTTP_502_BAD_GATEWAY, | |
| detail=str(e) | |
| ) | |
| except Exception as e: | |
| logger.error("Unexpected pipeline error: %s", e, exc_info=True) | |
| raise HTTPException( | |
| status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, | |
| detail="An internal error occurred during document intelligence processing." | |
| ) | |
| try: | |
| import spaces | |
| except ImportError: | |
| class spaces: | |
| def GPU(fn=None, **kwargs): | |
| if fn: | |
| return fn | |
| return lambda f: f | |
| def process_ui_image(image): | |
| """ZeroGPU-compatible top-level handler for Gradio document analysis.""" | |
| if image is None: | |
| return ( | |
| "Please upload an image.", | |
| "", | |
| None, | |
| None, | |
| None, | |
| None, | |
| "", | |
| "", | |
| {} | |
| ) | |
| try: | |
| res: DocumentIntelligenceResult = document_analyzer.analyze(image) | |
| # Format Entities table | |
| entities_df = pd.DataFrame([ | |
| {"Entity": e.name, "Value": e.value, "Type": e.type} | |
| for e in res.entities | |
| ]) if res.entities else pd.DataFrame(columns=["Entity", "Value", "Type"]) | |
| # Format Tables | |
| tables_df = None | |
| if res.tables and res.tables[0].headers: | |
| t = res.tables[0] | |
| tables_df = pd.DataFrame(t.rows, columns=t.headers) | |
| elif res.tables and res.tables[0].rows: | |
| tables_df = pd.DataFrame(res.tables[0].rows) | |
| # Format Components | |
| comp_df = pd.DataFrame([ | |
| {"Component": c.name, "Type": c.type, "Role": c.description or ""} | |
| for c in res.components | |
| ]) if res.components else pd.DataFrame(columns=["Component", "Type", "Role"]) | |
| # Format Relationships | |
| rel_df = pd.DataFrame([ | |
| {"Source": r.source, "Target": r.target, "Link Type": r.relationship_type, "Protocol/Label": r.label or ""} | |
| for r in res.relationships | |
| ]) if res.relationships else pd.DataFrame(columns=["Source", "Target", "Link Type", "Protocol/Label"]) | |
| notes_str = "\n".join([f"• {n}" for n in res.confidence_notes]) if res.confidence_notes else "• High extraction confidence." | |
| return ( | |
| f"📄 **Document Type:** `{res.document_type.upper()}`", | |
| res.summary, | |
| entities_df, | |
| tables_df, | |
| comp_df, | |
| rel_df, | |
| res.extracted_text, | |
| notes_str, | |
| res.model_dump() | |
| ) | |
| except Exception as err: | |
| logger.error("Gradio execution error: %s", err) | |
| return ( | |
| f"❌ Error: {str(err)}", | |
| "", | |
| None, | |
| None, | |
| None, | |
| None, | |
| "", | |
| f"• Failure occurred: {str(err)}", | |
| {"error": str(err)} | |
| ) | |
| def create_gradio_ui() -> gr.Blocks: | |
| """Build the Gradio interface for AI Document Intelligence.""" | |
| custom_css = """ | |
| .main-container { max-width: 1200px; margin: auto; } | |
| .header-box { text-align: center; padding: 20px; background: linear-gradient(135deg, #1e293b, #0f172a); color: white; border-radius: 12px; margin-bottom: 20px; } | |
| .header-box h1 { margin: 0; font-size: 2rem; font-weight: 700; } | |
| .header-box p { margin-top: 8px; color: #94a3b8; } | |
| .doc-badge { font-weight: bold; padding: 4px 12px; border-radius: 6px; background-color: #3b82f6; color: white; } | |
| """ | |
| with gr.Blocks(title="AI Document Intelligence & OCR", css=custom_css, theme=gr.themes.Soft()) as demo: | |
| with gr.Column(elem_classes=["main-container"]): | |
| gr.HTML( | |
| """ | |
| <div class="header-box"> | |
| <h1>🧠 AI Document Intelligence + OCR</h1> | |
| <p>Enterprise Multi-Modal Document Extraction & Architecture Understanding</p> | |
| </div> | |
| """ | |
| ) | |
| with gr.Row(): | |
| # LEFT SIDE: Upload & Trigger | |
| with gr.Column(scale=1): | |
| gr.Markdown("### 📥 Document Upload") | |
| image_input = gr.Image(type="pil", label="Upload Document (JPEG, PNG, Invoice, Architecture)") | |
| analyze_btn = gr.Button("⚡ Analyze Document", variant="primary", size="lg") | |
| gr.Markdown( | |
| """ | |
| **Supported Document Types:** | |
| - 🧾 Invoices & Receipts (Tables, Line Items, Totals) | |
| - 🏛 Architecture & System Diagrams (Nodes, Arrows, Tech) | |
| - 📄 Scanned Documents & Financial Forms | |
| """ | |
| ) | |
| # RIGHT SIDE: Structured Results | |
| with gr.Column(scale=2): | |
| doc_type_output = gr.Markdown("### 📄 Document Type: *Pending Analysis*") | |
| summary_output = gr.Textbox(label="Summary", interactive=False, lines=2) | |
| with gr.Tabs(): | |
| with gr.TabItem("📊 Entities & Fields"): | |
| entities_output = gr.Dataframe(label="Key-Value Entities", interactive=False) | |
| with gr.TabItem("📑 Tables"): | |
| tables_output = gr.Dataframe(label="Detected Tables", interactive=False) | |
| with gr.TabItem("🏗️ Architecture Diagram"): | |
| comp_output = gr.Dataframe(label="Components", interactive=False) | |
| rel_output = gr.Dataframe(label="Relationships & Links", interactive=False) | |
| with gr.TabItem("📝 OCR Extracted Text"): | |
| ocr_text_output = gr.Textbox(label="Raw OCR Extracted Content", interactive=False, lines=8) | |
| with gr.TabItem("🔍 Confidence & Notes"): | |
| notes_output = gr.Markdown("Confidence and uncertainty observations") | |
| with gr.TabItem("📦 Structured JSON"): | |
| json_output = gr.JSON(label="Full Pydantic Result") | |
| analyze_btn.click( | |
| fn=process_ui_image, | |
| inputs=[image_input], | |
| outputs=[ | |
| doc_type_output, | |
| summary_output, | |
| entities_output, | |
| tables_output, | |
| comp_output, | |
| rel_output, | |
| ocr_text_output, | |
| notes_output, | |
| json_output, | |
| ] | |
| ) | |
| return demo | |
| # Create Gradio UI instance | |
| demo = create_gradio_ui() | |
| # In Hugging Face Spaces with Gradio SDK, expose demo directly | |
| # For local / Docker / Cloud Run, mount on FastAPI app | |
| if not os.getenv("SPACE_ID"): | |
| app = gr.mount_gradio_app(app, demo, path="/") | |
| if __name__ == "__main__": | |
| port = int(os.getenv("PORT", "7860" if os.getenv("SPACE_ID") else str(settings.PORT))) | |
| logger.info("Starting server on port %d...", port) | |
| if os.getenv("SPACE_ID"): | |
| demo.launch(server_name="0.0.0.0", server_port=port) | |
| else: | |
| import uvicorn | |
| uvicorn.run("app:app", host="0.0.0.0", port=port) | |