app.py

"""
KYC Document Validator — FastAPI service (PAN + Aadhaar).

ONE model, loaded once, serves both document types. Runs fully locally:
images never leave this server.

Run:
    pip install fastapi "uvicorn[standard]" python-multipart transformers \
                accelerate qwen-vl-utils ultralytics opencv-python-headless pillow torch
    uvicorn app:app --host 0.0.0.0 --port 8000 --workers 1
"""
import os
os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
# os.environ.setdefault("HF_HUB_OFFLINE", "1")   # enable AFTER first download

import re
import json
import threading
from contextlib import asynccontextmanager
from typing import Optional

import numpy as np
import cv2
import torch
from PIL import Image
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from ultralytics import YOLO
from ultralytics.nn.tasks import DetectionModel
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info

MODEL_ID = "Qwen/Qwen2-VL-2B-Instruct"
_gpu_lock = threading.Lock()
_state = {}

# ---------------------------------------------------------------------
# PAN
# ---------------------------------------------------------------------
PAN_ENTITY_MAP = {
    'P': 'Person (Individual)', 'C': 'Company', 'F': 'Firm / Limited Liability Partnership (LLP)',
    'H': 'Hindu Undivided Family (HUF)', 'T': 'Trust', 'A': 'Association of Persons (AOP)',
    'B': 'Body of Individuals (BOI)', 'G': 'Government Agency', 'L': 'Local Authority',
    'J': 'Artificial Juridical Person'
}
PAN_REGEX = re.compile(r'^[A-Z]{5}[0-9]{4}[A-Z]$')
to_letter = str.maketrans({'0': 'O', '1': 'I', '2': 'Z', '5': 'S', '8': 'B'})
to_number = str.maketrans({'O': '0', 'I': '1', 'Z': '2', 'S': '5', 'B': '8', 'G': '6', 'Q': '0'})

def validate_pan(s):
    s = re.sub(r'[^A-Z0-9]', '', (s or '').upper())
    if len(s) != 10:
        return None
    if PAN_REGEX.match(s):
        return s
    fixed = (s[0:5].translate(to_letter)
             + s[5:9].translate(to_number)
             + s[9:10].translate(to_letter))
    return fixed if PAN_REGEX.match(fixed) else None

# ---------------------------------------------------------------------
# AADHAAR — validated with the Verhoeff checksum (last digit is a check digit)
# ---------------------------------------------------------------------
_VERHOEFF_D = [
    [0,1,2,3,4,5,6,7,8,9],[1,2,3,4,0,6,7,8,9,5],[2,3,4,0,1,7,8,9,5,6],
    [3,4,0,1,2,8,9,5,6,7],[4,0,1,2,3,9,5,6,7,8],[5,9,8,7,6,0,4,3,2,1],
    [6,5,9,8,7,1,0,4,3,2],[7,6,5,9,8,2,1,0,4,3],[8,7,6,5,9,3,2,1,0,4],
    [9,8,7,6,5,4,3,2,1,0],
]
_VERHOEFF_P = [
    [0,1,2,3,4,5,6,7,8,9],[1,5,7,6,2,8,3,0,9,4],[5,8,0,3,7,9,6,1,4,2],
    [8,9,1,6,0,4,3,5,2,7],[9,4,5,3,1,2,6,8,7,0],[4,2,8,6,5,7,3,9,0,1],
    [2,7,9,3,8,0,6,4,1,5],[7,0,4,6,9,1,3,2,5,8],
]

def _verhoeff_ok(num_str):
    c = 0
    for i, ch in enumerate(reversed(num_str)):
        c = _VERHOEFF_D[c][_VERHOEFF_P[i % 8][int(ch)]]
    return c == 0

def validate_aadhaar(s):
    s = re.sub(r'\D', '', s or '')
    if len(s) != 12:
        return None
    if s[0] in '01':          # Aadhaar never starts with 0 or 1
        return None
    return s if _verhoeff_ok(s) else None

def mask_aadhaar(s):
    return "XXXX XXXX " + s[-4:]

# ---------------------------------------------------------------------
# DOC CONFIGS (prompt + validator per type)
# ---------------------------------------------------------------------
PAN_PROMPT = (
    "Look at this image. If it is an Indian PAN card, extract its fields. "
    "Return ONLY a JSON object, no other text:\n"
    '{"is_pan_card": true or false, "pan_number": "", "name": "", '
    '"fathers_name": "", "date_of_birth": ""}\n'
    "The PAN number is exactly 10 characters: 5 letters, 4 digits, 1 letter "
    "(format ABCDE1234F). Copy it exactly. If not a PAN card, set is_pan_card to false."
)
AADHAAR_PROMPT = (
    "Look at this image. If it is an Indian Aadhaar card (UIDAI), extract its fields. "
    "Return ONLY a JSON object, no other text:\n"
    '{"is_aadhaar_card": true or false, "aadhaar_number": "", "name": "", '
    '"date_of_birth": "", "gender": "", "address": ""}\n'
    "The Aadhaar number is exactly 12 digits, usually in three groups of four "
    "(example 1234 5678 9012). Copy the digits exactly. "
    "If not an Aadhaar card, set is_aadhaar_card to false."
)
DOC_CONFIGS = {
    "pan":     {"prompt": PAN_PROMPT,     "type_key": "is_pan_card",     "num_field": "pan_number",     "validate": validate_pan},
    "aadhaar": {"prompt": AADHAAR_PROMPT, "type_key": "is_aadhaar_card", "num_field": "aadhaar_number", "validate": validate_aadhaar},
}

SPOOF_CLASSES = [62, 63, 67]   # COCO: tv(monitor), laptop, cell phone
SPOOF_CONFIDENCE = 0.35

# ---------------------------------------------------------------------
# MODEL LOADING (once, at startup)
# ---------------------------------------------------------------------
def _load_models():
    torch.serialization.add_safe_globals([DetectionModel])
    _orig = torch.load
    torch.load = lambda *a, **k: _orig(*a, **{**k, 'weights_only': False})
    spoof = YOLO('yolov8n.pt')
    torch.load = _orig
    use_cuda = torch.cuda.is_available()
    dtype = torch.bfloat16 if use_cuda else torch.float32
    vlm = Qwen2VLForConditionalGeneration.from_pretrained(
        MODEL_ID, torch_dtype=dtype,
        device_map="auto" if use_cuda else None)
    if not use_cuda:
        vlm = vlm.to("cpu")
    proc = AutoProcessor.from_pretrained(MODEL_ID)
    _state.update(spoof=spoof, vlm=vlm, proc=proc)

def _warmup():
    try:
        _run_vlm(Image.new("RGB", (64, 64), (255, 255, 255)), PAN_PROMPT)
    except Exception:
        pass

@asynccontextmanager
async def lifespan(app: FastAPI):
    print("Loading models (one time)...")
    _load_models()
    _warmup()
    print("Ready. Serving PAN + Aadhaar against one resident model.")
    yield
    _state.clear()

app = FastAPI(title="KYC Document Validator", lifespan=lifespan)

# For testing, allow any origin. In production, replace ["*"] with your
# frontend's exact origin, e.g. ["https://kyc.yourcompany.com"].
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["POST", "GET"],
    allow_headers=["*"],
)

# ---------------------------------------------------------------------
# PIPELINE
# ---------------------------------------------------------------------
def _detect_spoof(bgr):
    res = _state["spoof"].predict(bgr, verbose=False)
    for box in res[0].boxes:
        cid, conf = int(box.cls[0].item()), box.conf[0].item()
        if cid in SPOOF_CLASSES and conf > SPOOF_CONFIDENCE:
            return _state["spoof"].names[cid], float(conf)
    return None, 0.0

def _run_vlm(pil_img, prompt):
    proc, vlm = _state["proc"], _state["vlm"]
    messages = [{"role": "user", "content": [
        {"type": "image", "image": pil_img},
        {"type": "text", "text": prompt},
    ]}]
    text = proc.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = proc(text=[text], images=image_inputs, videos=video_inputs,
                  padding=True, return_tensors="pt").to(vlm.device)
    out = vlm.generate(**inputs, max_new_tokens=256, do_sample=False)
    trimmed = [o[len(i):] for i, o in zip(inputs.input_ids, out)]
    raw = proc.batch_decode(trimmed, skip_special_tokens=True)[0]
    raw = re.sub(r'```json|```', '', raw).strip()
    m = re.search(r'\{.*\}', raw, re.DOTALL)
    try:
        return json.loads(m.group()) if m else {}
    except json.JSONDecodeError:
        return {}

class KycResult(BaseModel):
    document_type: str
    accepted: bool
    status: str
    reason: Optional[str] = None
    number: Optional[str] = None          # PAN: full. Aadhaar: only if return_full_number=True
    number_masked: Optional[str] = None   # Aadhaar masked (XXXX XXXX 1234)
    name: Optional[str] = None
    date_of_birth: Optional[str] = None
    fathers_name: Optional[str] = None    # PAN only
    gender: Optional[str] = None          # Aadhaar only
    classification: Optional[str] = None  # PAN entity type
    routing: Optional[str] = None

def _process(bgr, doc_type, return_full_number=False) -> KycResult:
    cfg = DOC_CONFIGS[doc_type]
    with _gpu_lock:
        device, _ = _detect_spoof(bgr)
        if device:
            return KycResult(document_type=doc_type, accepted=False,
                status="REJECTED_GATE1_SPOOF",
                reason=f"Detected '{device}' — looks like a photo of a screen.")
        pil = Image.fromarray(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))
        data = _run_vlm(pil, cfg["prompt"])

    # is this actually the document type the frontend expected at this step?
    if not data.get(cfg["type_key"]):
        return KycResult(document_type=doc_type, accepted=False,
            status="REJECTED_TYPE_MISMATCH",
            reason=f"Expected a {doc_type.upper()} card, but the image isn't one.")

    number = cfg["validate"](data.get(cfg["num_field"]))
    if not number:
        return KycResult(document_type=doc_type, accepted=False,
            status="REJECTED_DATA_CHECK",
            reason=f"No valid {doc_type.upper()} number could be read.")

    if doc_type == "pan":
        ec = number[3]
        return KycResult(document_type="pan", accepted=True, status="ACCEPTED",
            number=number, name=data.get("name"), fathers_name=data.get("fathers_name"),
            date_of_birth=data.get("date_of_birth"),
            classification=PAN_ENTITY_MAP.get(ec, "Unknown Entity Classification"),
            routing="PERSONAL ROUTE" if ec == 'P' else "BUSINESS/ENTITY ROUTE")
    else:  # aadhaar — checksum-verified; default to MASKED for privacy
        return KycResult(document_type="aadhaar", accepted=True, status="ACCEPTED",
            number=number if return_full_number else None,
            number_masked=mask_aadhaar(number),
            name=data.get("name"), date_of_birth=data.get("date_of_birth"),
            gender=data.get("gender"))

# ---------------------------------------------------------------------
# ENDPOINTS
# ---------------------------------------------------------------------
@app.get("/health")
def health():
    return {"status": "ok", "model_loaded": "vlm" in _state}

@app.post("/validate-document", response_model=KycResult)
def validate_document(doc_type: str = Form(...), file: UploadFile = File(...),
                      return_full_number: bool = Form(False)):
    doc_type = (doc_type or "").lower().strip()
    if doc_type not in DOC_CONFIGS:
        raise HTTPException(400, "doc_type must be 'pan' or 'aadhaar'")
    img = cv2.imdecode(np.frombuffer(file.file.read(), np.uint8), cv2.IMREAD_COLOR)
    if img is None:
        raise HTTPException(400, "Could not decode image.")
    return _process(img, doc_type, return_full_number)

# convenience aliases so the frontend can call the obvious one per step
@app.post("/validate-pan", response_model=KycResult)
def validate_pan_endpoint(file: UploadFile = File(...)):
    img = cv2.imdecode(np.frombuffer(file.file.read(), np.uint8), cv2.IMREAD_COLOR)
    if img is None:
        raise HTTPException(400, "Could not decode image.")
    return _process(img, "pan")

@app.post("/validate-aadhaar", response_model=KycResult)
def validate_aadhaar_endpoint(file: UploadFile = File(...),
                              return_full_number: bool = Form(False)):
    img = cv2.imdecode(np.frombuffer(file.file.read(), np.uint8), cv2.IMREAD_COLOR)
    if img is None:
        raise HTTPException(400, "Could not decode image.")
    return _process(img, "aadhaar", return_full_number)