"""
╔══════════════════════════════════════════════════════════════════╗
║ CSM DUAL-CARD ID OCR SYSTEM — ARCHITECTURE NOTE ║
╠══════════════════════════════════════════════════════════════════╣
║ MODEL TASKS (8B VLM): ║
║ Step 1 → Raw OCR: All text, original script, no translate ║
║ Step 2 → Doc classify + non-English gap fill only ║
║ PYTHON TASKS (Authoritative): ║
║ MRZ parse+verify | Numeral convert | Calendar convert ║
║ English label extract | Script separate | Cross verify ║
╚══════════════════════════════════════════════════════════════════╝
"""
import os
import uuid
import time
import re
import datetime
from threading import Thread
from typing import Iterable, Dict, Any
import gradio as gr
import spaces
import torch
from PIL import Image
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
os.environ["HF_HOME"] = "/tmp/hf_home"
from transformers import (
AutoProcessor,
AutoModelForImageTextToText, # Universal VLM loader — Qwen2VL + Qwen3VL dono
TextIteratorStreamer,
BitsAndBytesConfig,
)
# Specific class imports — graceful fallback
try:
from transformers import Qwen3VLForConditionalGeneration
QWEN3_AVAILABLE = True
print("✅ Qwen3VLForConditionalGeneration available")
except ImportError:
QWEN3_AVAILABLE = False
print("⚠️ Qwen3VL direct import not available — using AutoModel fallback")
try:
from transformers import Qwen2VLForConditionalGeneration
QWEN2_AVAILABLE = True
except ImportError:
QWEN2_AVAILABLE = False
try:
from transformers import Qwen2_5_VLForConditionalGeneration
QWEN25_AVAILABLE = True
except ImportError:
QWEN25_AVAILABLE = False
try:
from peft import PeftModel, PeftConfig
PEFT_AVAILABLE = True
print("✅ PEFT available")
except ImportError:
PEFT_AVAILABLE = False
print("⚠️ PEFT not available")
from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes
# ===== THEME =====
colors.steel_blue = colors.Color(
name="steel_blue",
c50="#EBF3F8", c100="#D3E5F0", c200="#A8CCE1", c300="#7DB3D2",
c400="#529AC3", c500="#4682B4", c600="#3E72A0", c700="#36638C",
c800="#2E5378", c900="#264364", c950="#1E3450",
)
class SteelBlueTheme(Soft):
def __init__(self, *, primary_hue=colors.gray, secondary_hue=colors.steel_blue,
neutral_hue=colors.slate, text_size=sizes.text_lg,
font=(fonts.GoogleFont("Outfit"), "Arial", "sans-serif"),
font_mono=(fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace")):
super().__init__(primary_hue=primary_hue, secondary_hue=secondary_hue,
neutral_hue=neutral_hue, text_size=text_size, font=font, font_mono=font_mono)
super().set(
background_fill_primary="*primary_50",
background_fill_primary_dark="*primary_900",
body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
button_primary_text_color="white",
button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
button_secondary_text_color="black",
button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)",
button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)",
slider_color="*secondary_500",
block_title_text_weight="600",
block_border_width="3px",
block_shadow="*shadow_drop_lg",
button_primary_shadow="*shadow_drop_lg",
button_large_padding="11px",
color_accent_soft="*primary_100",
block_label_background_fill="*primary_200",
)
steel_blue_theme = SteelBlueTheme()
css = """
#main-title h1 { font-size: 2.3em !important; }
#output-title h2 { font-size: 2.2em !important; }
.ra-wrap{ width: fit-content; }
.ra-inner{ position: relative; display: inline-flex; align-items: center; gap: 0; padding: 6px;
background: var(--neutral-200); border-radius: 9999px; overflow: hidden; }
.ra-input{ display: none; }
.ra-label{ position: relative; z-index: 2; padding: 8px 16px; font-family: inherit; font-size: 14px;
font-weight: 600; color: var(--neutral-500); cursor: pointer; transition: color 0.2s; white-space: nowrap; }
.ra-highlight{ position: absolute; z-index: 1; top: 6px; left: 6px; height: calc(100% - 12px);
border-radius: 9999px; background: white; box-shadow: 0 2px 4px rgba(0,0,0,0.1);
transition: transform 0.2s, width 0.2s; }
.ra-input:checked + .ra-label{ color: black; }
.dark .ra-inner { background: var(--neutral-800); }
.dark .ra-label { color: var(--neutral-400); }
.dark .ra-highlight { background: var(--neutral-600); }
.dark .ra-input:checked + .ra-label { color: white; }
#gpu-duration-container { padding: 10px; border-radius: 8px;
background: var(--background-fill-secondary); border: 1px solid var(--border-color-primary); margin-top: 10px; }
"""
MAX_MAX_NEW_TOKENS = 4096
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("CUDA available:", torch.cuda.is_available())
if torch.cuda.is_available():
print("Device:", torch.cuda.get_device_name(0))
print("Using:", device)
# ╔══════════════════════════════════════════╗
# ║ UNIVERSAL PROMPTS ║
# ╚══════════════════════════════════════════╝
STEP1_EXTRACT_PROMPT = """You are a universal OCR engine. Transcribe ALL visible text from this document image.
OUTPUT FORMAT — fill exactly as shown:
PHOTO_PRESENT: yes/no
PHOTO_LOCATION: [describe position: top-left / top-right / center-left / not found]
SIGNATURE_PRESENT: yes/no
SIGNATURE_LOCATION: [describe position: bottom-left / bottom-right / not found]
MRZ_PRESENT: yes/no
DETECTED_LANGUAGE: [list all languages visible e.g. Arabic+English, Farsi+English, Hindi+English, Chinese, English]
---TEXT_START---
[Every word, number, symbol, label and value visible — line by line]
[Original script preserved: Arabic, Farsi, Hindi, Chinese, Cyrillic etc. — DO NOT translate here]
[Copy label AND its value together: e.g. "DATE OF BIRTH 12/05/2003"]
[MRZ lines: copy character-perfect including ALL < symbols]
[Include corner text, watermarks, small print]
---TEXT_END---
ABSOLUTE RULES:
- NEVER output pixel coordinates like (50,68) or bounding boxes — plain text ONLY
- DO NOT translate in this step — original script as-is
- DO NOT skip or summarize any field
- Copy every character exactly including < symbols in MRZ"""
STEP2_TEMPLATE = """You are a universal KYC document analyst.
The Python pipeline has already extracted English fields and parsed MRZ.
Your job is ONLY: classify document + fill gaps from non-English text.
━━━ ALREADY EXTRACTED BY PYTHON (DO NOT RE-EXTRACT) ━━━
English Fields Found Directly on Card:
{python_fields_table}
MRZ Python Parse Result:
{mrz_summary}
━━━ YOUR INPUT DATA ━━━
English text block from card:
{english_block}
Non-English original script block:
{original_block}
━━━ YOUR TASKS — ONLY THESE 3 ━━━
TASK 1: Identify document type and issuing info
- Read English block and original block
- Keywords: PASSPORT/RESIDENT CARD/NATIONAL ID/DRIVING LICENCE/بطاقة/جواز/رخصة/आधार/PAN
- Top of card = issuing country/institution (NOT person name)
TASK 2: Classify non-English labels → check if already in English fields above
- If نام (Farsi: Name) value already in Python English fields → SKIP
- If شماره ملی (National Number) already in Python fields → SKIP
- Only add fields GENUINELY missing from Python extraction
TASK 3: Transliterate non-English values NOT found in English block
- Example: محمد → Mohammad | چراغی → Cheraghi
- Dates in Shamsi/Hijri: write BOTH original AND note calendar type
(DO NOT convert — Python handles conversion)
RULES:
- NEVER copy template placeholders like [fill here] or [value]
- NEVER re-state what Python already found
- NEVER guess values not visible in card
- If all fields already covered → write "✅ All fields covered by Python extraction"
━━━ OUTPUT FORMAT ━━━
---
## 📋 Document Classification
| | |
|---|---|
| **Document Type** | |
| **Issuing Country** | |
| **Issuing Authority** | |
---
## ➕ Additional Fields (non-English only — genuinely new)
| Label (Original) | Label (English) | Value (Original) | Value (Transliterated) |
|---|---|---|---|
| [only if not in Python fields above] | | | |
---
## 🗓️ Calendar Note (if non-Gregorian dates found)
| Original Date | Calendar System | Note |
|---|---|---|
| [date as on card] | [Solar Hijri / Lunar Hijri / Buddhist] | Python will convert |
---"""
def load_vl_model(model_id: str, quantization_config=None, pre_quantized: bool = False):
"""
Universal VLM loader — Qwen2VL / Qwen3VL / any VLM
pre_quantized=True → model already has weights quantized, no extra config needed
pre_quantized=False → apply quantization_config during load
"""
load_kwargs = {
"torch_dtype": "auto",
"device_map": "auto",
"trust_remote_code": True,
}
if quantization_config is not None and not pre_quantized:
load_kwargs["quantization_config"] = quantization_config
# Try 1: Qwen3VL (newest)
if QWEN3_AVAILABLE:
try:
return Qwen3VLForConditionalGeneration.from_pretrained(
model_id, **load_kwargs).eval()
except Exception as e:
print(f" Qwen3VL failed: {e}, trying AutoModel...")
# Try 2: AutoModelForImageTextToText (universal fallback)
try:
return AutoModelForImageTextToText.from_pretrained(
model_id, **load_kwargs).eval()
except Exception as e:
print(f" AutoModel failed: {e}, trying Qwen2VL...")
# Try 3: Qwen2VL last resort
if QWEN2_AVAILABLE:
return Qwen2VLForConditionalGeneration.from_pretrained(
model_id, **load_kwargs).eval()
raise RuntimeError(f"No compatible loader found for {model_id}")
# ╔══════════════════════════════════════════╗
# ║ MODEL LOADING ║
# ╚══════════════════════════════════════════╝
print("\n" + "="*70)
print("🚀 LOADING 4 MODELS")
print("="*70)
# 4-bit BitsAndBytes config (shared for quantized models)
bnb_4bit_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True,
)
# ── Model 1: Chhagan_ML-VL-OCR-v1 (LoRA on Qwen2VL base) ──
print("\n1️⃣ Chhagan_ML-VL-OCR-v1 (LoRA Refined)...")
MODEL_ID_C1 = "Chhagan005/Chhagan_ML-VL-OCR-v1"
CHHAGAN_V1_AVAILABLE = False
processor_c1 = model_c1 = None
if PEFT_AVAILABLE:
try:
config = PeftConfig.from_pretrained(MODEL_ID_C1)
base_id = config.base_model_name_or_path
processor_c1 = AutoProcessor.from_pretrained(base_id, trust_remote_code=True)
base_c1 = load_vl_model(base_id)
model_c1 = PeftModel.from_pretrained(base_c1, MODEL_ID_C1).to(device).eval()
print(" ✅ Loaded!")
CHHAGAN_V1_AVAILABLE = True
except Exception as e:
print(f" ❌ Failed: {e}")
else:
print(" ⚠️ PEFT not available")
# ── Model 2: Chhagan-DocVL-Qwen3 (LoRA on Qwen3VL base) ──
print("\n2️⃣ Chhagan-DocVL-Qwen3 (Qwen3-VL Refined)...")
MODEL_ID_C2 = "Chhagan005/Chhagan-DocVL-Qwen3"
CHHAGAN_QWEN3_AVAILABLE = False
processor_c2 = model_c2 = None
if PEFT_AVAILABLE:
try:
config = PeftConfig.from_pretrained(MODEL_ID_C2)
base_id = config.base_model_name_or_path
processor_c2 = AutoProcessor.from_pretrained(base_id, trust_remote_code=True)
base_c2 = load_vl_model(base_id)
model_c2 = PeftModel.from_pretrained(base_c2, MODEL_ID_C2).to(device).eval()
print(" ✅ Loaded!")
CHHAGAN_QWEN3_AVAILABLE = True
except Exception as e:
print(f" ❌ Failed: {e}")
else:
print(" ⚠️ PEFT not available")
# ── Model 3: CSM-DocExtract-VL-Q4KM (Full Qwen3VL, pre-quantized) ──
print("\n3️⃣ CSM-DocExtract-VL-Q4KM (Full Qwen3VL, pre-quantized BNB)...")
MODEL_ID_Q4KM = "Chhagan005/CSM-DocExtract-VL-Q4KM"
CSM_Q4KM_AVAILABLE = False
processor_q4km = model_q4km = None
try:
processor_q4km = AutoProcessor.from_pretrained(
MODEL_ID_Q4KM, trust_remote_code=True
)
# Pre-quantized safetensors → torch_dtype=auto, NO extra quantization_config
model_q4km = Qwen3VLForConditionalGeneration.from_pretrained(
MODEL_ID_Q4KM,
torch_dtype="auto",
device_map="auto",
trust_remote_code=True,
).eval()
print(" ✅ Loaded! (Qwen3VL pre-quantized BNB ~6.4GB)")
CSM_Q4KM_AVAILABLE = True
except Exception as e:
try:
model_q4km = AutoModelForImageTextToText.from_pretrained(
MODEL_ID_Q4KM,
torch_dtype="auto",
device_map="auto",
trust_remote_code=True,
).eval()
print(" ✅ Loaded! (AutoModel fallback)")
CSM_Q4KM_AVAILABLE = True
except Exception as e2:
print(f" ❌ Failed: {e2}")
# ── Model 4: CSM-DocExtract-VL (Full Qwen3VL, BNB INT4 trained) ──
print("\n4️⃣ CSM-DocExtract-VL 4BNB (Full Qwen3VL, BNB INT4 trained)...")
MODEL_ID_4BNB = "Chhagan005/CSM-DocExtract-VL"
CSM_4BNB_AVAILABLE = False
processor_4bnb = model_4bnb = None
system_prompt_4bnb = "You are a helpful assistant." # default
try:
# Read custom system_prompt.txt — this model was trained with it
try:
from huggingface_hub import hf_hub_download
sp_path = hf_hub_download(repo_id=MODEL_ID_4BNB, filename="system_prompt.txt")
with open(sp_path, "r", encoding="utf-8") as f:
system_prompt_4bnb = f.read().strip()
print(f" 📋 system_prompt.txt loaded: {system_prompt_4bnb[:80]}...")
except Exception as sp_err:
print(f" ⚠️ system_prompt.txt not loaded: {sp_err} — using default")
processor_4bnb = AutoProcessor.from_pretrained(
MODEL_ID_4BNB, trust_remote_code=True
)
# BNB INT4 trained safetensors → torch_dtype=auto, NO extra quantization_config
# (ignore .gguf files — those are for llama.cpp, not transformers)
model_4bnb = Qwen3VLForConditionalGeneration.from_pretrained(
MODEL_ID_4BNB,
torch_dtype="auto",
device_map="auto",
trust_remote_code=True,
ignore_mismatched_sizes=True, # GGUF files present — ignore safely
).eval()
print(" ✅ Loaded! (Qwen3VL BNB INT4 trained ~6.4GB)")
CSM_4BNB_AVAILABLE = True
except Exception as e:
try:
model_4bnb = AutoModelForImageTextToText.from_pretrained(
MODEL_ID_4BNB,
torch_dtype="auto",
device_map="auto",
trust_remote_code=True,
).eval()
print(" ✅ Loaded! (AutoModel fallback)")
CSM_4BNB_AVAILABLE = True
except Exception as e2:
print(f" ❌ Failed: {e2}")
print("\n" + "="*70)
print("📊 MODEL STATUS")
print("="*70)
status = [
("Chhagan_ML-VL-OCR-v1", CHHAGAN_V1_AVAILABLE, "LoRA Fine-tuned"),
("Chhagan-DocVL-Qwen3", CHHAGAN_QWEN3_AVAILABLE, "Qwen3-VL Fine-tuned"),
("CSM-DocExtract-Q4KM", CSM_Q4KM_AVAILABLE, "Qwen3VL Q4KM pre-quantized"),
("CSM-DocExtract-4BNB", CSM_4BNB_AVAILABLE, "Qwen3VL BitsAndBytes 4-bit"),
]
for name, ok, note in status:
print(f" {'✅' if ok else '❌'} {name:<35} {note}")
print("="*70)
loaded = sum(x[1] for x in status)
print(f" Total loaded: {loaded}/4\n")
# ╔══════════════════════════════════════════╗
# ║ PYTHON PIPELINE FUNCTIONS ║
# ╚══════════════════════════════════════════╝
def convert_eastern_numerals(text: str) -> str:
"""P2: Convert Persian/Arabic/Devanagari numerals to Western 0-9"""
tables = [
str.maketrans('۰۱۲۳۴۵۶۷۸۹', '0123456789'), # Persian
str.maketrans('٠١٢٣٤٥٦٧٨٩', '0123456789'), # Arabic
str.maketrans('०१२३४५६७८९', '0123456789'), # Devanagari
str.maketrans('০১২৩৪৫৬৭৮৯', '0123456789'), # Bengali
str.maketrans('੦੧੨੩੪੫੬੭੮੯', '0123456789'), # Gurmukhi
]
for table in tables:
text = text.translate(table)
return text
def detect_calendar_system(raw_text: str) -> str:
"""Detect calendar system from country/language context"""
text_upper = raw_text.upper()
if any(kw in raw_text for kw in ['جمهوری اسلامی ایران', 'IRAN', 'AFGHANISTAN', 'افغانستان']):
return 'solar_hijri'
if any(kw in text_upper for kw in ['SAUDI', 'ARABIA', 'السعودية', 'KUWAIT', 'QATAR', 'BAHRAIN', 'JORDAN']):
return 'lunar_hijri'
return 'gregorian'
def convert_shamsi_to_gregorian(shamsi_date: str) -> str:
"""P3: Solar Hijri (Shamsi) → Gregorian using khayyam library"""
try:
import khayyam
parts = re.split(r'[/\-\.]', shamsi_date.strip())
if len(parts) == 3:
y, m, d = int(parts[0]), int(parts[1]), int(parts[2])
jd = khayyam.JalaliDate(y, m, d)
greg = jd.todate()
return f"{greg.day:02d}/{greg.month:02d}/{greg.year}"
except ImportError:
# Approximate manual conversion if khayyam not installed
try:
parts = re.split(r'[/\-\.]', shamsi_date.strip())
y, m, d = int(parts[0]), int(parts[1]), int(parts[2])
greg_year = y + 621
return f"{d:02d}/{m:02d}/{greg_year} (approx)"
except:
pass
except Exception:
pass
return f"{shamsi_date} (Shamsi)"
def convert_hijri_to_gregorian(hijri_date: str) -> str:
"""P3: Lunar Hijri → Gregorian using hijri library"""
try:
from hijri_converter import convert
parts = re.split(r'[/\-\.]', hijri_date.strip())
if len(parts) == 3:
y, m, d = int(parts[0]), int(parts[1]), int(parts[2])
greg = convert.Hijri(y, m, d).to_gregorian()
return f"{greg.day:02d}/{greg.month:02d}/{greg.year}"
except ImportError:
try:
parts = re.split(r'[/\-\.]', hijri_date.strip())
y, m, d = int(parts[0]), int(parts[1]), int(parts[2])
greg_year = y - 43 + 622
return f"{d:02d}/{m:02d}/{greg_year} (approx)"
except:
pass
except:
pass
return f"{hijri_date} (Hijri)"
def separate_scripts(raw_text: str) -> tuple:
"""P5: Separate English/Latin lines from non-Latin script lines"""
english_lines = []
original_lines = []
for line in raw_text.split('\n'):
line = line.strip()
if not line:
continue
non_latin = sum(1 for c in line if ord(c) > 591)
total_alpha = sum(1 for c in line if c.isalpha())
if total_alpha == 0:
english_lines.append(line)
elif non_latin / max(total_alpha, 1) > 0.4:
original_lines.append(line)
else:
english_lines.append(line)
return '\n'.join(english_lines), '\n'.join(original_lines)
def extract_english_fields(raw_text: str) -> list:
"""P4: Extract English label:value pairs directly from card text — no AI"""
results = []
patterns = [
(r'(?:FULL\s+)?NAME\s*[:\-.]?\s*([A-Za-z][A-Za-z\s\-\.\']{1,60})', 'NAME'),
(r'DATE\s+OF\s+BIRTH\s*[:\-.]?\s*(\d{1,2}[\s/\-\.]\d{1,2}[\s/\-\.]\d{2,4})', 'DATE OF BIRTH'),
(r'\bDOB\s*[:\-.]?\s*(\d{1,2}[\s/\-\.]\d{1,2}[\s/\-\.]\d{2,4})', 'DATE OF BIRTH'),
(r'BIRTH\s+DATE\s*[:\-.]?\s*(\d{1,2}[\s/\-\.]\d{1,2}[\s/\-\.]\d{2,4})', 'DATE OF BIRTH'),
(r'EXPIRY\s+DATE\s*[:\-.]?\s*(\d{1,2}[\s/\-\.]\d{1,2}[\s/\-\.]\d{2,4})', 'EXPIRY DATE'),
(r'DATE\s+OF\s+EXPIRY\s*[:\-.]?\s*(\d{1,2}[\s/\-\.]\d{1,2}[\s/\-\.]\d{2,4})', 'EXPIRY DATE'),
(r'VALID(?:\s+THRU|\s+UNTIL|ITY)?\s*[:\-.]?\s*(\d{1,2}[\s/\-\.]\d{1,2}[\s/\-\.]\d{2,4})', 'EXPIRY DATE'),
(r'EXPIRATION\s+DATE\s*[:\-.]?\s*(\d{1,2}[\s/\-\.]\d{1,2}[\s/\-\.]\d{2,4})', 'EXPIRY DATE'),
(r'(?:DATE\s+OF\s+)?ISSUE\s+DATE\s*[:\-.]?\s*(\d{1,2}[\s/\-\.]\d{1,2}[\s/\-\.]\d{2,4})', 'ISSUE DATE'),
(r'DATE\s+OF\s+ISSUE\s*[:\-.]?\s*(\d{1,2}[\s/\-\.]\d{1,2}[\s/\-\.]\d{2,4})', 'ISSUE DATE'),
(r'CIVIL\s+(?:NO\.?|NUMBER)\s*[:\-.]?\s*([A-Z0-9\-]{4,20})', 'CIVIL NUMBER'),
(r'PASSPORT\s+(?:NO\.?|NUMBER)\s*[:\-.]?\s*([A-Z0-9\-]{6,12})', 'PASSPORT NUMBER'),
(r'LICENCE\s+(?:NO\.?|NUMBER)\s*[:\-.]?\s*([A-Z0-9\-]{4,20})', 'LICENCE NUMBER'),
(r'LICENSE\s+(?:NO\.?|NUMBER)\s*[:\-.]?\s*([A-Z0-9\-]{4,20})', 'LICENCE NUMBER'),
(r'AADHAAR\s*(?:NO\.?|NUMBER)?\s*[:\-.]?\s*(\d{4}\s?\d{4}\s?\d{4})', 'AADHAAR NUMBER'),
(r'\bPAN\s*[:\-.]?\s*([A-Z]{5}\d{4}[A-Z])', 'PAN NUMBER'),
(r'EMIRATES\s+ID\s*[:\-.]?\s*(\d{3}-\d{4}-\d{7}-\d)', 'EMIRATES ID'),
(r'(?:NATIONAL\s+)?ID\s+(?:NO\.?|NUMBER)\s*[:\-.]?\s*([A-Z0-9\-]{4,20})', 'ID NUMBER'),
(r'DOCUMENT\s+(?:NO\.?|NUMBER)\s*[:\-.]?\s*([A-Z0-9\-]{4,20})', 'DOCUMENT NUMBER'),
(r'NATIONALITY\s*[:\-.]?\s*([A-Za-z]{3,30})', 'NATIONALITY'),
(r'(?:GENDER|SEX)\s*[:\-.]?\s*(MALE|FEMALE)', 'GENDER'),
(r'PLACE\s+OF\s+BIRTH\s*[:\-.]?\s*([A-Za-z\s,]{2,40})', 'PLACE OF BIRTH'),
(r'(?:PERMANENT\s+)?ADDRESS\s*[:\-.]?\s*(.{5,80})', 'ADDRESS'),
(r'BLOOD\s+(?:GROUP|TYPE)\s*[:\-.]?\s*([ABO]{1,2}[+-]?)', 'BLOOD GROUP'),
(r'(?:PROFESSION|OCCUPATION|JOB\s+TITLE)\s*[:\-.]?\s*(.{3,50})', 'PROFESSION'),
(r'FATHER(?:\'?S)?\s+NAME\s*[:\-.]?\s*([A-Za-z\s]{3,50})', "FATHER'S NAME"),
(r'MOTHER(?:\'?S)?\s+NAME\s*[:\-.]?\s*([A-Za-z\s]{3,50})', "MOTHER'S NAME"),
(r'EMPLOYER\s*[:\-.]?\s*(.{3,60})', 'EMPLOYER'),
]
seen = set()
for pattern, label in patterns:
m = re.search(pattern, raw_text, re.IGNORECASE)
if m and label not in seen:
val = m.group(1).strip()
if val and len(val) > 1 and '[' not in val:
results.append((label, val))
seen.add(label)
return results
def parse_mrz_lines(raw_text: str) -> dict:
"""P1: Authoritative Python MRZ parser — TD1, TD3, MRVA, MRVB"""
# Normalize: western numerals only
raw_text = convert_eastern_numerals(raw_text)
lines = []
for line in raw_text.split('\n'):
clean = re.sub(r'\s+', '', line.strip())
if re.match(r'^[A-Z0-9<]{25,50}$', clean):
lines.append(clean)
if not lines:
return {}
def decode_date(yymmdd: str, is_dob: bool = False) -> str:
try:
yy, mm, dd = int(yymmdd[0:2]), int(yymmdd[2:4]), int(yymmdd[4:6])
if not (1 <= mm <= 12 and 1 <= dd <= 31):
return f"Invalid ({yymmdd})"
cur_yy = datetime.datetime.now().year % 100
year = (1900 + yy) if (is_dob and yy > cur_yy) else (2000 + yy)
return f"{dd:02d}/{mm:02d}/{year}"
except:
return yymmdd
def clean_fill(s: str) -> str:
return re.sub(r'<+$', '', s).replace('<', ' ').strip()
def parse_name(line3: str) -> str:
name_clean = re.sub(r'<+$', '', line3)
if '<<' in name_clean:
parts = name_clean.split('<<')
surname = parts[0].replace('<', ' ').strip().title()
given = parts[1].replace('<', ' ').strip().title() if len(parts) > 1 else ''
return f"{given} {surname}".strip() if given else surname
return name_clean.replace('<', ' ').strip().title()
result = {}
# TD1: 3 lines, 28-36 chars
td1 = [l for l in lines if 28 <= len(l) <= 36]
if len(td1) >= 2:
l1, l2 = td1[0], td1[1]
l3 = td1[2] if len(td1) > 2 else ""
result['doc_type'] = clean_fill(l1[0:2])
result['country_code'] = clean_fill(l1[2:5])
result['doc_number'] = clean_fill(l1[5:14])
if len(l2) >= 19:
result['dob'] = decode_date(l2[0:6], is_dob=True)
sex = l2[7] if len(l2) > 7 else ''
result['sex'] = 'Male' if sex == 'M' else ('Female' if sex == 'F' else 'Unknown')
result['expiry'] = decode_date(l2[8:14], is_dob=False)
result['nationality'] = clean_fill(l2[15:18])
if l3:
result['name'] = parse_name(l3)
result['mrz_format'] = 'TD1'
return result
# TD3: 2 lines, 40-48 chars (Passports)
td3 = [l for l in lines if 40 <= len(l) <= 48]
if len(td3) >= 2:
l1, l2 = td3[0], td3[1]
result['doc_type'] = clean_fill(l1[0:2])
result['country_code'] = clean_fill(l1[2:5])
result['name'] = parse_name(l1[5:44])
if len(l2) >= 27:
result['doc_number'] = clean_fill(l2[0:9])
result['nationality'] = clean_fill(l2[10:13])
result['dob'] = decode_date(l2[13:19], is_dob=True)
sex = l2[20] if len(l2) > 20 else ''
result['sex'] = 'Male' if sex == 'M' else ('Female' if sex == 'F' else 'Unknown')
result['expiry'] = decode_date(l2[21:27], is_dob=False)
result['mrz_format'] = 'TD3'
return result
# MRVA/MRVB: 2 lines, 36 chars (Visas)
mrv = [l for l in lines if 36 <= len(l) <= 38]
if len(mrv) >= 2:
l1, l2 = mrv[0], mrv[1]
result['doc_type'] = clean_fill(l1[0:2])
result['country_code'] = clean_fill(l1[2:5])
result['name'] = parse_name(l1[5:36])
if len(l2) >= 27:
result['doc_number'] = clean_fill(l2[0:9])
result['nationality'] = clean_fill(l2[10:13])
result['dob'] = decode_date(l2[13:19], is_dob=True)
sex = l2[20] if len(l2) > 20 else ''
result['sex'] = 'Male' if sex == 'M' else ('Female' if sex == 'F' else 'Unknown')
result['expiry'] = decode_date(l2[21:27], is_dob=False)
result['mrz_format'] = 'MRVA/MRVB'
return result
return {}
def build_mrz_table(mrz_data: dict) -> str:
if not mrz_data:
return "No MRZ detected."
table = f"**Python Parsed MRZ — Authoritative ({mrz_data.get('mrz_format','?')} format):**\n\n"
table += "| Field | Verified Value |\n|---|---|\n"
fields = [
('mrz_format', 'MRZ Format'),
('doc_type', 'Document Type'),
('country_code', 'Issuing Country Code'),
('doc_number', 'Document / Civil Number'),
('name', 'Full Name'),
('dob', 'Date of Birth'),
('expiry', 'Expiry Date'),
('nationality', 'User Nationality'),
('sex', 'Gender'),
]
for key, label in fields:
if key in mrz_data:
table += f"| {label} | **{mrz_data[key]}** ✅ |\n"
return table
def build_unified_summary(front_result: str, back_result: str, mrz_data: dict) -> str:
"""P6: Merge front+back fields, MRZ as ground truth override"""
summary = "## 🔄 Unified Deduplicated Record\n\n"
if mrz_data:
summary += f"> ✅ *MRZ Python-parsed ({mrz_data.get('mrz_format','?')}) — MRZ values are **ground truth**.*\n\n"
summary += "### 🔐 MRZ Ground Truth\n\n"
summary += build_mrz_table(mrz_data) + "\n\n---\n\n"
else:
summary += "> *No MRZ — fields merged from front+back. Conflicts flagged ⚠️.*\n\n"
def get_rows(text):
rows = {}
m = re.search(r"## (?:✅|🗂️)[^\n]*\n\|[^\n]*\n\|[-| ]+\n(.*?)(?=\n---|\Z)", text, re.DOTALL)
if m:
for line in m.group(1).strip().split('\n'):
parts = [p.strip() for p in line.split('|') if p.strip()]
if len(parts) >= 2:
field = re.sub(r'[^\w\s/\']', '', parts[0]).strip()
val = parts[1].strip()
if val and val.lower() not in ('—', 'not on card', 'n/a', ''):
rows[field] = val
return rows
front_f = get_rows(front_result)
back_f = get_rows(back_result)
all_f = list(dict.fromkeys(list(front_f.keys()) + list(back_f.keys())))
# MRZ lookup
mrz_map = {}
if mrz_data:
kw_map = {
'name': ['name'],
'doc_number': ['civil', 'document', 'id', 'passport', 'licence'],
'dob': ['birth', 'dob'],
'expiry': ['expiry', 'expiration'],
'sex': ['gender', 'sex'],
'nationality':['nationality'],
}
for mk, keywords in kw_map.items():
if mk in mrz_data:
for kw in keywords:
mrz_map[kw] = mrz_data[mk]
def get_mrz(field):
fl = field.lower()
for kw, v in mrz_map.items():
if kw in fl:
return v
return None
summary += "### 📋 Field Comparison\n\n| Field | Value | Source |\n|---|---|---|\n"
for field in all_f:
fv = front_f.get(field, '')
bv = back_f.get(field, '')
mv = get_mrz(field)
if fv and bv:
if fv.lower() == bv.lower():
note = f"✅ MRZ Confirmed" if mv and any(x in fv.lower() for x in mv.lower().split()) else ("⚠️ MRZ differs: **" + mv + "**" if mv else "")
summary += f"| {field} | {fv} | Front+Back ✅ {note} |\n"
else:
if mv:
summary += f"| {field} | ~~{fv}~~ / ~~{bv}~~ → **{mv}** | ✅ MRZ Override |\n"
else:
summary += f"| {field} | F: **{fv}** / B: **{bv}** | ⚠️ Mismatch |\n"
elif fv:
note = f"✅ MRZ Confirmed" if mv and any(x in fv.lower() for x in mv.lower().split()) else (f"⚠️ MRZ: **{mv}**" if mv else "")
summary += f"| {field} | {fv} | Front only {note} |\n"
elif bv:
note = f"✅ MRZ Confirmed" if mv and any(x in bv.lower() for x in mv.lower().split()) else (f"⚠️ MRZ: **{mv}**" if mv else "")
summary += f"| {field} | {bv} | Back only {note} |\n"
return summary + "\n"
# ╔══════════════════════════════════════════╗
# ║ STEP PIPELINE FUNCTIONS ║
# ╚══════════════════════════════════════════╝
def run_step1_extraction(model, processor, image, device, temperature, top_p, top_k, repetition_penalty, system_prompt=None):
"""Step 1: LLM → Raw OCR, original script, NO translation, NO coordinates"""
def _generate(prompt_text):
try:
from qwen_vl_utils import process_vision_info
HAS_QWEN_VL_UTILS = True
except ImportError:
HAS_QWEN_VL_UTILS = False
sys_msg = system_prompt or "You are a helpful assistant."
messages = [
{"role": "system", "content": sys_msg},
{"role": "user", "content": [
{"type": "image", "image": image},
{"type": "text", "text": prompt_text},
]}
]
# Step A: Build prompt string
try:
prompt = processor.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
if not isinstance(prompt, str):
raise TypeError("non-string returned")
except Exception:
# Manual Qwen3VL token format — universal fallback
prompt = (
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
"<|im_start|>user\n"
"<|vision_start|><|image_pad|><|vision_end|>"
f"{prompt_text}<|im_end|>\n"
"<|im_start|>assistant\n"
)
# Step B: Build inputs — 3 fallback tiers
inputs = None
# Tier 1: qwen_vl_utils + images/videos kwargs (Qwen3VL standard)
if HAS_QWEN_VL_UTILS and inputs is None:
try:
image_inputs, video_inputs = process_vision_info(messages)
proc_kwargs = {
"text": [prompt],
"padding": True,
"return_tensors": "pt"
}
if image_inputs is not None and len(image_inputs) > 0:
proc_kwargs["images"] = image_inputs
if video_inputs is not None and len(video_inputs) > 0:
proc_kwargs["videos"] = video_inputs
inputs = processor(**proc_kwargs).to(device)
print(" ✅ Tier1: qwen_vl_utils")
except Exception as e:
print(f" Tier1 failed: {e}")
inputs = None
# Tier 2: Direct PIL image (Qwen2VL style)
if inputs is None:
try:
inputs = processor(
text=[prompt],
images=[image],
padding=True,
return_tensors="pt",
).to(device)
print(" ✅ Tier2: direct PIL")
except Exception as e:
print(f" Tier2 failed: {e}")
inputs = None
# Tier 3: Text-only (last resort)
if inputs is None:
print(" ⚠️ Tier3: text-only fallback (no image — degraded)")
inputs = processor(
text=[prompt],
padding=True,
return_tensors="pt",
).to(device)
with torch.no_grad():
out = model.generate(
**inputs,
max_new_tokens=600,
do_sample=True,
temperature=temperature,
top_p=top_p,
top_k=top_k,
repetition_penalty=repetition_penalty,
)
gen = out[:, inputs['input_ids'].shape[1]:]
decoded = processor.batch_decode(gen, skip_special_tokens=True)
if isinstance(decoded, list):
return decoded[0] if decoded else ""
return str(decoded) if decoded else ""
result = _generate(STEP1_EXTRACT_PROMPT)
# Coordinate output detect → retry with simpler prompt
if re.search(r'\(\d{1,4},\s*\d{1,4}\)', result) or '---TEXT_START---' not in result:
print(" ⚠️ Retrying with fallback prompt...")
fallback = (
"Read all text from this document image and write it line by line in plain text.\n"
"Do NOT output coordinates or bounding boxes.\n"
"Start output with:\n"
"PHOTO_PRESENT: yes or no\n"
"SIGNATURE_PRESENT: yes or no\n"
"MRZ_PRESENT: yes or no\n"
"DETECTED_LANGUAGE: name the language(s)\n"
"---TEXT_START---\n"
"[all text here exactly as printed]\n"
"---TEXT_END---"
)
result = _generate(fallback)
return result
def parse_step1_output(raw_output: str) -> dict:
"""Parse Step 1 structured output → metadata + original text"""
result = {
"photo_present": "❌ No",
"photo_location": "N/A",
"sig_present": "❌ No",
"sig_location": "N/A",
"mrz_present": "❌ No",
"detected_lang": "Unknown",
"original_text": raw_output,
}
def get(pattern, text, default="N/A"):
m = re.search(pattern, text, re.IGNORECASE)
return m.group(1).strip() if m else default
photo = get(r'PHOTO_PRESENT:\s*(yes|no)', raw_output)
result["photo_present"] = "✅ Yes" if photo.lower() == "yes" else "❌ No"
result["photo_location"] = get(r'PHOTO_LOCATION:\s*([^\n]+)', raw_output)
sig = get(r'SIGNATURE_PRESENT:\s*(yes|no)', raw_output)
result["sig_present"] = "✅ Yes" if sig.lower() == "yes" else "❌ No"
result["sig_location"] = get(r'SIGNATURE_LOCATION:\s*([^\n]+)', raw_output)
mrz = get(r'MRZ_PRESENT:\s*(yes|no)', raw_output)
result["mrz_present"] = "✅ Yes" if mrz.lower() == "yes" else "❌ No"
result["detected_lang"] = get(r'DETECTED_LANGUAGE:\s*([^\n]+)', raw_output, "Unknown")
m = re.search(r'---TEXT_START---\n?(.*?)---TEXT_END---', raw_output, re.DOTALL)
if m:
result["original_text"] = m.group(1).strip()
return result
def run_step2_structure(model, processor, metadata: dict, device,
max_new_tokens, temperature, top_p, top_k, repetition_penalty):
"""Step 2: Python extracts English fields + MRZ. LLM only classifies + fills gaps."""
raw_text = metadata.get('original_text', '')
# P2: Convert eastern numerals first
raw_text_normalized = convert_eastern_numerals(raw_text)
# P5: Separate scripts
english_block, original_block = separate_scripts(raw_text_normalized)
# P4: Direct English field extraction
english_fields = extract_english_fields(raw_text_normalized)
# P1: MRZ parse (authoritative)
mrz_data = parse_mrz_lines(raw_text_normalized)
# P3: Calendar detection + conversion (for display)
calendar_sys = detect_calendar_system(raw_text)
# Build python fields table
if english_fields:
tbl = "| Field (as printed on card) | Value (as printed) |\n|---|---|\n"
for label, val in english_fields:
tbl += f"| **{label}** | {val} |\n"
else:
tbl = "| — | No English label:value pairs detected |\n"
# MRZ summary
if mrz_data:
mrz_summary = " | ".join([f"{k}: {v}" for k, v in mrz_data.items() if k != 'mrz_format'])
mrz_summary = f"✅ {mrz_data.get('mrz_format','?')} parsed: {mrz_summary}"
else:
mrz_summary = "❌ No MRZ detected"
# Non-Gregorian note
cal_note = ""
if calendar_sys == 'solar_hijri':
cal_note = "\n> ⚠️ **Solar Hijri (Shamsi) calendar detected** — Python will convert dates to Gregorian."
elif calendar_sys == 'lunar_hijri':
cal_note = "\n> ⚠️ **Lunar Hijri calendar detected** — Python will convert dates to Gregorian."
# Build prompt for LLM (classification + gaps only)
prompt_text = STEP2_TEMPLATE.format(
python_fields_table=tbl,
mrz_summary=mrz_summary,
english_block=english_block or "None",
original_block=original_block or "None",
)
messages = [{"role": "user", "content": [{"type": "text", "text": prompt_text}]}]
try:
prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
except:
prompt = prompt_text
inputs = processor(
text=[prompt],
padding=True,
return_tensors="pt",
).to(device)
streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
gen_kwargs = {
**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens,
"do_sample": True, "temperature": temperature, "top_p": top_p,
"top_k": top_k, "repetition_penalty": repetition_penalty,
}
thread = Thread(target=model.generate, kwargs=gen_kwargs)
thread.start()
# Pre-build Python-verified sections
# ── Pre-compute outside f-string (backslash fix for Python < 3.12) ──
newline = "\n"
mrz_pattern = r'^[A-Z0-9<]{25,50}$'
ws_pattern = r'\s+'
mrz_raw_lines = []
for _l in raw_text.split("\n"):
_c = re.sub(ws_pattern, '', _l.strip())
if re.match(mrz_pattern, _c):
mrz_raw_lines.append(_c)
mrz_raw_display = newline.join(mrz_raw_lines) if mrz_raw_lines else "NOT PRESENT"
mrz_table_str = build_mrz_table(mrz_data) if mrz_data else "_No MRZ detected._"
# Pre-build Python-verified sections
python_sections = (
"## 🖼️ Visual Elements\n\n"
"| Element | Status | Location |\n"
"|---------|--------|----------|\n"
f"| 📷 Profile Photo | {metadata['photo_present']} | {metadata['photo_location']} |\n"
f"| ✍️ Signature | {metadata['sig_present']} | {metadata['sig_location']} |\n"
f"| 🔐 MRZ Zone | {metadata['mrz_present']} | Bottom strip |\n\n"
"---\n\n"
"## ✅ English Fields (Direct from Card — Not Modified)\n"
f"{cal_note}\n\n"
f"{tbl}\n\n"
"---\n\n"
"## 📜 Original Script\n\n"
"```\n"
f"{raw_text}\n"
"```\n\n"
"---\n\n"
"## 🔐 MRZ Data\n\n"
"```\n"
f"{mrz_raw_display}\n"
"```\n\n"
f"{mrz_table_str}\n\n"
"---\n\n"
)
return streamer, thread, mrz_data, python_sections
# ╔══════════════════════════════════════════╗
# ║ GRADIO HELPER CLASSES ║
# ╚══════════════════════════════════════════╝
class RadioAnimated(gr.HTML):
def __init__(self, choices, value=None, **kwargs):
if not choices or len(choices) < 2:
raise ValueError("RadioAnimated requires at least 2 choices.")
if value is None:
value = choices[0]
uid = uuid.uuid4().hex[:8]
group_name = f"ra-{uid}"
inputs_html = "\n".join(
f''
f''
for i, c in enumerate(choices)
)
html_template = f"""
"""
js_on_load = r"""
(() => {
const highlight = element.querySelector('.ra-highlight');
const inputs = Array.from(element.querySelectorAll('.ra-input'));
if (!inputs.length) return;
const choices = inputs.map(i => i.value);
function setHighlight(idx) {
highlight.style.width = `calc(${100/choices.length}% - 6px)`;
highlight.style.transform = `translateX(${idx * 100}%)`;
}
function setVal(val, trigger=false) {
const idx = Math.max(0, choices.indexOf(val));
inputs.forEach((inp, i) => { inp.checked = (i === idx); });
setHighlight(idx);
props.value = choices[idx];
if (trigger) trigger('change', props.value);
}
setVal(props.value ?? choices[0], false);
inputs.forEach(inp => inp.addEventListener('change', () => setVal(inp.value, true)));
})();"""
super().__init__(value=value, html_template=html_template, js_on_load=js_on_load, **kwargs)
def apply_gpu_duration(val: str):
return int(val)
def calc_timeout_duration(model_name, text, image_front, image_back,
max_new_tokens, temperature, top_p, top_k,
repetition_penalty, gpu_timeout):
try:
base = int(gpu_timeout)
return base * 2 if (image_front is not None and image_back is not None) else base
except:
return 180
# ╔══════════════════════════════════════════╗
# ║ MAIN PIPELINE FUNCTION ║
# ╚══════════════════════════════════════════╝
@spaces.GPU(duration=calc_timeout_duration)
def generate_dual_card_ocr(model_name: str, text: str,
image_front: Image.Image, image_back: Image.Image,
max_new_tokens: int, temperature: float, top_p: float,
top_k: int, repetition_penalty: float, gpu_timeout: int):
# Model selection
model_map = {
"Chhagan-ID-OCR-v1 ⭐": (CHHAGAN_V1_AVAILABLE, processor_c1, model_c1),
"Chhagan-DocVL-Qwen3 🔥": (CHHAGAN_QWEN3_AVAILABLE, processor_c2, model_c2),
"CSM-DocExtract-Q4KM 🏆": (CSM_Q4KM_AVAILABLE, processor_q4km, model_q4km),
"CSM-DocExtract-4BNB 💎": (CSM_4BNB_AVAILABLE, processor_4bnb, model_4bnb),
}
if model_name not in model_map:
yield "Invalid model.", "Invalid model."; return
available, processor, model = model_map[model_name]
if not available:
yield f"{model_name} not available.", f"{model_name} not available."; return
if image_front is None and image_back is None:
yield "Please upload at least one card image.", "Please upload at least one card image."; return
full_output = ""
front_result = ""
back_result = ""
all_mrz_data = {}
front_meta_saved = {}
back_meta_saved = {}
# ───── FRONT CARD ─────
if image_front is not None:
full_output += "# 🎴 FRONT CARD\n\n"
full_output += "⏳ **Step 1/2 — Raw OCR (original script, no translation)...**\n\n"
yield full_output, full_output
# Model 4 ke liye system prompt pass karo
sys_p = system_prompt_4bnb if model_name == "CSM-DocExtract-4BNB 💎" else None
step1_raw = run_step1_extraction(model, processor, image_front, device,
temperature, top_p, top_k, repetition_penalty,
system_prompt=sys_p)
front_meta = parse_step1_output(step1_raw)
front_meta_saved = front_meta
full_output += f"✅ **Step 1 Done** — 🌐 Language: **{front_meta['detected_lang']}**\n\n"
full_output += "⏳ **Step 2/2 — Python extract + LLM classify...**\n\n"
yield full_output, full_output
streamer_f, thread_f, mrz_f, python_sections_f = run_step2_structure(
model, processor, front_meta, device,
max_new_tokens, temperature, top_p, top_k, repetition_penalty)
if mrz_f:
all_mrz_data = mrz_f
buffer_f = python_sections_f
yield full_output + buffer_f, full_output + buffer_f
for new_text in streamer_f:
buffer_f += new_text.replace("<|im_end|>", "").replace("<|endoftext|>", "")
time.sleep(0.01)
yield full_output + buffer_f, full_output + buffer_f
full_output += buffer_f + "\n\n"
front_result = buffer_f
thread_f.join()
# ───── BACK CARD ─────
if image_back is not None:
full_output += "\n\n---\n\n# 🎴 BACK CARD\n\n"
full_output += "⏳ **Step 1/2 — Raw OCR (original script, no translation)...**\n\n"
yield full_output, full_output
step1_raw_back = run_step1_extraction(model, processor, image_back, device,
temperature, top_p, top_k, repetition_penalty)
back_meta = parse_step1_output(step1_raw_back)
back_meta_saved = back_meta
full_output += f"✅ **Step 1 Done** — 🌐 Language: **{back_meta['detected_lang']}**\n\n"
full_output += "⏳ **Step 2/2 — Python extract + LLM classify...**\n\n"
yield full_output, full_output
streamer_b, thread_b, mrz_b, python_sections_b = run_step2_structure(
model, processor, back_meta, device,
max_new_tokens, temperature, top_p, top_k, repetition_penalty)
if mrz_b and not all_mrz_data:
all_mrz_data = mrz_b
buffer_b = python_sections_b
yield full_output + buffer_b, full_output + buffer_b
for new_text in streamer_b:
buffer_b += new_text.replace("<|im_end|>", "").replace("<|endoftext|>", "")
time.sleep(0.01)
yield full_output + buffer_b, full_output + buffer_b
full_output += buffer_b
back_result = buffer_b
thread_b.join()
# ───── UNIFIED SUMMARY ─────
if image_front is not None and image_back is not None:
full_output += "\n\n---\n\n"
full_output += build_unified_summary(front_result, back_result, all_mrz_data)
mrz_note = f"MRZ: ✅ {all_mrz_data.get('mrz_format','?')} verified" if all_mrz_data else "MRZ: ❌ Not detected"
full_output += f"\n\n---\n\n**✨ Complete** | Model: `{model_name}` | {mrz_note} | Pipeline: OCR → Python Extract → LLM Classify\n"
yield full_output, full_output
# ╔══════════════════════════════════════════╗
# ║ MODEL CHOICES ║
# ╚══════════════════════════════════════════╝
model_choices = []
if CHHAGAN_V1_AVAILABLE: model_choices.append("Chhagan-ID-OCR-v1 ⭐")
if CHHAGAN_QWEN3_AVAILABLE: model_choices.append("Chhagan-DocVL-Qwen3 🔥")
if CSM_Q4KM_AVAILABLE: model_choices.append("CSM-DocExtract-Q4KM 🏆")
if CSM_4BNB_AVAILABLE: model_choices.append("CSM-DocExtract-4BNB 💎")
if not model_choices: model_choices = ["No models available"]
dual_card_examples = [
["Extract complete information", "examples/5.jpg", None],
["Multilingual OCR with MRZ", "examples/4.jpg", None],
["Extract profile photo and signature", "examples/2.jpg", None],
]
# ╔══════════════════════════════════════════╗
# ║ GRADIO UI ║
# ╚══════════════════════════════════════════╝
demo = gr.Blocks(css=css, theme=steel_blue_theme)
with demo:
gr.Markdown("# 🌍 **CSM Dual-Card ID OCR System**", elem_id="main-title")
gr.Markdown("### *Universal Document Extraction — MRZ + Multilingual + Auto Calendar*")
loaded_models = []
if CHHAGAN_V1_AVAILABLE: loaded_models.append("ID-OCR-v1 ⭐")
if CHHAGAN_QWEN3_AVAILABLE: loaded_models.append("DocVL-Qwen3 🔥")
if CSM_Q4KM_AVAILABLE: loaded_models.append("Q4KM 🏆")
if CSM_4BNB_AVAILABLE: loaded_models.append("4BNB 💎")
model_info = f"**Loaded ({len(loaded_models)}/4):** {', '.join(loaded_models)}" if loaded_models else "⚠️ No models"
gr.Markdown(f"**Status:** {model_info}")
gr.Markdown("**Pipeline:** ✅ Step1: Raw OCR → ✅ Python: MRZ+English Extract → ✅ LLM: Classify+Gaps → ✅ Deduplicate")
with gr.Row():
with gr.Column(scale=2):
image_query = gr.Textbox(
label="💬 Custom Query (Optional)",
placeholder="Leave empty for automatic full extraction...",
value=""
)
gr.Markdown("### 📤 Upload ID Cards")
with gr.Row():
image_front = gr.Image(type="pil", label="🎴 Front Card", height=250)
image_back = gr.Image(type="pil", label="🎴 Back Card (Optional)", height=250)
image_submit = gr.Button("🚀 Extract + Translate + Structure", variant="primary", size="lg")
gr.Examples(
examples=dual_card_examples,
inputs=[image_query, image_front, image_back],
label="📸 Sample ID Cards"
)
with gr.Accordion("⚙️ Advanced Settings", open=False):
max_new_tokens = gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS)
temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6)
top_p = gr.Slider(label="Top-p", minimum=0.05, maximum=1.0, step=0.05, value=0.9)
top_k = gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50)
repetition_penalty= gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.1)
with gr.Column(scale=3):
gr.Markdown("## 📄 Extraction Results", elem_id="output-title")
output = gr.Textbox(label="Raw Output (Streaming)", interactive=True, lines=15)
with gr.Accordion("📝 Structured Preview", open=True):
markdown_output = gr.Markdown(label="Formatted Result")
model_choice = gr.Radio(
choices=model_choices,
label="🤖 Select Model",
value=model_choices[0] if model_choices else None,
info="🏆💎 = 8B Quantized (best) | 🔥 = Qwen3 Fine-tuned | ⭐ = LoRA"
)
with gr.Row(elem_id="gpu-duration-container"):
with gr.Column():
gr.Markdown("**⏱️ GPU Duration (seconds)**")
radioanimated_gpu_duration = RadioAnimated(
choices=["60", "90", "120", "180", "240"],
value="180",
elem_id="radioanimated_gpu_duration"
)
gpu_duration_state = gr.Number(value=180, visible=False)
gr.Markdown("""
**✨ What This Extracts:**
- 🔐 MRZ: TD1/TD3/MRVA/MRVB — Python parsed, 100% accurate
- ✅ English fields: Direct from card, not modified
- 📜 Original script: Arabic/Farsi/Hindi/Chinese as-is
- 🗓️ Calendar: Shamsi/Hijri → Gregorian conversion
- 🔢 Eastern numerals: ۱۲۳ → 123 automatic
- 🔄 Front+Back: Deduplicated, MRZ-verified
""")
radioanimated_gpu_duration.change(
fn=apply_gpu_duration,
inputs=radioanimated_gpu_duration,
outputs=[gpu_duration_state],
api_visibility="private"
)
image_submit.click(
fn=generate_dual_card_ocr,
inputs=[model_choice, image_query, image_front, image_back,
max_new_tokens, temperature, top_p, top_k,
repetition_penalty, gpu_duration_state],
outputs=[output, markdown_output]
)
gr.Markdown("""
---
### 🎯 Feature Matrix
| Feature | Method | Accuracy |
|---------|--------|---------|
| MRZ Parse (TD1/TD3/MRVA) | Python | 100% |
| English Labels Extract | Python Regex | 100% |
| Eastern Numeral Convert | Python char map | 100% |
| Shamsi/Hijri Calendar | Python library | 100% |
| Raw OCR (32+ scripts) | 8B VLM | 90%+ |
| Doc Type Classification | 8B VLM | 95%+ |
| Non-English Translation | 8B VLM | 90%+ |
| Front+Back Deduplication | Python | 100% |
### 📋 Supported Documents
🇮🇳 Aadhaar, PAN, Passport | 🇦🇪 Emirates ID | 🇸🇦 Iqama | 🇴🇲 Oman Resident Card
🌍 International Passports (MRZ) | 🚗 Driving Licences | 🇮🇷 Iranian National ID (Shamsi)
### 🔒 Privacy
All processing on-device | No data stored | GDPR compliant
""")
if __name__ == "__main__":
print("\n🚀 STARTING...")
try:
demo.queue(max_size=50).launch(
server_name="0.0.0.0", server_port=7860, show_error=True, share=False)
except Exception as e:
import traceback
print(f"❌ {e}")
traceback.print_exc()