Spaces:
Sleeping
Sleeping
| import os | |
| import time | |
| import csv | |
| import json | |
| from datetime import datetime | |
| from typing import Dict, Any, Optional | |
| import pickle | |
| from pathlib import Path | |
| from typing import List | |
| from fastapi import FastAPI, File, UploadFile, HTTPException | |
| from fastapi.responses import FileResponse, JSONResponse | |
| from fastapi.staticfiles import StaticFiles | |
| import uvicorn | |
| app = FastAPI( | |
| title="Document Intelligence & Extraction API", | |
| description="Two-stage pipeline: Document Classification & Field Extraction with Confidence Routing", | |
| version="1.0.0" | |
| ) | |
| # Global pipeline instance | |
| PIPELINE = None | |
| PIPELINE_ERROR = None | |
| LOG_CSV = "service_requests.csv" | |
| PROJECT_ROOT = Path(__file__).resolve().parents[2] | |
| REPORTS_DIR = PROJECT_ROOT / "reports" | |
| FRONTEND_DIR = PROJECT_ROOT.parent / "frontend" / "dist" | |
| SAMPLE_DOCUMENTS = { | |
| "invoice-ffcy0002": { | |
| "label": "Commercial Invoice", | |
| "description": "Invoice with supplier, buyer, line items, and total amount.", | |
| "path": PROJECT_ROOT / "data/raw/invoice/ffcy0002.png", | |
| "doc_class": "invoice", | |
| }, | |
| "invoice-fflb0108": { | |
| "label": "Purchase Invoice", | |
| "description": "Invoice document for layout and field extraction testing.", | |
| "path": PROJECT_ROOT / "data/raw/invoice/fflb0108.png", | |
| "doc_class": "invoice", | |
| }, | |
| "receipt-cord-000": { | |
| "label": "Retail Receipt", | |
| "description": "Receipt with merchant details, date, and total balance.", | |
| "path": PROJECT_ROOT / "data/raw/receipt/cord_test-000_0.png", | |
| "doc_class": "receipt", | |
| }, | |
| "receipt-cord-001": { | |
| "label": "Receipt Variant", | |
| "description": "Receipt variation for OCR and classification comparison.", | |
| "path": PROJECT_ROOT / "data/raw/receipt/cord_test-000_1.png", | |
| "doc_class": "receipt", | |
| }, | |
| "application-funsd-1160": { | |
| "label": "Corporate Application Form", | |
| "description": "Structured application form with applicant fields and signature block.", | |
| "path": PROJECT_ROOT / "data/raw/application_form/funsd_0000971160.png", | |
| "doc_class": "application_form", | |
| }, | |
| "application-funsd-9556": { | |
| "label": "Official Evaluation Form", | |
| "description": "Form with headers, check boxes, and applicant fields.", | |
| "path": PROJECT_ROOT / "data/raw/application_form/funsd_0000989556.png", | |
| "doc_class": "application_form", | |
| }, | |
| } | |
| def log_api_request(doc_id: str, doc_class: str, confidence: float, needs_review: bool, latency_ms: float): | |
| file_exists = os.path.exists(LOG_CSV) | |
| with open(LOG_CSV, "a", newline="", encoding="utf-8") as f: | |
| writer = csv.writer(f) | |
| if not file_exists: | |
| writer.writerow(["timestamp", "doc_id", "doc_class", "confidence", "needs_review", "latency_ms"]) | |
| writer.writerow([datetime.utcnow().isoformat(), doc_id, doc_class, confidence, needs_review, latency_ms]) | |
| def load_champion_pipeline(): | |
| global PIPELINE, PIPELINE_ERROR | |
| if PIPELINE is not None: | |
| return PIPELINE | |
| started_at = time.time() | |
| try: | |
| from ..pipeline.end_to_end import DocumentIntelligencePipeline | |
| from ..pipeline.confidence_router import ConfidenceRouter | |
| from ..extraction.cohere_extractor import CohereFieldExtractor | |
| # Check for champion models on disk or initialize default | |
| models_dir = REPORTS_DIR / "classification_artifacts" | |
| # Logistic regression is the stable low-latency inference artifact for the | |
| # web app. XGBoost remains available in analytics but can block first-use | |
| # loading on small CPU Space instances. | |
| cls_path = models_dir / "logistic_regression.pkl" | |
| if not cls_path.exists(): | |
| cls_path = models_dir / "xgboost.pkl" | |
| if cls_path.exists(): | |
| print(f"[API] Loading classifier: {cls_path.name}", flush=True) | |
| with open(cls_path, "rb") as f: | |
| classifier = pickle.load(f) | |
| else: | |
| from ..classification.train_classical import DocumentClassificationPipeline | |
| classifier = DocumentClassificationPipeline("logistic_regression") | |
| extractors = { | |
| "invoice": CohereFieldExtractor("invoice"), | |
| "receipt": CohereFieldExtractor("receipt"), | |
| "application_form": CohereFieldExtractor("application_form") | |
| } | |
| PIPELINE = DocumentIntelligencePipeline( | |
| classifier=classifier, | |
| extractors_by_class=extractors, | |
| router=ConfidenceRouter(t1_class_confidence=0.75, t2_field_confidence=0.70) | |
| ) | |
| PIPELINE_ERROR = None | |
| print(f"[API] Pipeline initialized in {round(time.time() - started_at, 2)}s", flush=True) | |
| return PIPELINE | |
| except Exception as exc: | |
| PIPELINE_ERROR = f"{type(exc).__name__}: {exc}" | |
| print(f"[API] Pipeline initialization failed: {PIPELINE_ERROR}", flush=True) | |
| raise | |
| def health_check(): | |
| return {"status": "ok", "service": "document_intelligence_api", "timestamp": datetime.utcnow().isoformat()} | |
| def pipeline_status(): | |
| return { | |
| "status": "ready" if PIPELINE is not None else "not_initialized", | |
| "pipeline_initialized": PIPELINE is not None, | |
| "pipeline_error": PIPELINE_ERROR, | |
| } | |
| def list_samples(): | |
| return [ | |
| {"id": key, "label": value["label"], "description": value["description"], "doc_class": value["doc_class"]} | |
| for key, value in SAMPLE_DOCUMENTS.items() | |
| if value["path"].exists() | |
| ] | |
| def sample_preview(sample_id: str): | |
| sample = SAMPLE_DOCUMENTS.get(sample_id) | |
| if not sample or not sample["path"].exists(): | |
| raise HTTPException(status_code=404, detail="Sample document not found") | |
| return FileResponse(sample["path"]) | |
| def process_path(path: str, doc_id: str): | |
| try: | |
| pipeline = load_champion_pipeline() | |
| except Exception as exc: | |
| raise HTTPException(status_code=503, detail=f"Pipeline initialization failed: {exc}") from exc | |
| if pipeline is None: | |
| raise HTTPException(status_code=503, detail="Pipeline not initialized") | |
| res = pipeline.process_document(path, doc_id=doc_id) | |
| log_api_request( | |
| doc_id=doc_id, | |
| doc_class=res["doc_class"], | |
| confidence=res["class_confidence"], | |
| needs_review=res["needs_review"], | |
| latency_ms=res["processing_time_ms"] | |
| ) | |
| return res | |
| async def process_document_endpoint(file: UploadFile = File(...)): | |
| temp_path = f"temp_{int(time.time()*1000)}_{file.filename}" | |
| try: | |
| content = await file.read() | |
| if len(content) > 10 * 1024 * 1024: | |
| raise HTTPException(status_code=413, detail="File size exceeds maximum allowed upload limit of 10 MB") | |
| with open(temp_path, "wb") as f_out: | |
| f_out.write(content) | |
| return JSONResponse(content=process_path(temp_path, file.filename)) | |
| except HTTPException: | |
| raise | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| finally: | |
| if os.path.exists(temp_path): | |
| try: | |
| os.remove(temp_path) | |
| except Exception: | |
| pass | |
| def process_sample(sample_id: str): | |
| sample = SAMPLE_DOCUMENTS.get(sample_id) | |
| if not sample or not sample["path"].exists(): | |
| raise HTTPException(status_code=404, detail="Sample document not found") | |
| try: | |
| return JSONResponse(content=process_path(str(sample["path"]), sample_id)) | |
| except HTTPException: | |
| raise | |
| except Exception as exc: | |
| raise HTTPException(status_code=500, detail=str(exc)) from exc | |
| async def process_batch(files: List[UploadFile] = File(...)): | |
| if len(files) > 20: | |
| raise HTTPException(status_code=400, detail="Batch size cannot exceed 20 files") | |
| results = [] | |
| temp_paths = [] | |
| try: | |
| for file in files: | |
| content = await file.read() | |
| if len(content) > 10 * 1024 * 1024: | |
| results.append({"file_name": file.filename, "error": "File exceeds 10 MB"}) | |
| continue | |
| temp_path = PROJECT_ROOT / "data/processed/uploads" / f"batch_{int(time.time()*1000)}_{file.filename}" | |
| temp_path.parent.mkdir(parents=True, exist_ok=True) | |
| temp_path.write_bytes(content) | |
| temp_paths.append(temp_path) | |
| try: | |
| result = process_path(str(temp_path), file.filename) | |
| results.append({ | |
| "file_name": file.filename, | |
| "doc_class": result["doc_class"], | |
| "confidence": result["class_confidence"], | |
| "needs_review": result["needs_review"], | |
| "latency_ms": result["processing_time_ms"], | |
| "words": result["ocr_token_count"], | |
| }) | |
| except Exception as exc: | |
| results.append({"file_name": file.filename, "error": str(exc)}) | |
| return results | |
| finally: | |
| for path in temp_paths: | |
| path.unlink(missing_ok=True) | |
| def analytics(model: str = "xgboost"): | |
| artifact_dir = REPORTS_DIR / "classification_artifacts" | |
| matrix = artifact_dir / f"normalized_confusion_matrix_{model}.png" | |
| if not matrix.exists(): | |
| matrix = artifact_dir / f"confusion_matrix_{model}.png" | |
| metrics_path = artifact_dir / f"classification_report_{model}.json" | |
| metrics = {} | |
| if metrics_path.exists(): | |
| with metrics_path.open("r", encoding="utf-8") as handle: | |
| metrics = json.load(handle) | |
| return {"model": model, "metrics": metrics, "matrix_url": f"/api/analytics/{model}/matrix" if matrix.exists() else None} | |
| def analytics_matrix(model: str): | |
| artifact_dir = REPORTS_DIR / "classification_artifacts" | |
| matrix = artifact_dir / f"normalized_confusion_matrix_{model}.png" | |
| if not matrix.exists(): | |
| matrix = artifact_dir / f"confusion_matrix_{model}.png" | |
| if not matrix.exists(): | |
| raise HTTPException(status_code=404, detail="Matrix artifact not found") | |
| return FileResponse(matrix) | |
| if FRONTEND_DIR.exists(): | |
| app.mount("/assets", StaticFiles(directory=FRONTEND_DIR / "assets"), name="assets") | |
| def frontend(full_path: str): | |
| requested = FRONTEND_DIR / full_path | |
| if full_path and requested.is_file(): | |
| return FileResponse(requested) | |
| return FileResponse(FRONTEND_DIR / "index.html") | |
| if __name__ == "__main__": | |
| uvicorn.run("src.api.main:app", host="127.0.0.1", port=8000, reload=True) | |