Create app.py

app.py

@@ -0,0 +1,325 @@
+# app.py
+
+import os
+import warnings
+import numpy as np
+import cv2
+import onnxruntime as ort
+from insightface.app import FaceAnalysis
+from PIL import Image
+import gradio as gr
+import time
+import sys
+from typing import Optional, Tuple, Any
+
+# Suppress warnings for a cleaner output
+warnings.filterwarnings("ignore")
+
+# --- Configuration: Model Paths and Thresholds ---
+# NOTE: The user must ensure these files are downloaded to the same directory
+# before running the script (e.g., using a separate gdown script or Colab cell).
+TARGET_DIR = '.' # Current directory for a local setup
+
+# Deepfake Model Paths
+MODEL_PATHS = {
+    "mobilenetv3": os.path.join(TARGET_DIR, "mobilenetv3_small_100_final.onnx"),
+    "efficientnet_b0": os.path.join(TARGET_DIR, "efficientnet_b0_final.onnx"),
+    "edgenext": os.path.join(TARGET_DIR, "edgenext_small_final.onnx"),
+}
+
+# --- 2A. CONFIGURATION ---
+# We use the built-in insightface model 'buffalo_l' for better alignment/embedding
+SIM_MODEL_NAME = 'buffalo_l'
+CTX_ID = -1  # CPU
+ID_MATCH_THRESHOLD = 0.50  # Similarity Threshold
+FAKE_SCORE_THRESHOLD = 0.5 # Deepfake Score Threshold
+
+# Fixed paths for shared Deepfake models (Caffe files are no longer used for detection,
+# as we switch to the more robust insightface detector)
+# NOTE: We keep the ONNX paths as they are the deepfake models.
+ONNX_SESSIONS = {}
+
+# --- 2B. INITIALIZE MODELS ---
+app: Optional[FaceAnalysis] = None
+
+print("\n--- 1. Initializing Models ---")
+try:
+    # Use 'buffalo_l' for higher accuracy, which also provides the 5-point landmarks
+    app = FaceAnalysis(name=SIM_MODEL_NAME, providers=['CPUExecutionProvider'])
+    # Customize the detector to get landmarks ('lmk') along with the usual attributes
+    app.prepare(ctx_id=CTX_ID, det_size=(640, 640), det_thresh=0.5,
+                det_model="retinaface_r50_v1", allowed_modules=['detection', 'landmark', 'recognition'])
+
+    # Initialize ONNX Deepfake Classification Models
+    for model_name, path in MODEL_PATHS.items():
+        if os.path.exists(path):
+            ONNX_SESSIONS[model_name] = ort.InferenceSession(path, providers=['CPUExecutionProvider'])
+            print(f"Loaded {model_name.upper()} model.")
+        else:
+            print(f"Warning: Deepfake model {model_name.upper()} not found at {path}")
+
+    if not ONNX_SESSIONS:
+        raise FileNotFoundError("No ONNX deepfake models could be loaded.")
+
+except Exception as e:
+    print(f"❌ FATAL ERROR: Failed to load models. Detail: {e}")
+    app = None
+    sys.exit(1) # Exit if models fail to load
+
+print("✅ Model initialization complete.")
+
+# --- 2C. IDENTITY VERIFICATION HELPER FUNCTIONS (Similarity Check) ---
+
+def get_largest_face(faces: list) -> Optional[Any]:
+    """Returns the largest face detected."""
+    if not faces: return None
+    def get_area(face):
+        bbox = face.bbox.astype(np.int32)
+        return (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
+    return max(faces, key=get_area)
+
+def get_face_data(img_array_rgb: np.ndarray) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]:
+    """Extracts embedding, landmarks (5-point), BGR image, and bbox."""
+    if app is None: return None, None, None, None
+    img_bgr = cv2.cvtColor(img_array_rgb, cv2.COLOR_RGB2BGR)
+    all_faces = app.get(img_bgr)
+    if not all_faces:
+        return None, None, img_bgr, None
+    face = get_largest_face(all_faces)
+    # The 'face' object from FaceAnalysis now contains embedding and 5-point landmarks ('lmk')
+    return face.embedding, face.lmk, img_bgr, face.bbox
+
+def calculate_similarity(embedding1: Optional[np.ndarray], embedding2: Optional[np.ndarray]) -> float:
+    """Calculates cosine similarity."""
+    if embedding1 is None or embedding2 is None: return 0.0
+    # Ensure embeddings are normalized before dot product for true cosine sim
+    e1_norm = embedding1 / np.linalg.norm(embedding1)
+    e2_norm = embedding2 / np.linalg.norm(embedding2)
+    similarity = np.dot(e1_norm, e2_norm)
+    return float(similarity)
+
+# --- 2D. DEEPFAKE DETECTION HELPER FUNCTIONS (Liveness Check) ---
+
+def align_face_insightface(img_bgr: np.ndarray, landmarks_5pt: np.ndarray, output_size: int = 160) -> np.ndarray:
+    """
+    Simplified alignment using 5-point landmarks for deepfake model input.
+    This replaces the complex dlib/81-point alignment.
+    The goal is a centered, roughly aligned 160x160 face crop.
+    """
+    # Standard 5-point template for InsightFace's alignment for deepfake models
+    dst = np.array([
+        [30.2946, 51.6963],  # Left Eye
+        [65.5318, 51.6963],  # Right Eye
+        [48.0252, 71.7366],  # Nose Tip
+        [33.5493, 92.3655],  # Left Mouth Corner
+        [62.7299, 92.3655]   # Right Mouth Corner
+    ], dtype=np.float32) * (output_size / 96) # Scale for 160x160 output
+
+    src = landmarks_5pt.astype(np.float32)
+
+    # Calculate similarity transform matrix
+    M, _ = cv2.estimateAffinePartial2D(src, dst, method=cv2.LMEDS)
+    
+    # Apply transformation
+    aligned_face = cv2.warpAffine(img_bgr, M, (output_size, output_size), flags=cv2.INTER_CUBIC)
+    
+    return aligned_face
+
+def get_liveness_score(img_array_rgb: np.ndarray, landmarks_5pt: np.ndarray, model_choice: str) -> float:
+    """Extracts aligned face and runs deepfake inference using the selected model."""
+    if model_choice not in ONNX_SESSIONS: return 0.0
+    try:
+        session = ONNX_SESSIONS[model_choice]
+        img_bgr = cv2.cvtColor(img_array_rgb, cv2.COLOR_RGB2BGR)
+
+        # Align the face using the 5-point landmarks
+        face_crop_bgr = align_face_insightface(img_bgr, landmarks_5pt, output_size=160)
+
+        # Pre-process for ONNX model
+        face_crop_rgb = cv2.cvtColor(face_crop_bgr, cv2.COLOR_BGR2RGB)
+        # Assuming the deepfake model expects a [-1, 1] normalization
+        normalized_img = (face_crop_rgb / 255.0 - 0.5) / 0.5
+        input_tensor = np.transpose(normalized_img, (2, 0, 1)) # C, H, W
+        input_tensor = np.expand_dims(input_tensor, axis=0).astype("float32") # 1, C, H, W
+
+        input_name = session.get_inputs()[0].name
+        output_name = session.get_outputs()[0].name
+        logit = session.run([output_name], {input_name: input_tensor})[0]
+
+        # Convert logit to probability (Fake Confidence Score) using sigmoid: 1 / (1 + exp(-x))
+        probability = 1 / (1 + np.exp(-logit))
+        score = float(np.ravel(probability)[0]) if probability.size > 0 else 0.0
+
+        return score
+
+    except Exception as e:
+        print(f"Liveness check failed: {e}", file=sys.stderr)
+        return 0.0 # Fail safe
+
+# ==============================================================================
+# 3. UNIFIED eKYC ANALYSIS AND GRADIO INTERFACE
+# ==============================================================================
+
+def unified_ekyc_analysis(model_choice: str, img_A_pil: Image.Image, img_B_pil: Image.Image) -> Tuple[Optional[Image.Image], Optional[Image.Image], str]:
+    """
+    Performs Identity Verification (Step 1) then Forgery Check (Step 2).
+    """
+    if app is None or not ONNX_SESSIONS or model_choice not in ONNX_SESSIONS:
+        error_msg = f"""# ❌ CRITICAL FAILURE: Models failed to load. Please ensure all ONNX models are in the current directory and dependencies are installed."""
+        return None, None, error_msg
+
+    start_time = time.time()
+
+    # --- Data Prep and Feature Extraction ---
+    img_A_array = np.array(img_A_pil.convert('RGB'))
+    img_B_array = np.array(img_B_pil.convert('RGB'))
+
+    # e1, e2: Face Embeddings for Similarity
+    # lmk_A, lmk_B: 5-point Landmarks for Alignment
+    # vis_A_bgr, vis_B_bgr: BGR images for visualization
+    # bbox_A, bbox_B: Bounding Boxes
+    e1, lmk_A, vis_A_bgr, bbox_A = get_face_data(img_A_array)
+    e2, lmk_B, vis_B_bgr, bbox_B = get_face_data(img_B_array)
+
+    # 1. Basic Face Detection Check (Pre-Step)
+    if e1 is None or e2 is None:
+        report = "🛑 **PRE-CHECK FAILED:** Face detection failed on one or both images. Cannot proceed."
+        return Image.fromarray(img_A_array), Image.fromarray(img_B_array), report
+
+
+    # --- STEP 1: IDENTITY VERIFICATION (Similarity Check) ---
+    match_score = calculate_similarity(e1, e2)
+    match_is_success = match_score > ID_MATCH_THRESHOLD
+
+    liveness_score_A = 0.0
+    liveness_score_B = 0.0
+    is_acceptance_ok = False
+    rejection_reason = None
+
+    if not match_is_success:
+        # Stop here and REJECT
+        rejection_reason = "Identity Mismatch (Step 1 Failed)"
+
+    else:
+        # --- STEP 2: FORGERY/LIVENESS CHECK (Conditional) ---
+        liveness_score_A = get_liveness_score(img_A_array, lmk_A, model_choice)
+        liveness_score_B = get_liveness_score(img_B_array, lmk_B, model_choice)
+
+        # FAKE score <= 0.50 means it is classified as REAL
+        liveness_is_real_A = liveness_score_A <= FAKE_SCORE_THRESHOLD
+        liveness_is_real_B = liveness_score_B <= FAKE_SCORE_THRESHOLD
+
+        is_acceptance_ok = liveness_is_real_A and liveness_is_real_B
+        rejection_reason = "Forgery Detected (Step 2 Failed)" if not is_acceptance_ok else "All checks passed."
+
+    # --- Final Labeling and Visualization ---
+    final_label_text = "✅ **ACCEPT IDENTITY**" if is_acceptance_ok else "❌ **REJECT IDENTITY**"
+    final_label_color = "#28a745" if is_acceptance_ok else "#dc3545"
+
+    # Set BBox Color: Red for any rejection, Green for full acceptance
+    bbox_color = (0, 255, 0) if is_acceptance_ok else (0, 0, 255) # Green (BGR) or Red (BGR)
+
+    # Drawing BBoxes
+    bbox_A_int = bbox_A.astype(np.int32)
+    cv2.rectangle(vis_A_bgr, (bbox_A_int[0], bbox_A_int[1]), (bbox_A_int[2], bbox_A_int[3]), bbox_color, 4)
+    bbox_B_int = bbox_B.astype(np.int32)
+    cv2.rectangle(vis_B_bgr, (bbox_B_int[0], bbox_B_int[1]), (bbox_B_int[2], bbox_B_int[3]), bbox_color, 4)
+
+    end_time = time.time()
+
+    # --- Final Report Generation ---
+
+    match_status_text = "MATCH" if match_is_success else "MISMATCH"
+    liveness_is_real_A = liveness_score_A <= FAKE_SCORE_THRESHOLD
+    liveness_is_real_B = liveness_score_B <= FAKE_SCORE_THRESHOLD
+    liveness_A_text = "REAL" if liveness_is_real_A else "FAKE"
+    liveness_B_text = "REAL" if liveness_is_real_B else "FAKE"
+
+
+    final_report = f"""
+    ## 🏦 ZenTej eKYC Verification Report
+
+    <div style='border: 2px solid {final_label_color}; padding: 10px; border-radius: 5px; background-color: {'#e6ffe6' if is_acceptance_ok else '#ffe6e6'};'>
+    <h3 style='color: {final_label_color}; margin-top: 0;'>{final_label_text}</h3>
+    <p><strong>Reason:</strong> {rejection_reason}</p>
+    </div>
+
+    ---
+
+    ### Step 1: Identity Similarity Check
+
+    | Metric | Value | Status |
+    | :--- | :--- | :--- |
+    | **Match Score (Cosine)** | `{match_score:.4f}` | **<span style='color:{'#28a745' if match_is_success else '#dc3545'};'>{match_status_text}</span>** |
+    | **Required Threshold** | $>{ID_MATCH_THRESHOLD}$ | |
+
+    <details>
+      <summary style="font-weight: bold;">Metrics Interpretation</summary>
+      <p>The score is a **Cosine Similarity** value between 0.0 and 1.0. The system requires the score to be **greater than {ID_MATCH_THRESHOLD:.2f}** to proceed to the next step.</p>
+    </details>
+
+    ---
+
+    ### Step 2: Forgery / Liveness Check (Conditional)
+
+    (Run only if Step 1 is MATCH)
+
+    | Image | FAKE Confidence Score | Liveness Status |
+    | :--- | :--- | :--- |
+    | **Input 1 (Live)** | `{liveness_score_A:.4f}` | **<span style='color:{'#28a745' if liveness_is_real_A else '#dc3545'};'>{liveness_A_text}</span>** |
+    | **Input 2 (Document)** | `{liveness_score_B:.4f}` | **<span style='color:{'#28a745' if liveness_is_real_B else '#dc3545'};'>{liveness_B_text}</span>** |
+
+    <details>
+      <summary style="font-weight: bold;">Metrics Interpretation</summary>
+      <p>The score is the model's **Confidence that the image is a FAKE** (range: 0.0 to 1.0). A score greater than {FAKE_SCORE_THRESHOLD:.2f} results in a **FAKE** classification. Both images must be classified as **REAL** (score $\le {FAKE_SCORE_THRESHOLD:.2f}$) for Step 2 to pass.</p>
+      <p style="font-size: 0.9em;">*Model Used: {model_choice.upper()}*</p>
+    </details>
+
+    ---
+
+    *Total Processing Time: {end_time - start_time:.3f} seconds*
+    """
+
+    # Prepare images for Gradio (convert BGR to RGB)
+    vis_A_rgb = cv2.cvtColor(vis_A_bgr, cv2.COLOR_BGR2RGB)
+    vis_B_rgb = cv2.cvtColor(vis_B_bgr, cv2.COLOR_BGR2RGB)
+
+    return Image.fromarray(vis_A_rgb), Image.fromarray(vis_B_rgb), final_report
+
+
+# --- GRADIO FRONTEND ---
+print("\n--- 2. Initializing Gradio interface ---")
+
+available_models = list(ONNX_SESSIONS.keys())
+# Prioritize EdgeNeXt if available
+default_model = "edgenext" if "edgenext" in available_models else ("efficientnet_b0" if "efficientnet_b0" in available_models else (available_models[0] if available_models else ""))
+
+if not available_models:
+    print("FATAL: No deepfake models were loaded. Cannot launch Gradio.")
+else:
+    iface = gr.Interface(
+        fn=unified_ekyc_analysis,
+        inputs=[
+            gr.Dropdown(
+                label="Select Deepfake Model",
+                choices=available_models,
+                value=default_model,
+                interactive=True
+            ),
+            gr.Image(
+                type="pil",
+                label="Input 1: Live Selfie (Checked for Liveness)",
+                sources=['upload', 'webcam']
+            ),
+            gr.Image(type="pil", label="Input 2: Document Photo (Checked for Identity)")
+        ],
+        outputs=[
+            gr.Image(label="Input 1 Face (Final Status BBox)", height=256, width=256),
+            gr.Image(label="Input 2 Face (Final Status BBox)", height=256, width=256),
+            gr.Markdown(label="Final eKYC Report")
+        ],
+        title="Deepfake-Proof eKYC System (Unified Analysis)",
+        description="Performs two-step conditional verification: Step 1: Identity Match. Step 2 (if Match): Forgery/Liveness Check on both images. NOTE: Ensure all ONNX models are in the same directory.",
+    )
+
+    iface.launch(debug=False)