Team Ai
Apppublic

Explicite/api-test

sourceHugging Faceupdated 8mo agoView on Hugging Face
0likes
doc_classifier.py179 linesDownload Raw Back to services
1# app/services/doc_classifier.py2import re3from typing import Dict, List4 5import numpy as np6 7 8def _safe_import_pytesseract():9    try:10        import pytesseract  # type: ignore11        return pytesseract12    except Exception:13        return None14 15 16def _normalize_text(s: str) -> str:17    """Normaliza el texto OCR para búsqueda de patrones"""18    s = s.upper()19    s = re.sub(r"[^A-Z0-9<\s]", " ", s)20    s = re.sub(r"\s+", " ", s).strip()21    return s22 23 24def classify_document(np_rgb: np.ndarray) -> Dict:25    """26    Clasifica el tipo de documento mediante OCR y reglas heurísticas.27    28    NO bloquea el flujo - solo añade información al JSON.29    30    Tipos soportados:31      - ine_mx: INE (Credencial para votar) de México32      - passport: Pasaporte (detecta MRZ)33      - driver_license_mx: Licencia de conducir México34      - us_driver_license: Licencia de conducir USA35      - us_id_card: ID estatal USA36      - id_card_generic: Credencial genérica37      - unknown: No se pudo determinar38    39    Returns:40        Dict con: type (str), confidence (float 0-1), signals (List[str])41    """42    signals: List[str] = []43 44    pytesseract = _safe_import_pytesseract()45    if pytesseract is None:46        return {47            "type": "unknown",48            "confidence": 0.10,49            "signals": ["ocr:not_available (install pytesseract + tesseract)"],50        }51 52    # Ejecutar OCR53    try:54        raw_text = pytesseract.image_to_string(np_rgb, lang="eng+spa", config="--psm 6")55    except Exception as e:56        return {57            "type": "unknown",58            "confidence": 0.10,59            "signals": [f"ocr:error:{type(e).__name__}"],60        }61 62    text = _normalize_text(raw_text)63    64    if len(text) < 10:65        return {66            "type": "unknown",67            "confidence": 0.15,68            "signals": ["ocr:too_little_text"],69        }70 71    # ---------- REGLAS DE CLASIFICACIÓN ----------72 73    # 1) PASAPORTE: Detectar MRZ (Machine Readable Zone)74    #    Los pasaportes tienen líneas con muchos caracteres '<'75    mrz_detected = ("<<" in raw_text and len(re.findall(r"<", raw_text)) >= 20)76    if mrz_detected:77        signals.append("mrz:detected (<< and many <)")78 79    # 2) INE (México)80    ine_keywords = [81        "INSTITUTO NACIONAL ELECTORAL",82        "CREDENCIAL PARA VOTAR",83        "ELECTOR",84        "SECCION",85        "VIGENCIA",86        "CLAVE DE ELECTOR",87        "CURP",88    ]89    ine_hits = sum(1 for k in ine_keywords if k in text)90    if ine_hits:91        signals.append(f"ine:hits={ine_hits}")92 93    # 3) Licencia de conducir México94    mx_dl_keywords = [95        "LICENCIA",96        "CONDUCIR",97        "CONDUCTOR",98        "EXPEDICION",99        "VIGENCIA",100        "FOLIO",101    ]102    mx_dl_hits = sum(1 for k in mx_dl_keywords if k in text)103    if mx_dl_hits:104        signals.append(f"mx_dl:hits={mx_dl_hits}")105 106    # 4) Driver License USA107    us_dl_keywords = [108        "DRIVER LICENSE",109        "DRIVER LICENCE",110        "DL",111        "CLASS",112        "ISS",113        "EXP",114        "DOB",115        "HEIGHT",116        "EYES",117        "SEX",118    ]119    us_dl_hits = sum(1 for k in us_dl_keywords if k in text)120    if us_dl_hits:121        signals.append(f"us_dl:hits={us_dl_hits}")122 123    # 5) ID Card USA (estatal)124    us_id_keywords = [125        "IDENTIFICATION CARD",126        "IDENTITY CARD",127        "STATE ID",128        "ID CARD",129    ]130    us_id_hits = sum(1 for k in us_id_keywords if k in text)131    if us_id_hits:132        signals.append(f"us_id:hits={us_id_hits}")133 134    # 6) Credencial genérica135    generic_id_keywords = [136        "IDENTIFICACION",137        "IDENTIFICATION",138        "ID",139        "CREDENCIAL",140    ]141    generic_hits = sum(1 for k in generic_id_keywords if k in text)142    if generic_hits:143        signals.append(f"generic:hits={generic_hits}")144 145    # ---------- DECISIÓN (prioridad de arriba hacia abajo) ----------146 147    # PASAPORTE (mayor prioridad)148    if mrz_detected:149        return {"type": "passport", "confidence": 0.90, "signals": signals}150 151    # INE México152    if ine_hits >= 2:153        confidence = min(0.95, 0.85 + 0.03 * ine_hits)154        return {"type": "ine_mx", "confidence": confidence, "signals": signals}155    if ine_hits == 1:156        return {"type": "ine_mx", "confidence": 0.65, "signals": signals}157 158    # Licencia México (y que NO sea texto de USA)159    if mx_dl_hits >= 2 and "DRIVER LICENSE" not in text:160        confidence = min(0.92, 0.75 + 0.04 * mx_dl_hits)161        return {"type": "driver_license_mx", "confidence": confidence, "signals": signals}162 163    # Licencia USA164    if us_dl_hits >= 2:165        confidence = min(0.90, 0.78 + 0.03 * us_dl_hits)166        return {"type": "us_driver_license", "confidence": confidence, "signals": signals}167 168    # ID estatal USA169    if us_id_hits >= 1:170        confidence = min(0.88, 0.72 + 0.04 * us_id_hits)171        return {"type": "us_id_card", "confidence": confidence, "signals": signals}172 173    # Credencial genérica (catch-all)174    if generic_hits >= 1:175        return {"type": "id_card_generic", "confidence": 0.55, "signals": signals}176 177    # No pudimos clasificar178    return {"type": "unknown", "confidence": 0.20, "signals": signals if signals else ["no_keywords_matched"]}179