Update app.py

app.py

@@ -1,4 +1,47 @@
 # Install required packages
+import os
+import subprocess
+import sys
+import importlib
+import pkg_resources
+
+def install_package(package, version=None):
+    package_spec = f"{package}=={version}" if version else package
+    print(f"Installing {package_spec}...")
+    try:
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-cache-dir", package_spec])
+    except subprocess.CalledProcessError as e:
+        print(f"Failed to install {package_spec}: {e}")
+        raise
+
+def ensure_package(package, version=None):
+    try:
+        if version:
+            pkg_resources.require(f"{package}=={version}")
+        else:
+            importlib.import_module(package)
+        print(f"{package} is already installed with the correct version.")
+    except (ImportError, pkg_resources.VersionConflict, pkg_resources.DistributionNotFound) as e:
+        print(f"Package requirement failed: {e}")
+        install_package(package, version)
+
+# Check environment and install dependencies
+if not os.path.exists("/.dockerenv") and not os.path.exists("/kaggle"):
+    print("Setting up environment...")
+    
+    # Install core dependencies
+    ensure_package("numpy", "1.23.5")
+    ensure_package("protobuf", "3.20.3")
+    ensure_package("tensorflow", "2.10.0")
+    ensure_package("opencv-python-headless", "4.7.0.72")  # Fix for cv2 error
+    ensure_package("deepface", "0.0.79")
+    ensure_package("gradio", "3.50.2")
+    
+    # Install additional required packages
+    for pkg in ["matplotlib", "pillow", "pandas"]:
+        ensure_package(pkg)
+
+# Now import the required modules
 import gradio as gr
 import json
 import cv2
@@ -7,10 +50,14 @@ from deepface import DeepFace
 import matplotlib.pyplot as plt
 from PIL import Image
 import tempfile
-import os
 import pandas as pd
 import shutil
 
+# Google Drive integration (for Colab users)
+if 'google.colab' in sys.modules:
+    from google.colab import drive
+    drive.mount('/content/drive')
+
 def verify_faces(img1, img2, threshold=0.70, model="VGG-Face"):
     temp_dir = tempfile.mkdtemp()
     img1_path = os.path.join(temp_dir, "image1.jpg")
@@ -88,14 +135,23 @@ def find_faces(query_img, db_folder, threshold=0.70, model="VGG-Face"):
     else:
         query_img.save(query_path)
     
+    # Handle cloud storage paths and uploaded files
     if isinstance(db_folder, str):
-        db_path = db_folder
+        # Check if it's a Google Drive path
+        if db_folder.startswith("/content/drive"):
+            db_path = db_folder
+        else:
+            # Handle regular path
+            db_path = os.path.abspath(db_folder)
+            if not os.path.exists(db_path):
+                return None, "Invalid database path - directory does not exist"
     else:
+        # Handle uploaded files
         db_path = os.path.join(temp_dir, "db")
         os.makedirs(db_path, exist_ok=True)
         
         for i, file in enumerate(db_folder):
-            orig_filename = file.orig_name  # Get original filename with extension
+            orig_filename = file.orig_name
             file_ext = os.path.splitext(orig_filename)[1]
             new_filename = f"image_{i}{file_ext}"
             shutil.copy(file.name, os.path.join(db_path, new_filename))
@@ -107,7 +163,7 @@ def find_faces(query_img, db_folder, threshold=0.70, model="VGG-Face"):
             model_name=model,
             distance_metric="cosine",
             threshold=threshold,
-            silent=True  # Disable unnecessary logging
+            silent=True
         )
         
         if isinstance(dfs, list):
@@ -131,23 +187,34 @@ def find_faces(query_img, db_folder, threshold=0.70, model="VGG-Face"):
         axes[0].set_title("Query Image")
         axes[0].axis("off")
         
+        valid_matches = 0
         for i in range(num_matches):
+            if i >= len(df):
+                break
+                
             match_path = df.iloc[i]["identity"]
             if not os.path.exists(match_path):
-                continue  # Skip invalid paths
+                continue
                 
             distance = df.iloc[i]["distance"]
             confidence = round((1 - distance) * 100, 2)
             
-            match_img = cv2.imread(match_path)
-            if match_img is None:
-                continue  # Skip unreadable images
-                
-            match_img = cv2.cvtColor(match_img, cv2.COLOR_BGR2RGB)
-            
-            axes[i+1].imshow(match_img)
-            axes[i+1].set_title(f"Match #{i+1}\nConfidence: {confidence}%")
-            axes[i+1].axis("off")
+            try:
+                match_img = cv2.imread(match_path)
+                if match_img is None:
+                    continue
+                    
+                match_img = cv2.cvtColor(match_img, cv2.COLOR_BGR2RGB)
+                axes[valid_matches+1].imshow(match_img)
+                axes[valid_matches+1].set_title(f"Match #{valid_matches+1}\nConfidence: {confidence}%")
+                axes[valid_matches+1].axis("off")
+                valid_matches += 1
+            except:
+                continue
+        
+        # Hide empty axes
+        for j in range(valid_matches+1, num_matches+1):
+            axes[j].axis("off")
         
         plt.suptitle(f"Found {len(df)} matching faces", fontsize=16, fontweight='bold')
         plt.tight_layout()
@@ -238,7 +305,7 @@ def analyze_face(img, actions=['age', 'gender', 'race', 'emotion']):
                         emotion_conf = f"{conf:.1f}%"
                         break
             
-            ax = plt.subplot2grid((2, 4), (0 if i < 2 else 1, 2 + (i % 2)))
+            ax = plt.subplot2grid((2, 4), (0 if i < 2 else 1, 2 + (i % 2))
             
             text = (
                 f"Face #{i+1}\n\n"
@@ -297,7 +364,7 @@ with gr.Blocks(title="Complete Face Recognition Tool", theme=gr.themes.Soft()) a
     
     This tool provides three face recognition features:
     - **Verify Faces**: Compare two specific images to check if they contain the same person
-    - **Find Faces**: Search for matching faces in a database/folder
+    - **Find Faces**: Search for matching faces in a database/folder (supports Google Drive paths in Colab)
     - **Analyze Face**: Determine age, gender, race, and emotion from a facial image
     """)
     
@@ -320,18 +387,13 @@ with gr.Blocks(title="Complete Face Recognition Tool", theme=gr.themes.Soft()) a
             
             verify_result_plot = gr.Plot(label="Verification Result")
             verify_json = gr.JSON(label="Technical Details")
-            
-            verify_button.click(
-                verify_faces, 
-                inputs=[img1_input, img2_input, verify_threshold, verify_model], 
-                outputs=[verify_result_plot, verify_json]
-            )
-            
-            gr.Markdown("""... (keep existing markdown) ...""")
-        
+
         with gr.TabItem("Find Faces"):
             query_img = gr.Image(label="Query Image (Face to find)", type="pil")
-            db_path_input = gr.Textbox(label="Database Path (folder containing images to search in)")
+            db_path_input = gr.Textbox(
+                label="Database Path (folder path or Google Drive path in Colab)",
+                placeholder="/content/drive/MyDrive/your_folder"
+            )
             db_files_input = gr.File(label="Or upload images for database", file_count="multiple")
             
             with gr.Row():
@@ -347,21 +409,7 @@ with gr.Blocks(title="Complete Face Recognition Tool", theme=gr.themes.Soft()) a
             
             find_result_plot = gr.Plot(label="Search Results")
             find_results_table = gr.JSON(label="Detailed Results")
-            
-            find_button.click(
-                find_faces,
-                inputs=[query_img, db_path_input, find_threshold, find_model],
-                outputs=[find_result_plot, find_results_table]
-            )
-            
-            db_files_input.change(
-                lambda x: "",
-                inputs=db_files_input,
-                outputs=db_path_input
-            )
-            
-            gr.Markdown("""... (keep existing markdown) ...""")
-        
+
         with gr.TabItem("Analyze Face"):
             analyze_img = gr.Image(label="Upload Image for Analysis", type="pil")
             actions_checkboxes = gr.CheckboxGroup(
@@ -374,14 +422,31 @@ with gr.Blocks(title="Complete Face Recognition Tool", theme=gr.themes.Soft()) a
             
             analyze_result_plot = gr.Plot(label="Analysis Results")
             analyze_json = gr.JSON(label="Detailed Analysis")
-            
-            analyze_button.click(
-                analyze_face,
-                inputs=[analyze_img, actions_checkboxes],
-                outputs=[analyze_result_plot, analyze_json]
-            )
-            
-            gr.Markdown("""... (keep existing markdown) ...""")
+
+    # Setup all button clicks
+    verify_button.click(
+        verify_faces, 
+        inputs=[img1_input, img2_input, verify_threshold, verify_model], 
+        outputs=[verify_result_plot, verify_json]
+    )
+    
+    find_button.click(
+        find_faces,
+        inputs=[query_img, db_path_input, find_threshold, find_model],
+        outputs=[find_result_plot, find_results_table]
+    )
+    
+    db_files_input.change(
+        lambda x: "",
+        inputs=db_files_input,
+        outputs=db_path_input
+    )
+    
+    analyze_button.click(
+        analyze_face,
+        inputs=[analyze_img, actions_checkboxes],
+        outputs=[analyze_result_plot, analyze_json]
+    )
 
 # Launch the app
 demo.launch()