Update app.py

app.py

@@ -12,11 +12,17 @@ from typing import Optional, Dict, List, Tuple
 from dataclasses import dataclass, field
 from collections import Counter
 
-# 3rd-party
+# Gradio
 import gradio as gr
+
+# YOLO (Ultralytics)
 from ultralytics import YOLO
+
+# FaceNet
 from facenet_pytorch import InceptionResnetV1
 from torchvision import transforms
+
+# DeepSORT tracking
 from deep_sort_realtime.deepsort_tracker import DeepSort
 
 # --------------------------------------------------------------------
@@ -41,7 +47,7 @@ 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")
 
-# Example eye indices if you still want blink detection somewhere
+# If you still want blink detection or face mesh references, keep them
 LEFT_EYE_IDX = [33, 160, 158, 133, 153, 144]
 RIGHT_EYE_IDX = [263, 387, 385, 362, 380, 373]
 
@@ -238,7 +244,7 @@ class YOLOFaceDetector:
             return []
 
 # --------------------------------------------------------------------
-# FACE TRACKER (Used if you want tracking across frames - optional)
+# FACE TRACKER
 # --------------------------------------------------------------------
 class FaceTracker:
     def __init__(self, max_age: int = 30):
@@ -298,7 +304,7 @@ class FacePipeline:
     def initialize(self):
         try:
             self.detector = YOLOFaceDetector(
-                model_path=self.config.detector['model_path'], 
+                model_path=self.config.detector['model_path'],
                 device=self.config.detector['device']
             )
             self.tracker = FaceTracker(max_age=self.config.tracker['max_age'])
@@ -358,8 +364,8 @@ class FacePipeline:
                         similarity = 0.0
 
                     box_color_rgb = (
-                        self.config.bbox_color 
-                        if name != "Unknown" 
+                        self.config.bbox_color
+                        if name != "Unknown"
                         else self.config.unknown_bbox_color
                     )
                     box_color_bgr = box_color_rgb[::-1]
@@ -418,7 +424,7 @@ class FacePipeline:
         return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-6))
 
 # --------------------------------------------------------------------
-# GLOBAL LOAD PIPELINE
+# GLOBAL PIPELINE
 # --------------------------------------------------------------------
 pipeline = None
 def load_pipeline() -> FacePipeline:
@@ -431,7 +437,7 @@ def load_pipeline() -> FacePipeline:
     return pipeline
 
 # --------------------------------------------------------------------
-# GRADIO UTILS
+# HELPER FUNCTIONS FOR GRADIO
 # --------------------------------------------------------------------
 def hex_to_bgr(hex_str: str) -> Tuple[int,int,int]:
     if not hex_str.startswith('#'):
@@ -449,7 +455,7 @@ def bgr_to_hex(bgr: Tuple[int,int,int]) -> str:
     return f"#{r:02x}{g:02x}{b:02x}"
 
 # --------------------------------------------------------------------
-# TAB: CONFIGURATION
+# UPDATE CONFIG
 # --------------------------------------------------------------------
 def update_config(
     enable_recognition, enable_antispoof, enable_blink, enable_hand, enable_eyecolor, enable_facemesh,
@@ -503,22 +509,30 @@ def update_config(
     return "Configuration saved successfully!"
 
 # --------------------------------------------------------------------
-# TAB: DATABASE MANAGEMENT
+# ENROLL MULTIPLE IMAGES (using gr.File)
 # --------------------------------------------------------------------
-def enroll_user(name: str, images: List[np.ndarray]) -> str:
+def enroll_user(name: str, files: List[dict]) -> str:
+    """
+    Each item in `files` is a dict with keys:
+      - 'name': filename
+      - 'size': file size in bytes
+      - 'data': the binary file contents
+    We decode them into OpenCV images and process.
+    """
     pl = load_pipeline()
     if not name:
         return "Please provide a user name."
-
-    if not images or len(images) == 0:
+    if not files or len(files) == 0:
         return "No images provided."
 
     count_enrolled = 0
-    for img in images:
-        if img is None:
+    for f in files:
+        file_bytes = f["data"]
+        np_array = np.frombuffer(file_bytes, np.uint8)
+        img_bgr = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
+        if img_bgr is None:
             continue
-        # 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]
@@ -535,6 +549,9 @@ def enroll_user(name: str, images: List[np.ndarray]) -> str:
     else:
         return "No faces were detected or embedded. Enrollment failed."
 
+# --------------------------------------------------------------------
+# SEARCH / USER MGMT
+# --------------------------------------------------------------------
 def search_by_name(name: str) -> str:
     pl = load_pipeline()
     if not name:
@@ -550,6 +567,7 @@ def search_by_image(image: np.ndarray) -> str:
     if image is None:
         return "No image uploaded."
 
+    # Convert to BGR for pipeline
     img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
     detections = pl.detector.detect(img_bgr, pl.config.detection_conf_thres)
     if not detections:
@@ -589,7 +607,7 @@ def list_users() -> str:
     return "No users enrolled."
 
 # --------------------------------------------------------------------
-# TAB: IMAGE-BASED RECOGNITION
+# PROCESS SINGLE IMAGE (IMAGE TEST)
 # --------------------------------------------------------------------
 def process_test_image(img: np.ndarray) -> Tuple[np.ndarray, str]:
     if img is None:
@@ -607,6 +625,7 @@ def build_app():
     with gr.Blocks() as demo:
         gr.Markdown("# Face Recognition System (Image-Only)")
 
+        # -- Tab: Image Test --
         with gr.Tab("Image Test"):
             gr.Markdown("Upload a single image for face detection & recognition:")
             image_input = gr.Image(type="numpy", label="Upload Image")
@@ -620,6 +639,7 @@ def build_app():
                 outputs=[image_out, image_info],
             )
 
+        # -- Tab: Configuration --
         with gr.Tab("Configuration"):
             gr.Markdown("Modify pipeline settings & thresholds.")
 
@@ -668,7 +688,7 @@ def build_app():
             save_btn.click(
                 fn=update_config,
                 inputs=[
-                    enable_recognition, enable_antispoof, enable_blink, enable_hand, 
+                    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,
@@ -681,15 +701,21 @@ def build_app():
                 outputs=[save_msg],
             )
 
+        # -- Tab: Database Management --
         with gr.Tab("Database Management"):
             gr.Markdown("Enroll Users, Search by Name or Image, Remove or List.")
 
             with gr.Accordion("User Enrollment", open=False):
                 enroll_name = gr.Textbox(label="Enter name for enrollment")
-                enroll_images = gr.Image(type="numpy", label="Upload Enrollment Images", multiple=True)
+                # Use gr.File with multiple file support:
+                enroll_files = gr.File(
+                    file_count="multiple",
+                    type="file",
+                    label="Upload Enrollment Images (multiple allowed)"
+                )
                 enroll_btn = gr.Button("Enroll User")
                 enroll_result = gr.Textbox(label="", interactive=False)
-                enroll_btn.click(fn=enroll_user, inputs=[enroll_name, enroll_images], outputs=[enroll_result])
+                enroll_btn.click(fn=enroll_user, inputs=[enroll_name, enroll_files], outputs=[enroll_result])
 
             with gr.Accordion("User Search", open=False):
                 search_mode = gr.Radio(["Name", "Image"], value="Name", label="Search Database By")
@@ -699,15 +725,14 @@ 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], 
+                    fn=update_search_visibility,
+                    inputs=[search_mode],
                     outputs=[search_name_input, search_image_input]
                 )
 
@@ -718,8 +743,8 @@ def build_app():
                         return search_by_image(img)
 
                 search_btn.click(
-                    fn=search_user, 
-                    inputs=[search_mode, search_name_input, search_image_input], 
+                    fn=search_user,
+                    inputs=[search_mode, search_name_input, search_image_input],
                     outputs=[search_result]
                 )
 
@@ -749,5 +774,5 @@ def build_app():
 # --------------------------------------------------------------------
 if __name__ == "__main__":
     app = build_app()
-    # queue() is optional if you expect concurrency
+    # If concurrency or GPU usage is high, you can use .queue():
     app.queue().launch(server_name="0.0.0.0", server_port=7860)