app.py

"""
Visiting Card & Letterhead OCR API
===================================
Two-step pipeline: nemoretriever-ocr-v1 → nvidia-nemotron-nano-9b-v2
Deploy on Hugging Face Spaces (Docker or Python SDK):
  - Set secret  NVIDIA_API_KEY  in Space settings → Variables and secrets
  - The app serves the HTML frontend at  /  and the API at  /extract-card
  - HF Spaces exposes port 7860 by default (set via HF_PORT env var)
Local usage:
  pip install fastapi uvicorn requests python-multipart pdf2image
  NVIDIA_API_KEY=nvapi-xxx python visiting_card_api.py
  Open http://localhost:7860

PDF Support:
  - Upload a PDF (single or multi-page) to /extract-card just like an image.
  - Each page is converted to a PNG in memory, OCR'd, and the combined
    text is passed to the LLM — zero changes needed on the caller side.
  - Requires poppler to be installed on the server:
      Ubuntu/Debian : sudo apt-get install -y poppler-utils
      macOS         : brew install poppler
      Windows       : choco install poppler  (or download binaries)
"""

import io
import os
import re
import json
import base64
import requests
from pathlib import Path
from typing import List

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
from pydantic import BaseModel

# pdf2image converts PDF pages → PIL images (requires poppler on PATH)
import fitz  # pymupdf

# ── App ────────────────────────────────────────────────────────────────────────
app = FastAPI(
    title="Visiting Card & Letterhead OCR API",
    description="Two-step RAG pipeline: nemoretriever-ocr-v1 → nvidia-nemotron-nano-9b-v2",
)

# ── CORS — allow all origins (needed for HF Spaces iframe / custom domains) ───
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── Configuration ─────────────────────────────────────────────────────────────
NVIDIA_API_KEY = os.environ.get("NVIDIA_API_KEY", "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")

OCR_URL   = "https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-ocr-v1"
LLM_URL   = "https://integrate.api.nvidia.com/v1/chat/completions"
LLM_MODEL = "nvidia/nvidia-nemotron-nano-9b-v2"

OCR_HEADERS = {"Authorization": f"Bearer {NVIDIA_API_KEY}", "Accept": "application/json"}
LLM_HEADERS = {"Authorization": f"Bearer {NVIDIA_API_KEY}", "Content-Type": "application/json"}

# Maximum PDF pages to process (guards against large documents being sent by mistake)
PDF_MAX_PAGES = 10

# ── System prompt ──────────────────────────────────────────────────────────────
CARD_SYSTEM_PROMPT = """You are a business card and letterhead data extraction assistant.
You will receive raw OCR text extracted from a visiting card, business card, or the header/footer of a business letter.
Parse it carefully and return ONLY a valid JSON object.
No markdown fences, no explanation, no preamble — just the raw JSON object.
JSON schema (return exactly this structure):
{
  "company_name": "full name of the company or firm (string)",
  "contact_person": "name of the individual on the card or letter (string)",
  "designation": "job title or designation of the contact person (string)",
  "mobile": "mobile number(s) as a string; if multiple separate with comma (string)",
  "phone": "landline / office phone number(s); if multiple separate with comma (string)",
  "email": "email address(es); if multiple separate with comma (string)",
  "address": "full postal address as printed, preserving line breaks with a pipe | separator (string)",
  "pin": "PIN code / ZIP code / postal code as a string of digits (string)",
  "city": "city name (string)",
  "state": "state or province name (string)",
  "country": "country name (string)",
  "gst_number": "GST / GSTIN number; typically 15 alphanumeric characters (string)",
  "website": "website URL if present (string)",
  "fax": "fax number if present (string)"
}
Rules:
- company_name: usually the largest text or the text near a logo
- contact_person: individual's personal name distinct from company name
- designation: title like CEO, Manager, Director, Proprietor, Sales Executive, etc.
- mobile: numbers prefixed with M:, Mob:, Cell:, +91, or 10-digit numbers
- phone: numbers prefixed with Ph:, Tel:, T:, O:, or STD codes like (022), (080)
- email: look for @ symbol; may be prefixed with E:, Email:, Mail:
- address: collect all address lines; separate each line with ' | '
- pin: extract 6-digit Indian PIN code or 5/9-digit ZIP; digits only
- city: extract city name from address
- state: extract state name from address
- country: default to India if address looks Indian and country not stated
- gst_number: 15-character alphanumeric GSTIN
- website: any URL starting with www., http://, or https://
- fax: number prefixed with Fax:, F:, or similar
- If a field is not found return "" (empty string)
- Do NOT invent or hallucinate any information not present in the OCR text
- If multiple phone or mobile numbers are present, join them with ', '"""


# ── PDF helpers ────────────────────────────────────────────────────────────────

def _is_pdf_content_type(content_type: str | None) -> bool:
    """Return True if the MIME type indicates a PDF."""
    if not content_type:
        return False
    ct = content_type.lower()
    return ct in ("application/pdf", "application/x-pdf", "binary/octet-stream")


def _pdf_bytes_to_png_b64_list(pdf_bytes: bytes) -> list[str]:
    try:
        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    except Exception as exc:
        raise HTTPException(422, f"Could not read PDF: {exc}")

    if doc.page_count == 0:
        raise HTTPException(422, "PDF has no pages or could not be rendered.")

    if doc.page_count > PDF_MAX_PAGES:
        raise HTTPException(
            400,
            f"PDF has {doc.page_count} pages; maximum allowed is {PDF_MAX_PAGES}. "
            "Please send a single-page or trimmed PDF.",
        )

    b64_list: list[str] = []
    for page in doc:
        pix = page.get_pixmap(dpi=200)
        b64_list.append(base64.b64encode(pix.tobytes("png")).decode())
    doc.close()
    return b64_list


# ── Core OCR helper (image bytes → base64 → NVIDIA OCR API) ──────────────────

def _ocr_single_b64(image_b64: str) -> str:
    """
    Call the NVIDIA OCR endpoint for one base64-encoded PNG/JPEG image.
    Returns the extracted text lines joined by newlines.
    """
    if len(image_b64) >= 180_000:
        raise HTTPException(
            413,
            "Image too large (base64 must be < 180,000 chars). "
            "Reduce DPI or crop the image and retry.",
        )

    payload = {
        "input": [
            {
                "type": "image_url",
                "url": f"data:image/png;base64,{image_b64}",
            }
        ]
    }

    try:
        r = requests.post(OCR_URL, headers=OCR_HEADERS, json=payload, timeout=30)
        r.raise_for_status()
    except requests.exceptions.RequestException as e:
        raise HTTPException(502, f"NVIDIA OCR API error: {e}")

    ocr_json = r.json()
    detections = ocr_json.get("text_detections", [])
    if not detections:
        data = ocr_json.get("data", [])
        if isinstance(data, list) and data:
            detections = data[0].get("text_detections", [])

    lines: list[str] = []
    for det in detections:
        text = ""
        if isinstance(det, dict):
            if "text_prediction" in det:
                text = det["text_prediction"].get("text", "").strip()
            else:
                text = det.get("text", "").strip()
        if text:
            lines.append(text)
    return "\n".join(lines)


# ── Public run_ocr — handles BOTH images and PDFs transparently ───────────────

async def run_ocr(file: UploadFile) -> str:
    """
    Read the uploaded file, detect whether it is a PDF or an image,
    convert PDF pages to PNG images if needed, and return combined OCR text.

    The caller (extract_card / extract_card_batch) does NOT need to know
    whether the upload was a PDF or a direct image — the output is always
    a plain text string.
    """
    content: bytes = await file.read()

    # ── Detect PDF by magic bytes (%PDF) or MIME type ─────────────────────────
    is_pdf = content[:4] == b"%PDF" or _is_pdf_content_type(file.content_type)

    if is_pdf:
        # Convert every PDF page to PNG, OCR each, then join all page texts
        b64_pages = _pdf_bytes_to_png_b64_list(content)
        page_texts: list[str] = []
        for page_idx, b64 in enumerate(b64_pages, start=1):
            page_text = _ocr_single_b64(b64)
            if page_text.strip():
                page_texts.append(f"[Page {page_idx}]\n{page_text}")
        return "\n\n".join(page_texts)

    # ── Regular image path (unchanged behaviour) ──────────────────────────────
    image_b64 = base64.b64encode(content).decode()
    return _ocr_single_b64(image_b64)


# ── LLM call ──────────────────────────────────────────────────────────────────

def call_llm(ocr_text: str) -> dict:
    payload = {
        "model": LLM_MODEL,
        "max_tokens": 2048,
        "temperature": 0.1,
        "top_p": 0.9,
        "messages": [
            {"role": "system", "content": CARD_SYSTEM_PROMPT},
            {"role": "user", "content": (
                f"Here is the OCR text extracted from the business card or letterhead:\n\n"
                f"{ocr_text}\n\nExtract the required data and return ONLY the JSON object."
            )},
        ],
    }

    try:
        r = requests.post(LLM_URL, headers=LLM_HEADERS, json=payload, timeout=120)
        r.raise_for_status()
        llm_json = r.json()
    except requests.exceptions.RequestException as e:
        raise HTTPException(502, f"NVIDIA LLM API error: {e}")

    raw: str = llm_json.get("choices", [{}])[0].get("message", {}).get("content", "")
    if not raw:
        raise HTTPException(502, "LLM returned empty response")

    cleaned = re.sub(r"```json\s*", "", raw, flags=re.IGNORECASE)
    cleaned = re.sub(r"```\s*", "", cleaned).strip()

    try:
        parsed = json.loads(cleaned)
    except json.JSONDecodeError:
        m = re.search(r"\{[\s\S]*\}", cleaned)
        if not m:
            raise HTTPException(502, f"LLM did not return valid JSON. Preview: {raw[:400]}")
        try:
            parsed = json.loads(m.group(0))
        except json.JSONDecodeError as e:
            raise HTTPException(502, f"JSON parse error: {e}")

    if not isinstance(parsed, dict):
        raise HTTPException(502, f"LLM response not a JSON object. Got: {type(parsed).__name__}")
    return parsed


# ── Pydantic models ────────────────────────────────────────────────────────────

class CardData(BaseModel):
    company_name: str
    contact_person: str
    designation: str
    mobile: str
    phone: str
    email: str
    address: str
    pin: str
    city: str
    state: str
    country: str
    gst_number: str
    website: str
    fax: str


def build_card(parsed: dict) -> CardData:
    def s(k, n=300): return str(parsed.get(k, "")).strip()[:n]
    return CardData(
        company_name=s("company_name", 200), contact_person=s("contact_person", 100),
        designation=s("designation", 100),   mobile=s("mobile", 100),
        phone=s("phone", 100),               email=s("email", 200),
        address=s("address", 500),           pin=s("pin", 10),
        city=s("city", 100),                 state=s("state", 100),
        country=s("country", 100),           gst_number=s("gst_number", 20),
        website=s("website", 200),           fax=s("fax", 50),
    )


# ── API endpoints ──────────────────────────────────────────────────────────────

ALLOWED_IMAGE_TYPES = {"image/jpeg", "image/jpg", "image/png", "image/webp"}
ALLOWED_PDF_TYPES   = {"application/pdf", "application/x-pdf"}

def _validate_file_type(file: UploadFile, idx: int | None = None) -> None:
    """
    Allow images OR PDFs.  PDFs are also detected by magic bytes in run_ocr,
    so this is a friendly early rejection for clearly wrong MIME types.
    """
    prefix = f"File {idx}: " if idx is not None else ""
    ct = (file.content_type or "").lower()
    if ct and ct not in ALLOWED_IMAGE_TYPES | ALLOWED_PDF_TYPES:
        raise HTTPException(
            415,
            f"{prefix}Unsupported file type '{ct}'. "
            "Accepted: JPEG, PNG, WebP images or PDF documents.",
        )


@app.post("/extract-card", response_model=CardData)
async def extract_card(file: UploadFile = File(...)):
    _validate_file_type(file)
    ocr_text = await run_ocr(file)
    if not ocr_text.strip():
        raise HTTPException(422, "OCR produced no text. Check image/PDF quality.")
    return build_card(call_llm(ocr_text))


@app.post("/extract-card/batch", response_model=List[CardData])
async def extract_card_batch(files: List[UploadFile] = File(...)):
    if len(files) > 10:
        raise HTTPException(400, "Maximum 10 files per batch request.")
    empty = CardData(**{f: "" for f in CardData.model_fields})
    results: List[CardData] = []
    for idx, file in enumerate(files, start=1):
        _validate_file_type(file, idx)
        ocr_text = await run_ocr(file)
        results.append(build_card(call_llm(ocr_text)) if ocr_text.strip() else empty)
    return results


@app.get("/health")
async def health():
    return {"status": "healthy", "model": LLM_MODEL}


# ── Serve index.html at root (must be placed alongside this script) ────────────
HTML_PATH = Path(__file__).parent / "index.html"

@app.get("/", response_class=HTMLResponse)
async def serve_ui():
    if not HTML_PATH.exists():
        return HTMLResponse(
            "<h2 style='font-family:sans-serif;padding:40px'>"
            "index.html not found — place it next to visiting_card_api.py</h2>", 500
        )
    return HTMLResponse(HTML_PATH.read_text(encoding="utf-8"))


# ── Entry point ────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("HF_PORT", 7860))
    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=False)