Update app.py

app.py

@@ -1,10 +1,3 @@
-# app.py
-# --------------------------------------------------------------------
-# A Gradio-based face recognition system that mimics most features of
-# your Streamlit app: real-time webcam, image tests, configuration,
-# database enrollment, searching, user removal, etc.
-# --------------------------------------------------------------------
-
 import os
 import sys
 import math
@@ -19,7 +12,7 @@ from typing import Optional, Dict, List, Tuple
 from dataclasses import dataclass, field
 from collections import Counter
 
-# 3rd-party modules
+# 3rd-party
 import gradio as gr
 from ultralytics import YOLO
 from facenet_pytorch import InceptionResnetV1
@@ -27,7 +20,7 @@ from torchvision import transforms
 from deep_sort_realtime.deepsort_tracker import DeepSort
 
 # --------------------------------------------------------------------
-# GLOBALS & CONSTANTS
+# LOGGING
 # --------------------------------------------------------------------
 logging.basicConfig(
     level=logging.INFO,
@@ -36,20 +29,24 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# Mute some debug logs from third-party libraries
 logging.getLogger('torch').setLevel(logging.ERROR)
 logging.getLogger('deep_sort_realtime').setLevel(logging.ERROR)
 
+# --------------------------------------------------------------------
+# CONSTANTS
+# --------------------------------------------------------------------
 DEFAULT_MODEL_URL = "https://github.com/wuhplaptop/face-11-n/blob/main/face2.pt?raw=true"
 DEFAULT_DB_PATH = os.path.expanduser("~/.face_pipeline/known_faces.pkl")
 MODEL_DIR = os.path.expanduser("~/.face_pipeline/models")
 CONFIG_PATH = os.path.expanduser("~/.face_pipeline/config.pkl")
 
-# If you need blink detection or face mesh, keep or define your landmarks:
+# Example eye indices if you still want blink detection somewhere
 LEFT_EYE_IDX = [33, 160, 158, 133, 153, 144]
 RIGHT_EYE_IDX = [263, 387, 385, 362, 380, 373]
 
 # --------------------------------------------------------------------
-# PIPELINE CONFIG DATACLASS
+# PIPELINE CONFIG
 # --------------------------------------------------------------------
 @dataclass
 class PipelineConfig:
@@ -191,11 +188,11 @@ class FaceDatabase:
 
     def search_by_image(self, query_embedding: np.ndarray, threshold: float = 0.7) -> List[Tuple[str, float]]:
         results = []
-        for label, embeddings in self.embeddings.items():
-            for db_emb in embeddings:
+        for lbl, embs in self.embeddings.items():
+            for db_emb in embs:
                 similarity = FacePipeline.cosine_similarity(query_embedding, db_emb)
                 if similarity >= threshold:
-                    results.append((label, similarity))
+                    results.append((lbl, similarity))
         return sorted(results, key=lambda x: x[1], reverse=True)
 
 # --------------------------------------------------------------------
@@ -241,7 +238,7 @@ class YOLOFaceDetector:
             return []
 
 # --------------------------------------------------------------------
-# FACE TRACKER
+# FACE TRACKER (Used if you want tracking across frames - optional)
 # --------------------------------------------------------------------
 class FaceTracker:
     def __init__(self, max_age: int = 30):
@@ -287,7 +284,7 @@ class FaceNetEmbedder:
             return None
 
 # --------------------------------------------------------------------
-# MAIN PIPELINE
+# FACE PIPELINE
 # --------------------------------------------------------------------
 class FacePipeline:
     def __init__(self, config: PipelineConfig):
@@ -351,7 +348,6 @@ class FacePipeline:
                     name = "Spoofed"
                     similarity = 0.0
                 else:
-                    # Face recognition
                     embedding = self.facenet.get_embedding(face_roi)
                     if embedding is not None and self.config.recognition['enable']:
                         name, similarity = self.recognize_face(
@@ -371,7 +367,7 @@ class FacePipeline:
                 label_text = f"{name}"
                 cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), box_color_bgr, 2)
                 cv2.putText(
-                    annotated_frame, label_text, (x1, y1 - 10), 
+                    annotated_frame, label_text, (x1, y1 - 10),
                     cv2.FONT_HERSHEY_SIMPLEX, 0.5, box_color_bgr, 2
                 )
 
@@ -422,10 +418,9 @@ class FacePipeline:
         return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-6))
 
 # --------------------------------------------------------------------
-# GLOBAL pipeline instance (we can store it in a lazy loader)
+# GLOBAL LOAD PIPELINE
 # --------------------------------------------------------------------
 pipeline = None
-
 def load_pipeline() -> FacePipeline:
     global pipeline
     if pipeline is None:
@@ -436,31 +431,25 @@ def load_pipeline() -> FacePipeline:
     return pipeline
 
 # --------------------------------------------------------------------
-# GRADIO HELPER FUNCTIONS
+# GRADIO UTILS
 # --------------------------------------------------------------------
 def hex_to_bgr(hex_str: str) -> Tuple[int,int,int]:
-    """
-    Convert a hex string (#RRGGBB) into a BGR tuple as used in OpenCV.
-    """
     if not hex_str.startswith('#'):
         hex_str = f"#{hex_str}"
     hex_str = hex_str.lstrip('#')
     if len(hex_str) != 6:
-        return (255, 0, 0)  # fallback to something
+        return (255, 0, 0)
     r = int(hex_str[0:2], 16)
     g = int(hex_str[2:4], 16)
     b = int(hex_str[4:6], 16)
     return (b,g,r)
 
 def bgr_to_hex(bgr: Tuple[int,int,int]) -> str:
-    """
-    Convert a BGR tuple (as stored in pipeline config) to a #RRGGBB hex string.
-    """
     b,g,r = bgr
     return f"#{r:02x}{g:02x}{b:02x}"
 
 # --------------------------------------------------------------------
-# TAB: Configuration
+# TAB: CONFIGURATION
 # --------------------------------------------------------------------
 def update_config(
     enable_recognition, enable_antispoof, enable_blink, enable_hand, enable_eyecolor, enable_facemesh,
@@ -473,11 +462,10 @@ def update_config(
     mesh_hex, contour_hex, iris_hex,
     eye_color_text_hex
 ):
-    # Load pipeline
     pl = load_pipeline()
     cfg = pl.config
 
-    # Update toggles
+    # Toggles
     cfg.recognition['enable'] = enable_recognition
     cfg.anti_spoof['enable'] = enable_antispoof
     cfg.blink['enable'] = enable_blink
@@ -489,7 +477,7 @@ def update_config(
     cfg.face_mesh_options['contours'] = show_contours
     cfg.face_mesh_options['irises'] = show_irises
 
-    # Update thresholds
+    # Thresholds
     cfg.detection_conf_thres = detection_conf
     cfg.recognition_conf_thres = recognition_thresh
     cfg.anti_spoof['lap_thresh'] = antispoof_thresh
@@ -497,8 +485,8 @@ def update_config(
     cfg.hand['min_detection_confidence'] = hand_det_conf
     cfg.hand['min_tracking_confidence'] = hand_track_conf
 
-    # Update color fields
-    cfg.bbox_color = hex_to_bgr(bbox_hex)[::-1]         # store in (R,G,B)
+    # Colors
+    cfg.bbox_color = hex_to_bgr(bbox_hex)[::-1]
     cfg.spoofed_bbox_color = hex_to_bgr(spoofed_hex)[::-1]
     cfg.unknown_bbox_color = hex_to_bgr(unknown_hex)[::-1]
     cfg.eye_outline_color = hex_to_bgr(eye_hex)[::-1]
@@ -511,19 +499,13 @@ def update_config(
     cfg.iris_color = hex_to_bgr(iris_hex)[::-1]
     cfg.eye_color_text_color = hex_to_bgr(eye_color_text_hex)[::-1]
 
-    # Save config
     cfg.save(CONFIG_PATH)
-
     return "Configuration saved successfully!"
 
 # --------------------------------------------------------------------
-# TAB: Database Management
+# TAB: DATABASE MANAGEMENT
 # --------------------------------------------------------------------
 def enroll_user(name: str, images: List[np.ndarray]) -> str:
-    """
-    Enroll user by name using one or more images. images is a list of
-    NxMx3 numpy arrays in BGR or RGB depending on Gradio type.
-    """
     pl = load_pipeline()
     if not name:
         return "Please provide a user name."
@@ -535,13 +517,8 @@ def enroll_user(name: str, images: List[np.ndarray]) -> str:
     for img in images:
         if img is None:
             continue
-        # Gradio provides images in RGB by default, let's ensure BGR for pipeline
-        if img.shape[-1] == 3:  # RGB
-            img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-        else:
-            img_bgr = img
-
-        # Run YOLO detection on each image
+        # Gradio typically supplies images in RGB
+        img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
         detections = pl.detector.detect(img_bgr, pl.config.detection_conf_thres)
         for x1, y1, x2, y2, conf, cls in detections:
             face_roi = img_bgr[y1:y2, x1:x2]
@@ -569,9 +546,6 @@ def search_by_name(name: str) -> str:
         return f"No embeddings found for user '{name}'."
 
 def search_by_image(image: np.ndarray) -> str:
-    """
-    Search database by face in the uploaded image.
-    """
     pl = load_pipeline()
     if image is None:
         return "No image uploaded."
@@ -611,27 +585,11 @@ def list_users() -> str:
     pl = load_pipeline()
     labels = pl.db.list_labels()
     if labels:
-        return f"Enrolled users:\n{', '.join(labels)}"
+        return "Enrolled users:\n" + ", ".join(labels)
     return "No users enrolled."
 
 # --------------------------------------------------------------------
-# TAB: Real-Time Recognition
-# --------------------------------------------------------------------
-def process_webcam_frame(frame: np.ndarray) -> Tuple[np.ndarray, str]:
-    """
-    Called for every incoming webcam frame. Return annotated frame + textual info.
-    Gradio delivers frames in RGB.
-    """
-    if frame is None:
-        return None, "No frame."
-    pl = load_pipeline()
-    frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
-    annotated_bgr, detections = pl.process_frame(frame_bgr)
-    annotated_rgb = cv2.cvtColor(annotated_bgr, cv2.COLOR_BGR2RGB)
-    return annotated_rgb, str(detections)
-
-# --------------------------------------------------------------------
-# TAB: Image Test
+# TAB: IMAGE-BASED RECOGNITION
 # --------------------------------------------------------------------
 def process_test_image(img: np.ndarray) -> Tuple[np.ndarray, str]:
     if img is None:
@@ -647,21 +605,10 @@ def process_test_image(img: np.ndarray) -> Tuple[np.ndarray, str]:
 # --------------------------------------------------------------------
 def build_app():
     with gr.Blocks() as demo:
-        gr.Markdown("# Face Recognition System (Gradio)")
-
-        with gr.Tab("Real-Time Recognition"):
-            gr.Markdown("Live face recognition from your webcam (roughly 'real-time').")
-            webcam_input = gr.Video(source="webcam", mirror=True, streaming=True)
-            webcam_output = gr.Image()
-            webcam_info = gr.Textbox(label="Detections", interactive=False)
-            webcam_input.change(
-                fn=process_webcam_frame, 
-                inputs=webcam_input, 
-                outputs=[webcam_output, webcam_info],
-            )
+        gr.Markdown("# Face Recognition System (Image-Only)")
 
         with gr.Tab("Image Test"):
-            gr.Markdown("Upload a single image for face detection and recognition.")
+            gr.Markdown("Upload a single image for face detection & recognition:")
             image_input = gr.Image(type="numpy", label="Upload Image")
             image_out = gr.Image()
             image_info = gr.Textbox(label="Detections", interactive=False)
@@ -674,8 +621,8 @@ def build_app():
             )
 
         with gr.Tab("Configuration"):
-            gr.Markdown("Modify the pipeline settings and thresholds here.")
-            
+            gr.Markdown("Modify pipeline settings & thresholds.")
+
             with gr.Row():
                 enable_recognition = gr.Checkbox(label="Enable Face Recognition", value=True)
                 enable_antispoof = gr.Checkbox(label="Enable Anti-Spoof", value=True)
@@ -684,7 +631,7 @@ def build_app():
                 enable_eyecolor = gr.Checkbox(label="Enable Eye Color Detection", value=False)
                 enable_facemesh = gr.Checkbox(label="Enable Face Mesh", value=False)
 
-            gr.Markdown("**Face Mesh Options** (only if Face Mesh is enabled):")
+            gr.Markdown("**Face Mesh Options**")
             with gr.Row():
                 show_tesselation = gr.Checkbox(label="Show Tesselation", value=False)
                 show_contours = gr.Checkbox(label="Show Contours", value=False)
@@ -713,7 +660,6 @@ def build_app():
             mesh_hex = gr.Textbox(label="Mesh Color", value="#64ff64")
             contour_hex = gr.Textbox(label="Contour Color", value="#c8c800")
             iris_hex = gr.Textbox(label="Iris Color", value="#ff00ff")
-
             eye_color_text_hex = gr.Textbox(label="Eye Color Text Color", value="#ffffff")
 
             save_btn = gr.Button("Save Configuration")
@@ -722,8 +668,8 @@ def build_app():
             save_btn.click(
                 fn=update_config,
                 inputs=[
-                    enable_recognition, enable_antispoof, enable_blink,
-                    enable_hand, enable_eyecolor, enable_facemesh,
+                    enable_recognition, enable_antispoof, enable_blink, enable_hand, 
+                    enable_eyecolor, enable_facemesh,
                     show_tesselation, show_contours, show_irises,
                     detection_conf, recognition_thresh, antispoof_thresh, blink_thresh,
                     hand_det_conf, hand_track_conf,
@@ -753,14 +699,17 @@ def build_app():
                 search_result = gr.Textbox(label="", interactive=False)
 
                 def update_search_visibility(mode):
+                    # Show name dropdown if "Name", else show image upload
                     if mode == "Name":
                         return gr.update(visible=True), gr.update(visible=False)
                     else:
                         return gr.update(visible=False), gr.update(visible=True)
 
-                search_mode.change(fn=update_search_visibility, 
-                                   inputs=[search_mode], 
-                                   outputs=[search_name_input, search_image_input])
+                search_mode.change(
+                    fn=update_search_visibility, 
+                    inputs=[search_mode], 
+                    outputs=[search_name_input, search_image_input]
+                )
 
                 def search_user(mode, name, img):
                     if mode == "Name":
@@ -768,21 +717,21 @@ def build_app():
                     else:
                         return search_by_image(img)
 
-                search_btn.click(fn=search_user, 
-                                 inputs=[search_mode, search_name_input, search_image_input], 
-                                 outputs=[search_result])
+                search_btn.click(
+                    fn=search_user, 
+                    inputs=[search_mode, search_name_input, search_image_input], 
+                    outputs=[search_result]
+                )
 
             with gr.Accordion("User Management Tools", open=False):
                 list_btn = gr.Button("List Enrolled Users")
                 list_output = gr.Textbox(label="", interactive=False)
                 list_btn.click(fn=lambda: list_users(), inputs=[], outputs=[list_output])
 
-                # Reload user list dropdown
                 def get_user_list():
                     pl = load_pipeline()
                     return gr.update(choices=pl.db.list_labels())
 
-                # A dedicated button to refresh the dropdown
                 refresh_users_btn = gr.Button("Refresh User List")
                 refresh_users_btn.click(fn=get_user_list, inputs=[], outputs=[search_name_input])
 
@@ -800,4 +749,5 @@ def build_app():
 # --------------------------------------------------------------------
 if __name__ == "__main__":
     app = build_app()
+    # queue() is optional if you expect concurrency
     app.queue().launch(server_name="0.0.0.0", server_port=7860)