Update app.py

app.py

@@ -14,7 +14,6 @@ import numpy as np
 from PIL import Image
 import cv2
 
-# Clear cache conflicts
 os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
 os.environ["HF_HOME"] = "/tmp/hf_home"
 
@@ -26,7 +25,6 @@ from transformers import (
     AutoConfig
 )
 
-# PEFT for loading LoRA adapters
 try:
     from peft import PeftModel, PeftConfig
     PEFT_AVAILABLE = True
@@ -34,7 +32,6 @@ except:
     PEFT_AVAILABLE = False
     print("⚠️ PEFT not available, LoRA adapters cannot be loaded")
 
-# Try importing Qwen3VL
 try:
     from transformers import Qwen3VLForConditionalGeneration
     QWEN3_AVAILABLE = True
@@ -46,6 +43,7 @@ from transformers.image_utils import load_image
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
 
+# ===== THEME SETUP =====
 colors.steel_blue = colors.Color(
     name="steel_blue",
     c50="#EBF3F8",
@@ -53,7 +51,7 @@ colors.steel_blue = colors.Color(
     c200="#A8CCE1",
     c300="#7DB3D2",
     c400="#529AC3",
-    c500="#4682B4",  
+    c500="#4682B4",
     c600="#3E72A0",
     c700="#36638C",
     c800="#2E5378",
@@ -111,16 +109,12 @@ class SteelBlueTheme(Soft):
             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;
-}
+#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;
@@ -166,147 +160,108 @@ print("cuda device count:", torch.cuda.device_count())
 if torch.cuda.is_available():
     print("current device:", torch.cuda.current_device())
     print("device name:", torch.cuda.get_device_name(torch.cuda.current_device()))
-
 print("Using device:", device)
 
-# Enhanced multilingual OCR prompt with embedded image extraction
-DUAL_CARD_OCR_PROMPT = """Perform comprehensive OCR extraction on this ID card image. Extract ALL information with maximum English translation accuracy:
-
-EXTRACTION REQUIREMENTS:
-
-1. TEXT EXTRACTION: Extract ALL text in original language with accurate English translation
-2. EMBEDDED IMAGES: 
-   - Locate and describe profile photo/headshot (if present)
-   - Locate and describe signature (if present)
-   - Extract any logos or official seals
-3. MRZ DATA: If Machine Readable Zone is present (usually at bottom):
-   - Extract complete MRZ lines
-   - Parse: Document Type, Country Code, Document Number, Date of Birth, Expiry Date, Nationality
-4. STRUCTURED FIELDS: Extract with English labels:
-   - Full Name (in English)
-   - ID/Document Number
-   - Date of Birth
-   - Issue Date & Expiry Date
-   - Nationality/Country
-   - Address (if present)
-   - Document Type
-
-OUTPUT FORMAT:
-Document Type: [Type: Passport/ID Card/License/etc.]
-
-Embedded Images:
-- Profile Photo Location: [describe position]
-- Profile Photo Description: [describe photo]
-- Signature Present: [Yes/No]
-- Signature Location: [describe position if present]
-
-Original Text:
-[All text in original language with layout preserved]
-
-English Translation:
-[Complete accurate English translation]
-
-Key Fields (English):
-- Full Name: 
-- ID Number: 
-- Date of Birth: 
-- Issue Date: 
-- Expiry Date: 
-- Nationality: 
-- Address: 
-
-MRZ Data (if present):
-Raw MRZ Lines: [extract here]
-Parsed MRZ:
-- Document Type: 
-- Country Code: 
-- Document Number: 
-- Date of Birth: 
-- Expiry Date: 
-- Nationality: 
-- Sex: 
-
-ACCURACY REQUIREMENTS:
-- English translation must be 95%+ accurate
-- Preserve all numbers and dates exactly
-- MRZ must be character-perfect
-- Do not skip any fields"""
-
-SINGLE_SIDE_PROMPT = """Extract all information from this ID card side:
-- All visible text (original + English translation)
-- Profile photo location and description
-- Signature (if present)
-- MRZ data (if present at bottom)
-- All key fields in structured format
-
-Provide complete extraction with high English accuracy."""
 
-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]
+# ===== PROMPTS =====
 
-        uid = uuid.uuid4().hex[:8]
-        group_name = f"ra-{uid}"
+STEP1_EXTRACT_PROMPT = """You are a precision OCR engine. Your ONLY job is to extract raw text from this ID card image.
 
-        inputs_html = "\\n".join(
-            f"""
-            <input class="ra-input" type="radio" name="{group_name}" id="{group_name}-{i}" value="{c}">
-            <label class="ra-label" for="{group_name}-{i}">{c}</label>
-            """
-            for i, c in enumerate(choices)
-        )
+STRICT RULES:
+- Copy ALL text EXACTLY as it appears in original language/script (Hindi, Arabic, Urdu, Chinese, Devanagari, etc.)
+- DO NOT translate anything in this step
+- DO NOT add any interpretation or explanation
+- Preserve layout and line breaks exactly
+- Extract every number, date, code, and character precisely
+- Also detect visual element presence
 
-        html_template = f"""
-        <div class="ra-wrap" data-ra="{uid}">
-          <div class="ra-inner">
-            <div class="ra-highlight"></div>
-            {inputs_html}
-          </div>
-        </div>
-        """
+Output ONLY in this exact structured format, nothing else:
 
-        js_on_load = r"""
-        (() => {
-          const wrap = element.querySelector('.ra-wrap');
-          const inner = element.querySelector('.ra-inner');
-          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 setHighlightByIndex(idx) {
-            const n = choices.length;
-            const pct = 100 / n;
-            highlight.style.width = `calc(${pct}% - 6px)`;
-            highlight.style.transform = `translateX(${idx * 100}%)`;
-          }
-          function setCheckedByValue(val, shouldTrigger=false) {
-            const idx = Math.max(0, choices.indexOf(val));
-            inputs.forEach((inp, i) => { inp.checked = (i === idx); });
-            setHighlightByIndex(idx);
-            props.value = choices[idx];
-            if (shouldTrigger) trigger('change', props.value);
-          }
-          setCheckedByValue(props.value ?? choices[0], false);
-          inputs.forEach((inp) => {
-            inp.addEventListener('change', () => {
-              setCheckedByValue(inp.value, true);
-            });
-          });
-        })();
-        """
+PHOTO_PRESENT: yes/no
+PHOTO_LOCATION: [top-left / top-right / center-left / center-right / bottom-left / not found]
+SIGNATURE_PRESENT: yes/no
+SIGNATURE_LOCATION: [bottom-center / bottom-right / bottom-left / not found]
+MRZ_PRESENT: yes/no
+DETECTED_LANGUAGE: [Hindi / Arabic / Urdu / Chinese / English / Mixed / etc.]
+---TEXT_START---
+[Every piece of text in original script, line by line, layout preserved exactly]
+---TEXT_END---"""
 
-        super().__init__(
-            value=value,
-            html_template=html_template,
-            js_on_load=js_on_load,
-            **kwargs
-        )
 
-def apply_gpu_duration(val: str):
-    return int(val)
+STEP2_TEMPLATE = """You are a multilingual KYC document expert with 95%+ translation accuracy.
+
+DOCUMENT METADATA (from Step 1 analysis):
+- Photo Present: {photo_present} | Location: {photo_location}
+- Signature Present: {sig_present} | Location: {sig_location}
+- MRZ Present: {mrz_present}
+- Detected Language: {detected_lang}
+
+RAW EXTRACTED TEXT (original script):
+{raw_text}
+
+YOUR TASKS:
+1. If text is non-English → translate to English with 95%+ accuracy
+2. If text is already English → copy as-is
+3. Extract all key KYC fields
+4. Output EXACTLY in the format below — no extra commentary
+
+---
+
+## 🖼️ Visual Elements
+
+| Element | Status | Location |
+|---------|--------|----------|
+| 📷 Profile Photo | {photo_present} | {photo_location} |
+| ✍️ Signature | {sig_present} | {sig_location} |
+| 🔐 MRZ Zone | {mrz_present} | Bottom strip |
+
+---
+
+## 📜 Original Script
+
+{raw_text}
+
+---
+
+## 🌐 English Translation
+[Write complete English translation here. If already English, write: Already in English — then copy text]
+
+---
+
+## 🗂️ Key Fields (English)
+
+| Field | Value |
+|-------|-------|
+| 📄 Document Type | |
+| 👤 Full Name | |
+| 🔢 ID / Document Number | |
+| 🎂 Date of Birth | |
+| 📅 Issue Date | |
+| ⏳ Expiry Date | |
+| 🌍 Nationality | |
+| ⚧️ Gender | |
+| 🏠 Address | |
+| 👨 Father / Guardian | |
+| 🏛️ Issuing Authority | |
+
+---
+
+## 🔐 MRZ Data
+[Raw MRZ lines here — if not present write: NOT PRESENT]
+
+**Parsed MRZ:**
+| Field | Value |
+|-------|-------|
+| Document Type | |
+| Country Code | |
+| Document Number | |
+| Date of Birth | |
+| Expiry Date | |
+| Nationality | |
+| Sex | |
+
+---"""
+
 
 # ===== MODEL LOADING =====
 
@@ -314,7 +269,7 @@ print("\n" + "="*70)
 print("🚀 LOADING ALL 4 MODELS")
 print("="*70 + "\n")
 
-# Model 1: Chhagan_ML-VL-OCR-v1 (LoRA Fine-tuned for ID Cards)
+# Model 1: Chhagan_ML-VL-OCR-v1 (LoRA Fine-tuned)
 print("1️⃣ Loading Chhagan_ML-VL-OCR-v1 (LoRA Refined)...")
 MODEL_ID_C1 = "Chhagan005/Chhagan_ML-VL-OCR-v1"
 CHHAGAN_V1_AVAILABLE = False
@@ -330,7 +285,7 @@ if PEFT_AVAILABLE:
         except:
             base_model_id = "Qwen/Qwen2.5-VL-2B-Instruct"
             print(f"   Using default base model: {base_model_id}")
-        
+
         processor_c1 = AutoProcessor.from_pretrained(base_model_id, trust_remote_code=True)
         base_model_c1 = Qwen2VLForConditionalGeneration.from_pretrained(
             base_model_id,
@@ -340,15 +295,14 @@ if PEFT_AVAILABLE:
         )
         model_c1 = PeftModel.from_pretrained(base_model_c1, MODEL_ID_C1)
         model_c1 = model_c1.to(device).eval()
-        
-        print("   ✅ Chhagan_ML-VL-OCR-v1 (Refined) loaded successfully!")
+        print("   ✅ Chhagan_ML-VL-OCR-v1 loaded successfully!")
         CHHAGAN_V1_AVAILABLE = True
     except Exception as e:
         print(f"   ❌ Chhagan_ML-VL-OCR-v1 failed: {e}")
 else:
     print("   ⚠️ PEFT not available, skipping LoRA model")
 
-# Model 2: Chhagan-DocVL-Qwen3 (Qwen3-VL Refined for Documents)
+# Model 2: Chhagan-DocVL-Qwen3
 print("\n2️⃣ Loading Chhagan-DocVL-Qwen3 (Qwen3-VL Refined)...")
 MODEL_ID_C2 = "Chhagan005/Chhagan-DocVL-Qwen3"
 CHHAGAN_QWEN3_AVAILABLE = False
@@ -362,7 +316,6 @@ if QWEN3_AVAILABLE:
                 config = PeftConfig.from_pretrained(MODEL_ID_C2)
                 base_model_id = config.base_model_name_or_path
                 print(f"   Detected as LoRA adapter, base: {base_model_id}")
-                
                 processor_c2 = AutoProcessor.from_pretrained(base_model_id, trust_remote_code=True)
                 base_model_c2 = Qwen3VLForConditionalGeneration.from_pretrained(
                     base_model_id,
@@ -384,15 +337,14 @@ if QWEN3_AVAILABLE:
                 device_map="auto",
                 trust_remote_code=True
             ).to(device).eval()
-        
-        print("   ✅ Chhagan-DocVL-Qwen3 (Refined) loaded successfully!")
+        print("   ✅ Chhagan-DocVL-Qwen3 loaded successfully!")
         CHHAGAN_QWEN3_AVAILABLE = True
     except Exception as e:
         print(f"   ❌ Chhagan-DocVL-Qwen3 failed: {e}")
 else:
     print("   ⚠️ Qwen3VL not available in transformers version")
 
-# Model 3: Qwen3-VL-2B-Instruct (Baseline for Comparison)
+# Model 3: Qwen3-VL-2B-Instruct (Baseline)
 print("\n3️⃣ Loading Qwen3-VL-2B-Instruct (Baseline)...")
 MODEL_ID_Q3 = "Qwen/Qwen3-VL-2B-Instruct"
 QWEN3_BASELINE_AVAILABLE = False
@@ -409,14 +361,14 @@ if QWEN3_AVAILABLE:
             device_map="auto",
             trust_remote_code=True
         ).to(device).eval()
-        print("   ✅ Qwen3-VL-2B-Instruct (Baseline) loaded successfully!")
+        print("   ✅ Qwen3-VL-2B-Instruct loaded successfully!")
         QWEN3_BASELINE_AVAILABLE = True
     except Exception as e:
         print(f"   ❌ Qwen3-VL-2B-Instruct failed: {e}")
 else:
     print("   ⚠️ Qwen3VL not available in transformers version")
 
-# Model 4: Nanonets-OCR2-3B (General OCR Fallback)
+# Model 4: Nanonets-OCR2-3B
 print("\n4️⃣ Loading Nanonets-OCR2-3B (General OCR)...")
 MODEL_ID_V = "nanonets/Nanonets-OCR2-3B"
 NANONETS_AVAILABLE = False
@@ -436,7 +388,6 @@ try:
 except Exception as e:
     print(f"   ❌ Nanonets-OCR2-3B failed: {e}")
 
-# Summary
 print("\n" + "="*70)
 print("📊 MODEL STATUS SUMMARY (4 Models)")
 print("="*70)
@@ -447,14 +398,79 @@ print(f"{'Chhagan-DocVL-Qwen3':<40} {'✅ Loaded' if CHHAGAN_QWEN3_AVAILABLE els
 print(f"{'Qwen3-VL-2B-Instruct':<40} {'✅ Loaded' if QWEN3_BASELINE_AVAILABLE else '❌ Failed':<15} {'Baseline'}")
 print(f"{'Nanonets-OCR2-3B':<40} {'✅ Loaded' if NANONETS_AVAILABLE else '❌ Failed':<15} {'General OCR'}")
 print("="*70)
-
 loaded_count = sum([CHHAGAN_V1_AVAILABLE, CHHAGAN_QWEN3_AVAILABLE, QWEN3_BASELINE_AVAILABLE, NANONETS_AVAILABLE])
 print(f"\n✨ Total models loaded: {loaded_count}/4")
 
-def calc_timeout_duration(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):
-    """Calculate GPU timeout duration - doubled for dual card processing"""
+
+# ===== HELPER: RadioAnimated =====
+
+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"""
+            <input class="ra-input" type="radio" name="{group_name}" id="{group_name}-{i}" value="{c}">
+            <label class="ra-label" for="{group_name}-{i}">{c}</label>
+            """
+            for i, c in enumerate(choices)
+        )
+        html_template = f"""
+        <div class="ra-wrap" data-ra="{uid}">
+          <div class="ra-inner">
+            <div class="ra-highlight"></div>
+            {inputs_html}
+          </div>
+        </div>
+        """
+        js_on_load = r"""
+        (() => {
+          const wrap = element.querySelector('.ra-wrap');
+          const inner = element.querySelector('.ra-inner');
+          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 setHighlightByIndex(idx) {
+            const n = choices.length;
+            const pct = 100 / n;
+            highlight.style.width = `calc(${pct}% - 6px)`;
+            highlight.style.transform = `translateX(${idx * 100}%)`;
+          }
+          function setCheckedByValue(val, shouldTrigger=false) {
+            const idx = Math.max(0, choices.indexOf(val));
+            inputs.forEach((inp, i) => { inp.checked = (i === idx); });
+            setHighlightByIndex(idx);
+            props.value = choices[idx];
+            if (shouldTrigger) trigger('change', props.value);
+          }
+          setCheckedByValue(props.value ?? choices[0], false);
+          inputs.forEach((inp) => {
+            inp.addEventListener('change', () => {
+              setCheckedByValue(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_timeout = int(gpu_timeout)
         if image_front is not None and image_back is not None:
@@ -464,67 +480,195 @@ def calc_timeout_duration(model_name: str, text: str, image_front: Image.Image,
         return 120
 
 
-def extract_embedded_images_info(text_output: str) -> Dict[str, Any]:
-    """Parse extracted text to identify profile photo and signature mentions"""
+# ===== STEP 1: RAW EXTRACTION (NO TRANSLATION) =====
+
+def run_step1_extraction(model, processor, image, device, temperature, top_p, top_k, repetition_penalty):
+    messages = [{
+        "role": "user",
+        "content": [
+            {"type": "image"},
+            {"type": "text", "text": STEP1_EXTRACT_PROMPT},
+        ]
+    }]
+    try:
+        prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    except:
+        prompt = STEP1_EXTRACT_PROMPT
+
+    inputs = processor(
+        text=[prompt],
+        images=[image],
+        return_tensors="pt",
+        padding=True
+    ).to(device)
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=512,
+            do_sample=True,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            repetition_penalty=repetition_penalty,
+        )
+    input_len = inputs['input_ids'].shape[1]
+    generated = output_ids[:, input_len:]
+    return processor.batch_decode(generated, skip_special_tokens=True)[0]
+
+
+# ===== PARSE STEP 1 OUTPUT =====
+
+def parse_step1_output(raw_output: str) -> dict:
     result = {
-        "has_profile_photo": False,
-        "profile_location": "",
-        "has_signature": False,
-        "signature_location": "",
-        "mrz_data": ""
+        "photo_present": "❌ Not detected",
+        "photo_location": "N/A",
+        "sig_present": "❌ Not detected",
+        "sig_location": "N/A",
+        "mrz_present": "❌ Not detected",
+        "detected_lang": "Unknown",
+        "original_text": raw_output
     }
-    
-    if re.search(r"(profile|photo|picture|image|headshot)", text_output, re.IGNORECASE):
-        result["has_profile_photo"] = True
-        photo_match = re.search(r"(top|bottom|left|right|corner|center).{0,20}(photo|image|picture)", text_output, re.IGNORECASE)
-        if photo_match:
-            result["profile_location"] = photo_match.group(0)
-    
-    if re.search(r"signature", text_output, re.IGNORECASE):
-        result["has_signature"] = True
-        sig_match = re.search(r"(signature).{0,50}", text_output, re.IGNORECASE)
-        if sig_match:
-            result["signature_location"] = sig_match.group(0)
-    
-    mrz_pattern = r"^[A-Z][<0-9A-Z]{20,}$"
-    mrz_lines = [line.strip() for line in text_output.split("\n") if re.match(mrz_pattern, line.strip())]
-    if mrz_lines:
-        result["mrz_data"] = "\n".join(mrz_lines)
-    
+
+    def extract_field(pattern, text, default="N/A"):
+        match = re.search(pattern, text, re.IGNORECASE)
+        return match.group(1).strip() if match else default
+
+    photo = extract_field(r"PHOTO_PRESENT:\s*(yes|no)", raw_output)
+    result["photo_present"] = "✅ Yes" if photo.lower() == "yes" else "❌ No"
+    result["photo_location"] = extract_field(r"PHOTO_LOCATION:\s*([^\n]+)", raw_output)
+
+    sig = extract_field(r"SIGNATURE_PRESENT:\s*(yes|no)", raw_output)
+    result["sig_present"] = "✅ Yes" if sig.lower() == "yes" else "❌ No"
+    result["sig_location"] = extract_field(r"SIGNATURE_LOCATION:\s*([^\n]+)", raw_output)
+
+    mrz = extract_field(r"MRZ_PRESENT:\s*(yes|no)", raw_output)
+    result["mrz_present"] = "✅ Yes" if mrz.lower() == "yes" else "❌ No"
+    result["detected_lang"] = extract_field(r"DETECTED_LANGUAGE:\s*([^\n]+)", raw_output, "Unknown")
+
+    text_match = re.search(r"---TEXT_START---\n?(.*?)---TEXT_END---", raw_output, re.DOTALL)
+    if text_match:
+        result["original_text"] = text_match.group(1).strip()
+
     return result
 
 
+# ===== STEP 2: TRANSLATE + STRUCTURE (STREAMING) =====
+
+def run_step2_structure(model, processor, metadata: dict, device,
+                        max_new_tokens, temperature, top_p, top_k, repetition_penalty):
+    step2_prompt = STEP2_TEMPLATE.format(
+        photo_present=metadata["photo_present"],
+        photo_location=metadata["photo_location"],
+        sig_present=metadata["sig_present"],
+        sig_location=metadata["sig_location"],
+        mrz_present=metadata["mrz_present"],
+        detected_lang=metadata["detected_lang"],
+        raw_text=metadata["original_text"],
+    )
+
+    messages = [{"role": "user", "content": [{"type": "text", "text": step2_prompt}]}]
+    try:
+        prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    except:
+        prompt = step2_prompt
+
+    inputs = processor(text=[prompt], return_tensors="pt", padding=True).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()
+    return streamer, thread
+
+
+# ===== UNIFIED DEDUPLICATED SUMMARY =====
+
+def build_unified_summary(front_result: str, back_result: str) -> str:
+    summary = "## 🔄 Unified Deduplicated Record\n\n"
+    summary += "> *Unique fields from both sides merged. Conflicts flagged with ⚠️.*\n\n"
+
+    def extract_table_rows(text):
+        rows = {}
+        table_match = re.search(
+            r"## 🗂️ Key Fields.*?\n\|.*?\n\|[-| ]+\n(.*?)(?=\n---|\Z)", text, re.DOTALL
+        )
+        if table_match:
+            for line in table_match.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()
+                    value = parts[1].strip()
+                    if value and value != "—":
+                        rows[field] = value
+        return rows
+
+    front_fields = extract_table_rows(front_result)
+    back_fields = extract_table_rows(back_result)
+    all_fields = list(dict.fromkeys(list(front_fields.keys()) + list(back_fields.keys())))
+
+    summary += "| Field | Value | Source |\n"
+    summary += "|-------|-------|--------|\n"
+
+    for field in all_fields:
+        f_val = front_fields.get(field, "")
+        b_val = back_fields.get(field, "")
+
+        if f_val and b_val:
+            if f_val.lower() == b_val.lower():
+                summary += f"| {field} | {f_val} | Front + Back ✅ |\n"
+            else:
+                summary += f"| {field} | Front: **{f_val}** / Back: **{b_val}** | ⚠️ Mismatch |\n"
+        elif f_val:
+            summary += f"| {field} | {f_val} | Front only |\n"
+        elif b_val:
+            summary += f"| {field} | {b_val} | Back only |\n"
+
+    return summary + "\n"
+
+
+# ===== MAIN OCR 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,
+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):
-    """
-    Enhanced OCR processing for front and back ID cards with embedded image extraction
-    """
+
+    # Model selection
     if model_name == "Chhagan-ID-OCR-v1 ⭐":
         if not CHHAGAN_V1_AVAILABLE:
             yield "Chhagan_ML-VL-OCR-v1 model is not available.", "Chhagan_ML-VL-OCR-v1 model is not available."
             return
-        processor = processor_c1
-        model = model_c1
+        processor, model = processor_c1, model_c1
+
     elif model_name == "Chhagan-DocVL-Qwen3 🔥":
         if not CHHAGAN_QWEN3_AVAILABLE:
             yield "Chhagan-DocVL-Qwen3 model is not available.", "Chhagan-DocVL-Qwen3 model is not available."
             return
-        processor = processor_c2
-        model = model_c2
+        processor, model = processor_c2, model_c2
+
     elif model_name == "Qwen3-VL-2B (Baseline) 📊":
         if not QWEN3_BASELINE_AVAILABLE:
             yield "Qwen3-VL-2B-Instruct baseline model is not available.", "Qwen3-VL-2B-Instruct baseline model is not available."
             return
-        processor = processor_q3
-        model = model_q3
+        processor, model = processor_q3, model_q3
+
     elif model_name == "Nanonets-OCR2-3B":
         if not NANONETS_AVAILABLE:
             yield "Nanonets-OCR2-3B model is not available.", "Nanonets-OCR2-3B model is not available."
             return
-        processor = processor_v
-        model = model_v
+        processor, model = processor_v, model_v
+
     else:
         yield "Invalid model selected.", "Invalid model selected."
         return
@@ -533,131 +677,85 @@ def generate_dual_card_ocr(model_name: str, text: str, image_front: Image.Image,
         yield "Please upload at least one card image (front or back).", "Please upload at least one card image (front or back)."
         return
 
-    if not text or text.strip().lower() in ["ocr", "extract", "read", ""]:
-        text = DUAL_CARD_OCR_PROMPT
-
     full_output = ""
-    
-    # Process Front Card
+    front_result = ""
+    back_result = ""
+
+    # ===== FRONT CARD =====
     if image_front is not None:
-        full_output += "# 🎴 FRONT CARD EXTRACTION\n\n"
+        full_output += "# 🎴 FRONT CARD\n\n"
+        full_output += "⏳ **Step 1 / 2 — Extracting raw text (original script, no translation)...**\n\n"
         yield full_output, full_output
-        
-        messages_front = [{
-            "role": "user",
-            "content": [
-                {"type": "image"},
-                {"type": "text", "text": text},
-            ]
-        }]
-        
-        try:
-            prompt_front = processor.apply_chat_template(messages_front, tokenize=False, add_generation_prompt=True)
-        except:
-            prompt_front = text
-
-        inputs_front = processor(
-            text=[prompt_front],
-            images=[image_front],
-            return_tensors="pt",
-            padding=True).to(device)
-
-        streamer_front = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
-        generation_kwargs_front = {
-            **inputs_front,
-            "streamer": streamer_front,
-            "max_new_tokens": max_new_tokens,
-            "do_sample": True,
-            "temperature": temperature,
-            "top_p": top_p,
-            "top_k": top_k,
-            "repetition_penalty": repetition_penalty,
-        }
-        thread_front = Thread(target=model.generate, kwargs=generation_kwargs_front)
-        thread_front.start()
-        
-        buffer_front = ""
-        for new_text in streamer_front:
-            buffer_front += new_text
-            buffer_front = buffer_front.replace("<|im_end|>", "").replace("<|endoftext|>", "")
+
+        step1_raw = run_step1_extraction(
+            model, processor, image_front, device,
+            temperature, top_p, top_k, repetition_penalty
+        )
+        front_meta = parse_step1_output(step1_raw)
+
+        full_output += f"✅ **Step 1 Complete** — 🌐 Detected Language: **{front_meta['detected_lang']}**\n\n"
+        full_output += "⏳ **Step 2 / 2 — Translating to English & building structured output...**\n\n"
+        yield full_output, full_output
+
+        streamer_f, thread_f = run_step2_structure(
+            model, processor, front_meta, device,
+            max_new_tokens, temperature, top_p, top_k, repetition_penalty
+        )
+
+        buffer_f = ""
+        for new_text in streamer_f:
+            buffer_f += new_text
+            buffer_f = buffer_f.replace("<|im_end|>", "").replace("<|endoftext|>", "")
             time.sleep(0.01)
-            current_output = full_output + buffer_front
-            yield current_output, current_output
-        
-        full_output += buffer_front + "\n\n"
-        thread_front.join()
-    
-    # Process Back Card
+            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 EXTRACTION\n\n"
+        full_output += "\n\n---\n\n# 🎴 BACK CARD\n\n"
+        full_output += "⏳ **Step 1 / 2 — Extracting raw text (original script, no translation)...**\n\n"
         yield full_output, full_output
-        
-        messages_back = [{
-            "role": "user",
-            "content": [
-                {"type": "image"},
-                {"type": "text", "text": text},
-            ]
-        }]
-        
-        try:
-            prompt_back = processor.apply_chat_template(messages_back, tokenize=False, add_generation_prompt=True)
-        except:
-            prompt_back = text
-
-        inputs_back = processor(
-            text=[prompt_back],
-            images=[image_back],
-            return_tensors="pt",
-            padding=True).to(device)
-
-        streamer_back = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
-        generation_kwargs_back = {
-            **inputs_back,
-            "streamer": streamer_back,
-            "max_new_tokens": max_new_tokens,
-            "do_sample": True,
-            "temperature": temperature,
-            "top_p": top_p,
-            "top_k": top_k,
-            "repetition_penalty": repetition_penalty,
-        }
-        thread_back = Thread(target=model.generate, kwargs=generation_kwargs_back)
-        thread_back.start()
-        
-        buffer_back = ""
-        for new_text in streamer_back:
-            buffer_back += new_text
-            buffer_back = buffer_back.replace("<|im_end|>", "").replace("<|endoftext|>", "")
+
+        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)
+
+        full_output += f"✅ **Step 1 Complete** — 🌐 Detected Language: **{back_meta['detected_lang']}**\n\n"
+        full_output += "⏳ **Step 2 / 2 — Translating to English & building structured output...**\n\n"
+        yield full_output, full_output
+
+        streamer_b, thread_b = run_step2_structure(
+            model, processor, back_meta, device,
+            max_new_tokens, temperature, top_p, top_k, repetition_penalty
+        )
+
+        buffer_b = ""
+        for new_text in streamer_b:
+            buffer_b += new_text
+            buffer_b = buffer_b.replace("<|im_end|>", "").replace("<|endoftext|>", "")
             time.sleep(0.01)
-            current_output = full_output + buffer_back
-            yield current_output, current_output
-        
-        full_output += buffer_back
-        thread_back.join()
-    
-    # Add summary section
-    full_output += "\n\n---\n\n## 📊 Extraction Summary\n"
-    
-    embedded_info = extract_embedded_images_info(full_output)
-    
-    full_output += f"\n### 🖼️ Embedded Content Detection:\n"
-    full_output += f"- **Profile Photo**: {'✅ Detected' if embedded_info['has_profile_photo'] else '❌ Not found'}\n"
-    if embedded_info['profile_location']:
-        full_output += f"  - Location: {embedded_info['profile_location']}\n"
-    full_output += f"- **Signature**: {'✅ Detected' if embedded_info['has_signature'] else '❌ Not found'}\n"
-    if embedded_info['signature_location']:
-        full_output += f"  - Details: {embedded_info['signature_location']}\n"
-    
-    if embedded_info['mrz_data']:
-        full_output += f"\n### 🔐 MRZ Data Extracted:\n```\n{embedded_info['mrz_data']}\n```\n"
-    
-    full_output += f"\n**✨ Extraction Complete** | Model: {model_name}\n"
-    
+            yield full_output + buffer_b, full_output + buffer_b
+
+        full_output += buffer_b
+        back_result = buffer_b
+        thread_b.join()
+
+    # ===== UNIFIED SUMMARY (only when both sides uploaded) =====
+    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)
+
+    full_output += f"\n\n---\n\n**✨ Extraction Complete** | Model: `{model_name}` | Pipeline: OCR → Language Detect → Translate → Structure\n"
     yield full_output, full_output
 
 
-# Build model choices dynamically
+# ===== BUILD MODEL CHOICES =====
+
 model_choices = []
 if CHHAGAN_V1_AVAILABLE:
     model_choices.append("Chhagan-ID-OCR-v1 ⭐")
@@ -671,18 +769,20 @@ if NANONETS_AVAILABLE:
 if not model_choices:
     model_choices = ["No models available"]
 
-# Example images
 dual_card_examples = [
     ["Extract complete information from both sides", "examples/5.jpg", None],
     ["Multilingual OCR with MRZ extraction", "examples/4.jpg", None],
     ["Extract profile photo and signature locations", "examples/2.jpg", None],
 ]
 
+
+# ===== GRADIO UI =====
+
 demo = gr.Blocks(css=css, theme=steel_blue_theme)
 with demo:
     gr.Markdown("# 🌍 **Chhagan Dual-Card ID OCR System**", elem_id="main-title")
-    gr.Markdown("### *Advanced OCR with Profile Image, Signature & MRZ Extraction*")
-    
+    gr.Markdown("### *Advanced OCR • Auto Language Detection • English Translation • MRZ Parsing*")
+
     loaded_models = []
     if CHHAGAN_V1_AVAILABLE:
         loaded_models.append("ID-OCR-v1 ⭐")
@@ -692,53 +792,52 @@ with demo:
         loaded_models.append("Qwen3-Baseline 📊")
     if NANONETS_AVAILABLE:
         loaded_models.append("Nanonets")
-    
+
     model_info = f"**Loaded Models ({len(loaded_models)}/4):** {', '.join(loaded_models)}" if loaded_models else "⚠️ No models loaded"
-    
     gr.Markdown(f"**Status:** {model_info}")
-    gr.Markdown("**Features**: ✅ Dual Card Upload | ✅ Profile Photo Detection | ✅ Signature Extraction | ✅ MRZ Reading | ✅ 95%+ English Accuracy")
-    
+    gr.Markdown("**Pipeline:** ✅ Step 1: Raw OCR (original script) → ✅ Step 2: Auto Translate to English → ✅ Structured Output → ✅ Front+Back Deduplication")
+
     with gr.Row():
         with gr.Column(scale=2):
             image_query = gr.Textbox(
-                label="💬 Custom Query (Optional)", 
-                placeholder="Leave empty for automatic extraction of all data including images and MRZ...",
+                label="💬 Custom Query (Optional)",
+                placeholder="Leave empty for automatic full extraction (OCR + translate + structure)...",
                 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 Complete Data", variant="primary", size="lg")
-            
+            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 (nucleus sampling)", 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("## 📄 Complete Extraction Results", elem_id="output-title")
-            output = gr.Textbox(label="OCR Output (Streaming)", interactive=True, lines=15)
-            with gr.Accordion("📝 Markdown Preview", open=True):
+            gr.Markdown("## 📄 Extraction Results", elem_id="output-title")
+            output = gr.Textbox(label="Raw Output (Streaming)", interactive=True, lines=15)
+            with gr.Accordion("📝 Markdown Preview (Structured)", open=True):
                 markdown_output = gr.Markdown(label="Formatted Result")
 
             model_choice = gr.Radio(
                 choices=model_choices,
                 label="🤖 Select OCR Model",
                 value=model_choices[0] if model_choices else None,
-                info="⭐🔥 = Refined for high accuracy | 📊 = Baseline"
+                info="⭐🔥 = Fine-tuned for ID Cards | 📊 = Baseline | General OCR = Nanonets"
             )
-            
+
             with gr.Row(elem_id="gpu-duration-container"):
                 with gr.Column():
                     gr.Markdown("**⏱️ GPU Duration (seconds)**")
@@ -748,62 +847,72 @@ with demo:
                         elem_id="radioanimated_gpu_duration"
                     )
                     gpu_duration_state = gr.Number(value=120, visible=False)
-            
+
             gr.Markdown("""
-            **✨ Extraction Includes:**
-            - 📝 Complete text extraction (original + English)
-            - 🖼️ Profile photo location detection
-            - ✍️ Signature identification
-            - 🔐 MRZ data parsing
-            - 🎯 Structured key fields
+            **✨ What This Extracts:**
+            - 📜 Original script (Hindi, Arabic, Urdu, Chinese, etc.)
+            - 🌐 Auto English translation (95%+ accuracy)
+            - 🖼️ Profile photo location & description
+            - ✍️ Signature detection & location
+            - 🔐 MRZ raw lines + parsed fields
+            - 🗂️ Structured key fields (Name, DOB, ID No., etc.)
+            - 🔄 Front + Back unified deduplicated record
             """)
 
     radioanimated_gpu_duration.change(
-        fn=apply_gpu_duration, 
-        inputs=radioanimated_gpu_duration, 
-        outputs=[gpu_duration_state], 
+        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],
+        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("""
     ---
-    ### 🎯 Key Features
-    
+    ### 🎯 Feature Matrix
+
     | Feature | Status | Description |
     |---------|--------|-------------|
-    | **Dual Card Upload** | ✅ | Process front and back simultaneously |
-    | **Profile Photo Detection** | ✅ | Automatically locates and describes headshot |
-    | **Signature Extraction** | ✅ | Identifies signature presence and location |
-    | **MRZ Reading** | ✅ | Parses Machine Readable Zone data |
-    | **English Translation** | ✅ | 95%+ accuracy for non-English text |
-    | **Multilingual Support** | ✅ | 30+ languages including Arabic, Hindi, Urdu |
-    
-    ### 📋 Supported Document Types
-    - Government ID Cards (front + back)
-    - Passports (with MRZ)
-    - Driver's Licenses
-    - Residence Permits
-    - Visas and Travel Documents
-    
-    ### 🔒 Privacy & Security
-    - All processing on-device
+    | **Two-Step Pipeline** | ✅ | Step 1 = Raw OCR, Step 2 = Translate + Structure |
+    | **Auto Language Detect** | ✅ | Hindi, Arabic, Urdu, Chinese, 30+ languages |
+    | **English Translation** | ✅ | 95%+ accuracy, only when non-English detected |
+    | **Original Script Preserved** | ✅ | Both original + translated shown side by side |
+    | **Profile Photo Detection** | ✅ | Location described in visual elements box |
+    | **Signature Extraction** | ✅ | Detected and located per card side |
+    | **MRZ Parsing** | ✅ | Raw lines + structured parsed fields |
+    | **Dual Card Deduplication** | ✅ | Front + Back merged, mismatches flagged ⚠️ |
+    | **Markdown Structured Output** | ✅ | Tables, code blocks, section headers |
+
+    ### 📋 Supported Documents
+    - 🇮🇳 Aadhaar Card, PAN Card, Voter ID
+    - 🌍 International Passports (with MRZ)
+    - 🪪 Driver's Licenses
+    - 🏛️ Government ID Cards (30+ countries)
+    - 📋 Residence Permits & Visas
+
+    ### 🔒 Privacy
+    - All processing on-device (GPU)
     - No data stored or transmitted
     - GDPR compliant
-    
-    **💡 Pro Tip**: Upload both front and back for complete extraction including hidden MRZ data on back side!
+
+    **💡 Pro Tip**: Upload both front and back for full deduplication and MRZ cross-validation!
     """)
 
+
 if __name__ == "__main__":
     print("\n" + "="*70)
     print("🚀 STARTING GRADIO INTERFACE...")
     print("="*70 + "\n")
-    
     try:
         demo.queue(max_size=50).launch(
             server_name="0.0.0.0",
@@ -816,3 +925,5 @@ if __name__ == "__main__":
         print(f"❌ Launch error: {e}")
         import traceback
         traceback.print_exc()
+
+