kyc_extract.py

"""
Generic KYC OCR pipeline.

Handles:
  - PDF or image input (PyMuPDF for PDF rendering)
  - Auto-rotation detection (cheques are often shot sideways)
  - Document type auto-classification (PAN / Cheque / Aadhaar / Passport)
  - Per-document field extractors with regex + heuristics + validation

Nothing is hardcoded to specific banks, names, or documents.
All extraction is rule-based and works on any same-type document.
"""
from __future__ import annotations
import re
import time
import io
from pathlib import Path

import cv2
import numpy as np
import fitz  # PyMuPDF
from rapidocr_onnxruntime import RapidOCR

# ──────────────────────────────────────────────────────────────────────
# Regexes (generic — based on official format specs, not specific docs)
# ──────────────────────────────────────────────────────────────────────
PAN_RX     = re.compile(r"\b([A-Z]{5}[0-9]{4}[A-Z])\b")
IFSC_RX    = re.compile(r"\b([A-Z]{4}0[A-Z0-9]{6})\b")
AADHAAR_RX = re.compile(r"\b(\d{4}\s?\d{4}\s?\d{4})\b")
DATE_RX    = re.compile(
    r"\b("
    r"\d{2}[/\-. ]\d{2}[/\-. ]\d{4}"   # 28/10/2010
    r"|\d{4}[/\-. ]\d{4}"              # OCR dropped one separator
    r"|\d{2}[/\-. ]\d{6}"
    r")\b"
)
ACCOUNT_RX = re.compile(r"\b(\d{9,18})\b")        # bank A/c numbers: 9-18 digits
MICR_RX    = re.compile(r"\b(\d{9})\b")           # 9-digit MICR codes

PAN_HOLDER_TYPE = {
    "P": "Individual", "F": "Firm", "C": "Company", "H": "HUF",
    "A": "AOP", "T": "Trust", "B": "BOI", "L": "Local Authority",
    "J": "Artificial Juridical Person", "G": "Government",
}

LABEL_WORDS = {
    "photo","signature","name","father","fathers","father's","date","birth",
    "incorporation","formation","permanent","account","number","card",
    "income","tax","department","govt","india","of","pay","rupees","valid",
    "months","bearer","bank","branch","ifsc","micr","rtgs","neft","payable",
    "for","or","at","par","through","clearing","transfer","all","branches",
    "ltd","authorised","signatories","please","sign","above","ground","floor",
}


# ──────────────────────────────────────────────────────────────────────
# Input: load PDF or image into a list of numpy BGR images
# ──────────────────────────────────────────────────────────────────────
def load_pages(path: str, dpi: int = 300) -> list[np.ndarray]:
    p = Path(path)
    if p.suffix.lower() == ".pdf":
        doc = fitz.open(p)
        pages = []
        for page in doc:
            pix = page.get_pixmap(dpi=dpi)
            arr = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.h, pix.w, pix.n)
            if pix.n == 4:
                arr = cv2.cvtColor(arr, cv2.COLOR_RGBA2BGR)
            elif pix.n == 3:
                arr = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
            elif pix.n == 1:
                arr = cv2.cvtColor(arr, cv2.COLOR_GRAY2BGR)
            pages.append(arr)
        return pages
    else:
        img = cv2.imread(str(p), cv2.IMREAD_COLOR)
        if img is None:
            raise ValueError(f"Could not read image: {p}")
        return [img]


# ──────────────────────────────────────────────────────────────────────
# OCR with rotation auto-correction
# ──────────────────────────────────────────────────────────────────────
def _rotate(img: np.ndarray, angle: int) -> np.ndarray:
    if angle == 0:   return img
    if angle == 90:  return cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
    if angle == 180: return cv2.rotate(img, cv2.ROTATE_180)
    if angle == 270: return cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
    raise ValueError(angle)


def _score(result) -> float:
    """Quality score = sum(confidence * text_length). Bigger = better OCR."""
    if not result: return 0.0
    return sum(conf * len(text) for _, text, conf in result)


def smart_ocr(img: np.ndarray, engine: RapidOCR) -> tuple[list, int, float]:
    """
    Try orientation 0 first. If score is suspiciously low, try 90/180/270.
    Returns: (result, rotation_used, score)
    """
    base_result, _ = engine(img)
    base_score = _score(base_result)

    # Heuristic threshold: if score is high, trust it. If low, try other rotations.
    if base_score > 100:
        return base_result, 0, base_score

    best = (base_result, 0, base_score)
    for rot in (90, 180, 270):
        rotated = _rotate(img, rot)
        result, _ = engine(rotated)
        score = _score(result)
        if score > best[2]:
            best = (result, rot, score)
    return best


# ──────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────
def normalise_date(s: str) -> str:
    digits = re.sub(r"\D", "", s)
    if len(digits) == 8:
        return f"{digits[0:2]}/{digits[2:4]}/{digits[4:8]}"
    return s


def to_lines(result, pass_id: str = "primary") -> list[dict]:
    """Convert raw RapidOCR boxes to sorted line dicts (top-to-bottom, left-to-right)."""
    lines = []
    for bbox, text, conf in result:
        y = sum(p[1] for p in bbox) / 4
        x = sum(p[0] for p in bbox) / 4
        lines.append({"y": y, "x": x, "text": text.strip(), "conf": conf, "pass": pass_id})
    lines.sort(key=lambda l: (l["y"], l["x"]))
    return lines


def looks_like_name(text: str, strict: bool = True) -> bool:
    if PAN_RX.search(text.upper()): return False
    if re.search(r"\d", text):       return False
    words = re.findall(r"[A-Za-z]+", text)
    if not words: return False
    if any(w.lower() in LABEL_WORDS for w in words): return False
    alpha = [c for c in text if c.isalpha()]
    if sum(c.isupper() for c in alpha) / max(len(alpha), 1) < 0.7: return False
    if strict and len(words) < 2: return False
    if not strict and len(words[0]) < 3: return False
    return True


# ──────────────────────────────────────────────────────────────────────
# Document type classifier (scoring)
# ──────────────────────────────────────────────────────────────────────
def classify(lines: list[dict]) -> str:
    text = " ".join(l["text"] for l in lines).lower()
    upper = text.upper()
    scores = {"pan": 0, "aadhaar": 0, "cheque": 0, "passport": 0}

    if "income tax" in text or "permanent account" in text: scores["pan"] += 3
    if PAN_RX.search(upper) and "pan" in text or PAN_RX.search(upper) and "account number" in text:
        scores["pan"] += 2
    if "uidai" in text or "aadhaar" in text or "unique identification" in text:
        scores["aadhaar"] += 3
    if AADHAAR_RX.search(text) and scores["aadhaar"] > 0: scores["aadhaar"] += 1
    if "ifsc" in text or "rtgs" in text or "neft" in text: scores["cheque"] += 2
    if "bank" in text and ("pay" in text or "rupees" in text): scores["cheque"] += 2
    if IFSC_RX.search(upper): scores["cheque"] += 2
    if "passport" in text or "republic of india" in text: scores["passport"] += 3
    if "given name" in text or "surname" in text: scores["passport"] += 1

    best = max(scores, key=scores.get)
    return best if scores[best] >= 2 else "unknown"


# ──────────────────────────────────────────────────────────────────────
# Extractors
# ──────────────────────────────────────────────────────────────────────
def extract_pan(lines: list[dict]) -> dict:
    full = " ".join(l["text"] for l in lines)
    pan = (PAN_RX.search(full.upper()) or [None, None])[1] if PAN_RX.search(full.upper()) else None

    # Date
    dob = None
    for l in lines:
        m = DATE_RX.search(l["text"])
        if m:
            dob = normalise_date(m.group(1))
            break

    # Name / Father's Name via labels
    name, father = None, None
    skip_next = {"photo", "signature"}
    for i, l in enumerate(lines):
        t = l["text"].lower()
        if re.search(r"\bname\b", t) and "father" not in t:
            for j in range(i + 1, len(lines)):
                cand = lines[j]["text"]
                if cand.lower() in skip_next: continue
                if looks_like_name(cand, strict=False):
                    name = name or cand
                    break
        if "father" in t:
            for j in range(i + 1, len(lines)):
                cand = lines[j]["text"]
                if cand.lower() in skip_next: continue
                if looks_like_name(cand, strict=False):
                    father = cand
                    break

    # Fallback when bilingual labels are mangled
    if not name:
        for l in lines:
            if looks_like_name(l["text"], strict=True) and l["text"] != father:
                name = l["text"]
                break

    is_company = pan and pan[3] == "C"
    return {
        "document_type": "PAN",
        "pan_number": pan,
        "pan_valid": bool(pan and re.fullmatch(r"[A-Z]{5}[0-9]{4}[A-Z]", pan)),
        "entity_type": PAN_HOLDER_TYPE.get(pan[3], "unknown") if pan else None,
        "name": name,
        "father_name": None if is_company else father,
        "dob_or_incorporation": dob,
    }


def extract_cheque(lines: list[dict]) -> dict:
    """
    Cheque field extraction.
    - Bank name: line near the top with 'BANK' in it
    - Address: lines between bank name and the IFSC/RTGS section
    - IFSC: regex
    - MICR: cluster of digits at the bottom of the cheque, specifically the 9-digit code
    - Account number: long digit string near an A/c. label, or the 14-digit-ish account near top
    """
    full = " ".join(l["text"] for l in lines)
    upper_full = full.upper()

    # ── Bank name: first line containing 'BANK' that isn't a generic phrase
    bank_name = None
    for l in lines:
        t = l["text"].strip()
        if re.search(r"\bbank\b", t, re.I):
            if re.search(r"branches of|payable", t, re.I): continue
            bank_name = t
            break

    # ── IFSC
    ifsc_match = IFSC_RX.search(upper_full)
    ifsc = ifsc_match.group(1) if ifsc_match else None

    # ── Account number: prefer long digit strings near an A/c label
    account = None
    for i, l in enumerate(lines):
        if re.search(r"a\s*/?\s*c\.?\s*no", l["text"], re.I):
            # look at this line + next 2
            for j in range(i, min(i + 3, len(lines))):
                m = re.search(r"\b(\d{9,18})\b", lines[j]["text"])
                if m:
                    account = m.group(1); break
            if account: break
    # Fallback: any long digit string (10-16) that isn't the MICR line content
    if not account:
        candidates = []
        for l in lines:
            for m in re.finditer(r"\b(\d{10,18})\b", l["text"]):
                candidates.append((m.group(1), l["y"]))
        if candidates:
            # Pick the one with most digits (account numbers tend to be longest)
            candidates.sort(key=lambda c: -len(c[0]))
            account = candidates[0][0]

    # ── MICR line: standard layout reading left-to-right is
    #     <cheque_no(6)> <MICR(9)> <account_short(6)> <txn(2 or 3)>
    # OCR may split this row into multiple boxes. Strategy:
    #   1. Restrict to the absolute bottom ~20% of the document (where MICR lives)
    #   2. Keep digit-heavy boxes with few letters or MICR-marker chars (' " = : |)
    #   3. Sort those LEFT-TO-RIGHT by x-coord, concatenate digits, parse positionally
    micr = None
    cheque_no = None
    txn_code = None
    # Only use spatially-consistent lines (primary or rotated pass) for MICR.
    spatial_lines = [l for l in lines if l.get("pass") != "original"]
    if spatial_lines:
        ys = [l["y"] for l in spatial_lines]
        y_range = max(ys) - min(ys)
        bottom_cutoff = min(ys) + 0.80 * y_range
        micr_segments = []
        for l in spatial_lines:
            if l["y"] < bottom_cutoff: continue
            t = l["text"]
            n_digit = sum(c.isdigit() for c in t)
            n_alpha = sum(c.isalpha() for c in t)
            has_marker = any(c in '"\'=:|' for c in t)
            if n_digit >= 3 and (n_alpha == 0 or has_marker) and n_alpha <= 2:
                micr_segments.append(l)
        micr_segments.sort(key=lambda l: l["x"])
        merged = "".join(re.sub(r"\D", "", s["text"]) for s in micr_segments)

        if 21 <= len(merged) <= 28:
            cheque_no = merged[0:6]
            micr      = merged[6:15]
            txn_code  = merged[-2:]
        elif len(merged) >= 15:
            cheque_no = merged[:6]
            micr      = merged[6:15]
            if len(merged) >= 17:
                txn_code = merged[-2:]

    # ── Address: lines that look like address content — long alphanumeric,
    # not a label, not the bank name, not the IFSC line.
    # OCR sometimes joins words without spaces (e.g. 'Validfor3monthsonly'),
    # so we match against the space-stripped lowercase version.
    EXCLUDE_SUBSTRINGS = [
        "ifsc", "rtgs", "neft", "rupees", "validfor", "monthsonly",
        "payableatpar", "throughclearing", "transferatall",
        "branchesof", "authorised", "signator", "bearer", "pleasesign",
        "cts-", "accountnumber", "permanentaccount", "incometax",
        "govt", "ofindia", "department",
    ]
    address_lines = []
    for l in lines:
        t = l["text"].strip()
        if l["conf"] < 0.5: continue
        if len(t) < 15: continue
        if t == bank_name: continue
        norm = re.sub(r"\s+", "", t).lower()
        if any(sub in norm for sub in EXCLUDE_SUBSTRINGS): continue
        if sum(c.isalpha() for c in t) < 8: continue
        # Skip "For COMPANY NAME" payee lines (with or without space after 'For')
        if re.match(r"^for[\s]?[A-Z]", t, re.I): continue
        address_lines.append(t)
    address = ", ".join(address_lines) if address_lines else None

    return {
        "document_type": "Cheque",
        "bank_name": bank_name,
        "address": address,
        "ifsc_code": ifsc,
        "ifsc_valid": bool(ifsc),
        "account_number": account,
        "micr_code": micr,
        "cheque_number": cheque_no,
        "transaction_code": txn_code,
    }


EXTRACTORS = {"pan": extract_pan, "cheque": extract_cheque}


# ──────────────────────────────────────────────────────────────────────
# Main pipeline
# ──────────────────────────────────────────────────────────────────────
def process(file_path: str, engine: RapidOCR) -> list[dict]:
    pages = load_pages(file_path)
    out = []
    for page_idx, img in enumerate(pages):
        t0 = time.time()
        result, rot, score = smart_ocr(img, engine)
        ocr_time = time.time() - t0

        if not result:
            out.append({"page": page_idx + 1, "error": "no text detected"})
            continue

        lines = to_lines(result)
        doc_type = classify(lines)

        # ── Cheque rotation correction ─────────────────────────────────
        # Cheques are landscape (wider than tall). If we detect a cheque on a
        # portrait page, the underlying scan/photo was rotated. RapidOCR's
        # per-line angle classifier still reads text correctly, but the spatial
        # layout is sideways — the MICR row ends up on the left/right edge
        # instead of the bottom. Rotate the image and re-OCR for correct layout,
        # then MERGE the new lines with the original (deduped by text) so that
        # text-only fields like address benefit from both OCR passes.
        if doc_type == "cheque":
            h, w = img.shape[:2]
            if h > 1.2 * w:
                for cv_rot, deg in [(cv2.ROTATE_90_COUNTERCLOCKWISE, 90),
                                    (cv2.ROTATE_90_CLOCKWISE, -90)]:
                    rotated_img = cv2.rotate(img, cv_rot)
                    r2, _ = engine(rotated_img)
                    if r2:
                        new_lines = to_lines(r2, pass_id="rotated")
                        xs = [l["x"] for l in new_lines]
                        ys = [l["y"] for l in new_lines]
                        if (max(xs) - min(xs)) > 1.3 * (max(ys) - min(ys)):
                            # Merge: rotated lines first (spatial truth), then
                            # original-pass lines tagged so spatial logic can
                            # skip them but content-matching can still use them.
                            for l in lines:
                                l["pass"] = "original"
                            seen = {l["text"].strip().lower() for l in new_lines}
                            for l in lines:
                                if l["text"].strip().lower() not in seen:
                                    new_lines.append(l)
                                    seen.add(l["text"].strip().lower())
                            lines = new_lines
                            rot = deg
                            ocr_time = time.time() - t0
                            break

        extractor = EXTRACTORS.get(doc_type)
        fields = extractor(lines) if extractor else {"document_type": "unknown"}

        out.append({
            "page": page_idx + 1,
            "rotation_applied": rot,
            "ocr_time_s": round(ocr_time, 2),
            "ocr_score": round(score, 1),
            "detected_type": doc_type,
            "fields": fields,
            "_raw_lines": [(round(l["conf"], 2), l["text"]) for l in lines],
        })
    return out


def pretty(file_path: str, engine: RapidOCR):
    print(f"\n{'='*70}\nFILE: {Path(file_path).name}\n{'='*70}")
    pages = process(file_path, engine)
    for p in pages:
        print(f"\nPage {p.get('page')}  |  rotation={p.get('rotation_applied')}°"
              f"  |  ocr={p.get('ocr_time_s')}s  |  type={p.get('detected_type')}")
        print("-" * 70)
        for k, v in p["fields"].items():
            print(f"  {k:24s}: {v}")
        print(f"\n  Raw OCR lines ({len(p['_raw_lines'])}):")
        for conf, text in p["_raw_lines"]:
            print(f"    [{conf:.2f}] {text}")


if __name__ == "__main__":
    import sys
    files = sys.argv[1:] if len(sys.argv) > 1 else []
    engine = RapidOCR()
    for f in files:
        pretty(f, engine)