models/document_classifier.py

"""
InvoiceForge AI — models/document_classifier.py
Keyword + layout heuristic document type classifier.
Classifies into: gst_invoice, purchase_invoice, sales_invoice,
retail_receipt, handwritten_bill, delivery_challan, credit_note, debit_note, unknown
"""
from __future__ import annotations
import logging
import re
from dataclasses import dataclass
import cv2
import numpy as np

logger = logging.getLogger(__name__)

DOCUMENT_TYPES = [
    "gst_invoice","purchase_invoice","sales_invoice","retail_receipt",
    "handwritten_bill","delivery_challan","credit_note","debit_note","unknown",
]

TYPE_KEYWORDS: dict[str, list[str]] = {
    "gst_invoice": ["tax invoice","gst","gstin","cgst","sgst","igst","e-invoice",
                    "irn","ack no","hsn","sac","taxable value","place of supply"],
    "purchase_invoice": ["purchase","purchase order","p.o.","po no","vendor",
                         "supplier","bill to","grn","goods receipt"],
    "sales_invoice": ["sales invoice","sale invoice","sold to","customer","client","receivable"],
    "retail_receipt": ["receipt","cash memo","retail","pos","cashier","change","tender",
                       "thank you","store","outlet"],
    "handwritten_bill": ["hand written","handwritten","manual bill"],
    "delivery_challan": ["delivery challan","challan","consignment","goods dispatch",
                         "delivery note","d.c.","dc no","vehicle no","lorry","transport"],
    "credit_note": ["credit note","credit memo","return note","c.n.","cn no","credit"],
    "debit_note": ["debit note","debit memo","d.n.","dn no","penalty","debit"],
}

STRONG_KEYWORDS: dict[str, list[str]] = {
    "gst_invoice": ["tax invoice"],
    "delivery_challan": ["delivery challan","challan"],
    "credit_note": ["credit note","credit memo"],
    "debit_note": ["debit note","debit memo"],
    "retail_receipt": ["cash memo","receipt"],
}


@dataclass
class ClassificationResult:
    doc_type: str
    confidence: float
    scores: dict[str, float]
    is_handwritten: bool
    is_low_quality: bool


class DocumentClassifier:
    """
    Keyword + layout heuristic document classifier.
    Usage:
        result = DocumentClassifier().classify(full_text, img_bgr)
    """

    def classify(self, full_text: str, img_bgr: np.ndarray | None = None) -> ClassificationResult:
        lower = full_text.lower()
        scores: dict[str, float] = {dt: 0.0 for dt in DOCUMENT_TYPES}

        # Keyword scoring
        for dtype, kws in TYPE_KEYWORDS.items():
            for kw in kws:
                if kw in lower:
                    scores[dtype] += 1.0

        # Strong keyword bonus (3×)
        for dtype, kws in STRONG_KEYWORDS.items():
            for kw in kws:
                if kw in lower:
                    scores[dtype] += 3.0

        is_handwritten = False
        is_low_quality = False

        if img_bgr is not None:
            sigs = self._analyse_layout(img_bgr)
            is_handwritten = sigs["is_handwritten"]
            is_low_quality = sigs["is_low_quality"]
            if is_handwritten:
                scores["handwritten_bill"] += 5.0

        active = {k: v for k, v in scores.items() if k != "unknown"}
        if all(v == 0.0 for v in active.values()):
            doc_type, confidence = "unknown", 0.0
        else:
            doc_type = max(active, key=lambda k: active[k])
            total = sum(active.values())
            confidence = round(active[doc_type] / total if total > 0 else 0.0, 4)

        logger.info("Classified '%s' conf=%.2f handwritten=%s", doc_type, confidence, is_handwritten)
        return ClassificationResult(doc_type, confidence, scores, is_handwritten, is_low_quality)

    @staticmethod
    def _analyse_layout(img_bgr: np.ndarray) -> dict:
        gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
        lap_var = float(cv2.Laplacian(gray, cv2.CV_64F).var())
        is_low_quality = lap_var < 80.0
        edges = cv2.Canny(gray, 50, 150)
        lines = cv2.HoughLinesP(edges, 1, np.pi/180, 50, minLineLength=30, maxLineGap=10)
        is_handwritten = False
        if lines is not None and len(lines) > 5:
            angles = [np.degrees(np.arctan2(y2-y1, x2-x1)) % 180
                      for x1,y1,x2,y2 in lines[:,0]]
            is_handwritten = float(np.std(angles)) > 40
        elif lines is None or len(lines) < 5:
            is_handwritten = True
        return {"is_handwritten": is_handwritten, "is_low_quality": is_low_quality,
                "blur_score": round(lap_var, 2)}