Update app.py

app.py

@@ -1,4 +1,6 @@
-# app.py
+# Combined Script for Colab/Notebook Execution
+# This script handles installation, model download via gdown, 
+# and runs the dlib-free eKYC application with the buffalo_l model.
 
 import os
 import warnings
@@ -15,46 +17,102 @@ 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
+# ==============================================================================
+# 1. SETUP, INSTALLATION, AND MODEL DOWNLOAD (Combined Cell)
+# ==============================================================================
+
+print("--- 1. Installing Required Libraries ---")
+# Install core libraries and gdown, gradio. dlib is removed from the requirements
+!pip install insightface==0.7.3 numpy onnxruntime opencv-python matplotlib tqdm gdown gradio --quiet
 
-# Deepfake Model Paths
+# --- Configuration: Model File IDs and Target Paths ---
+TARGET_DIR = '/content/'
+
+# Deepfake Model Paths (All ONNX models are downloaded for runtime switching)
 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"),
+    "mobilenetv3": "/content/mobilenetv3_small_100_final.onnx",
+    "efficientnet_b0": "/content/efficientnet_b0_final.onnx",
+    "edgenext": "/content/edgenext_small_final.onnx",
+}
+
+# Mapping of file names to their corresponding Google Drive File IDs
+# NOTE: The DLIB and Caffe files are kept for download continuity, 
+# but are NOT used in the dlib-free logic.
+MODEL_FILES = {
+    # Deepfake Detector Components (Downloaded but not used in final logic)
+    "deploy.prototxt": "1V02QA7eOnrkKixTdnP6cvIBx4Qxqwhmw",
+    "res10_300x300_ssd_iter_140000_fp16.caffemodel": "14n7DryxHqwqac9z0HzpIqtipBp5EfRvA",
+    "shape_predictor_81_face_landmarks.dat": "1sixwbA4oOn7Ijmm85sAODL8AtwjCq6a9",
+
+    # Deepfake Classification ONNX Models
+    "mobilenetv3_small_100_final.onnx": "1spFbTIL8nRmIBG_F6j6-aF01fWGVGo_f",
+    "efficientnet_b0_final.onnx": "1TsHUbx0cd-55XDygQIAmEbXFUGHxBT_x",
+    "edgenext_small_final.onnx": "15hnhznZVyASYhSOYOFSsgMGEfsyh1MBY"
 }
 
+def download_models_from_drive():
+    """Downloads all required model files from the provided Drive IDs."""
+    print(f"\n--- 2. Starting Deepfake Model Download to {TARGET_DIR} ---")
+
+    try:
+        import gdown
+    except ImportError:
+        # Should be installed by pip above, but double check
+        !pip install gdown --quiet
+        import gdown
+
+    os.makedirs(TARGET_DIR, exist_ok=True)
+    downloaded_files = 0
+
+    for filename, file_id in MODEL_FILES.items():
+        local_path = os.path.join(TARGET_DIR, filename)
+
+        if os.path.exists(local_path) and os.path.getsize(local_path) > 0:
+            # print(f"Skipping download, {filename} already exists.")
+            downloaded_files += 1
+            continue
+
+        try:
+            print(f"Downloading {filename}...")
+            gdown.download(id=file_id, output=local_path, quiet=True, fuzzy=True)
+            if os.path.exists(local_path) and os.path.getsize(local_path) > 0:
+                downloaded_files += 1
+        except Exception as e:
+            print(f"Warning: Failed to download {filename}. Error: {e}", file=sys.stderr)
+
+download_models_from_drive()
+print("✅ Initial setup complete. Proceeding to model initialization.")
+
+# --- 2. MODEL INITIALIZATION AND CORE HELPER FUNCTIONS ---
+
 # --- 2A. CONFIGURATION ---
-# We use the built-in insightface model 'buffalo_l' for better alignment/embedding
-SIM_MODEL_NAME = 'buffalo_l'
+# IMPORTANT: Using 'buffalo_l' as requested, which provides 5-point landmarks (lmk)
+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 = {}
+app: Optional[FaceAnalysis] = None
 
 # --- 2B. INITIALIZE MODELS ---
-app: Optional[FaceAnalysis] = None
 
-print("\n--- 1. Initializing Models ---")
+print("\n--- 3. Initializing Models ---")
 try:
-    # Use 'buffalo_l' for higher accuracy, which also provides the 5-point landmarks
+    # Face Analysis/Recognition Model (downloads buffalo_l)
+    print(f"Initializing FaceAnalysis model: {SIM_MODEL_NAME}")
     app = FaceAnalysis(name=SIM_MODEL_NAME, providers=['CPUExecutionProvider'])
-    # Customize the detector to get landmarks ('lmk') along with the usual attributes
+    
+    # Configure app to get detection, 5-point landmark (lmk), and recognition (embedding)
+    # The 'det_model' argument is removed for compatibility with insightface==0.7.3
     app.prepare(ctx_id=CTX_ID, det_size=(640, 640), det_thresh=0.5,
-                det_model="retinaface_r50_v1", allowed_modules=['detection', 'landmark', 'recognition'])
+                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.")
+            print(f"Loaded {model_name.upper()} deepfake model.")
         else:
             print(f"Warning: Deepfake model {model_name.upper()} not found at {path}")
 
@@ -64,7 +122,8 @@ try:
 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
+    # We don't sys.exit(1) here to allow Gradio to show the error, but the function will handle it.
+
 
 print("✅ Model initialization complete.")
 
@@ -86,7 +145,7 @@ def get_face_data(img_array_rgb: np.ndarray) -> Tuple[Optional[np.ndarray], Opti
     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')
+    # The 'face' object from FaceAnalysis provides 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:
@@ -98,22 +157,20 @@ def calculate_similarity(embedding1: Optional[np.ndarray], embedding2: Optional[
     similarity = np.dot(e1_norm, e2_norm)
     return float(similarity)
 
-# --- 2D. DEEPFAKE DETECTION HELPER FUNCTIONS (Liveness Check) ---
+# --- 2D. DEEPFAKE DETECTION HELPER FUNCTIONS (Liveness Check - DLIB FREE) ---
 
 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.
+    Alignment using 5-point landmarks provided by insightface.
     """
-    # Standard 5-point template for InsightFace's alignment for deepfake models
+    # Standard 5-point template for alignment scaled for 160x160 output
     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
+    ], dtype=np.float32) * (output_size / 96) 
 
     src = landmarks_5pt.astype(np.float32)
 
@@ -135,18 +192,17 @@ def get_liveness_score(img_array_rgb: np.ndarray, landmarks_5pt: np.ndarray, mod
         # 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
+        # Pre-process for ONNX model (C, H, W, normalization [-1, 1])
         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_tensor = np.transpose(normalized_img, (2, 0, 1))
+        input_tensor = np.expand_dims(input_tensor, axis=0).astype("float32")
 
         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))
+        # Convert logit to probability (Fake Confidence Score) using sigmoid
         probability = 1 / (1 + np.exp(-logit))
         score = float(np.ravel(probability)[0]) if probability.size > 0 else 0.0
 
@@ -165,7 +221,7 @@ def unified_ekyc_analysis(model_choice: str, img_A_pil: Image.Image, img_B_pil:
     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."""
+        error_msg = f"""# ❌ CRITICAL FAILURE: Models failed to load. Please check console output and ensure all ONNX models are present and FaceAnalysis initialized successfully."""
         return None, None, error_msg
 
     start_time = time.time()
@@ -174,15 +230,12 @@ def unified_ekyc_analysis(model_choice: str, img_A_pil: Image.Image, img_B_pil:
     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, e2: Embeddings; lmk_A, lmk_B: Landmarks (5-point); vis_A_bgr, vis_B_bgr: BGR images; 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:
+    if e1 is None or e2 is None or lmk_A is None or lmk_B 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
 
@@ -288,13 +341,12 @@ def unified_ekyc_analysis(model_choice: str, img_A_pil: Image.Image, img_B_pil:
 
 
 # --- GRADIO FRONTEND ---
-print("\n--- 2. Initializing Gradio interface ---")
+print("\n--- 4. 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 ""))
+default_model = "edgenext" if "edgenext" in available_models else (available_models[0] if available_models else None)
 
-if not available_models:
+if default_model is None:
     print("FATAL: No deepfake models were loaded. Cannot launch Gradio.")
 else:
     iface = gr.Interface(
@@ -319,7 +371,7 @@ else:
             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.",
+        description="Performs two-step conditional verification: Step 1: Identity Match. Step 2 (if Match): Forgery/Liveness Check on both images. Includes automatic model download and is DLIB-free.",
     )
 
     iface.launch(debug=False)