wuhp commited on
Commit
3a57d6b
·
verified ·
1 Parent(s): 4166933

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1245 -3
app.py CHANGED
@@ -1,6 +1,1248 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  import gradio as gr
2
- from face_pipeline import FacePipeline
3
 
4
- pipeline = FacePipeline()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5
  if __name__ == "__main__":
6
- iface.launch()
 
1
+ import os
2
+ import sys
3
+ import math
4
+ import requests
5
+ import numpy as np
6
+ import cv2
7
+ import torch
8
+ import pickle
9
+ import logging
10
+ from PIL import Image
11
+ from typing import Optional, Dict, List, Tuple
12
+ from dataclasses import dataclass, field
13
+ from collections import Counter
14
+ import io
15
+ import tempfile # Import tempfile
16
+
17
  import gradio as gr
 
18
 
19
+ from ultralytics import YOLO
20
+ from facenet_pytorch import InceptionResnetV1
21
+ from torchvision import transforms
22
+ from deep_sort_realtime.deepsort_tracker import DeepSort
23
+
24
+ import mediapipe as mp
25
+
26
+ # Configure logging
27
+ logging.basicConfig(
28
+ level=logging.DEBUG, # Changed to DEBUG for more detailed logs
29
+ format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
30
+ handlers=[logging.FileHandler('face_pipeline.log'), logging.StreamHandler()],
31
+ )
32
+ logger = logging.getLogger(__name__)
33
+
34
+ # Suppress verbose logs from libraries
35
+ logging.getLogger('torch').setLevel(logging.ERROR)
36
+ logging.getLogger('mediapipe').setLevel(logging.ERROR)
37
+ logging.getLogger('deep_sort_realtime').setLevel(logging.ERROR)
38
+
39
+ # Constants and default paths
40
+ DEFAULT_MODEL_URL = "https://github.com/wuhplaptop/face-11-n/blob/main/face2.pt?raw=true"
41
+ DEFAULT_DB_PATH = os.path.expanduser("~/.face_pipeline/known_faces.pkl")
42
+ MODEL_DIR = os.path.expanduser("~/.face_pipeline/models")
43
+ CONFIG_PATH = os.path.expanduser("~/.face_pipeline/config.pkl")
44
+
45
+ # Mediapipe indices for eye landmarks
46
+ LEFT_EYE_IDX = [33, 160, 158, 133, 153, 144]
47
+ RIGHT_EYE_IDX = [263, 387, 385, 362, 380, 373]
48
+
49
+ # Initialize Mediapipe drawing utilities
50
+ mp_drawing = mp.solutions.drawing_utils
51
+ mp_face_mesh = mp.solutions.face_mesh
52
+ mp_hands = mp.solutions.hands
53
+
54
+ @dataclass
55
+ class PipelineConfig:
56
+ detector: Dict = field(default_factory=dict)
57
+ tracker: Dict = field(default_factory=dict)
58
+ recognition: Dict = field(default_factory=dict)
59
+ anti_spoof: Dict = field(default_factory=dict)
60
+ blink: Dict = field(default_factory=dict)
61
+ face_mesh_options: Dict = field(default_factory=dict)
62
+ hand: Dict = field(default_factory=dict)
63
+ eye_color: Dict = field(default_factory=dict)
64
+ enabled_components: Dict = field(default_factory=dict)
65
+
66
+ detection_conf_thres: float = 0.4
67
+ recognition_conf_thres: float = 0.85
68
+
69
+ bbox_color: Tuple[int, int, int] = (0, 255, 0)
70
+ spoofed_bbox_color: Tuple[int, int, int] = (0, 0, 255)
71
+ unknown_bbox_color: Tuple[int, int, int] = (0, 0, 255)
72
+ eye_outline_color: Tuple[int, int, int] = (255, 255, 0)
73
+ blink_text_color: Tuple[int, int, int] = (0, 0, 255)
74
+ hand_landmark_color: Tuple[int, int, int] = (255, 210, 77)
75
+ hand_connection_color: Tuple[int, int, int] = (204, 102, 0)
76
+ hand_text_color: Tuple[int, int, int] = (255, 255, 255)
77
+ mesh_color: Tuple[int, int, int] = (100, 255, 100)
78
+ contour_color: Tuple[int, int, int] = (200, 200, 0)
79
+ iris_color: Tuple[int, int, int] = (255, 0, 255)
80
+ eye_color_text_color: Tuple[int, int, int] = (255, 255, 255)
81
+
82
+ def __post_init__(self):
83
+ self.detector = self.detector or {
84
+ 'model_path': os.path.join(MODEL_DIR, "face2.pt"),
85
+ 'device': 'cuda' if torch.cuda.is_available() else 'cpu',
86
+ }
87
+ self.tracker = self.tracker or {'max_age': 30}
88
+ self.recognition = self.recognition or {'enable': True}
89
+ self.anti_spoof = self.anti_spoof or {'enable': True, 'lap_thresh': 80.0}
90
+ self.blink = self.blink or {'enable': True, 'ear_thresh': 0.25}
91
+ self.face_mesh_options = self.face_mesh_options or {
92
+ 'enable': False,
93
+ 'tesselation': False,
94
+ 'contours': False,
95
+ 'irises': False,
96
+ }
97
+ self.hand = self.hand or {
98
+ 'enable': True,
99
+ 'min_detection_confidence': 0.5,
100
+ 'min_tracking_confidence': 0.5,
101
+ }
102
+ self.eye_color = self.eye_color or {'enable': False}
103
+ self.enabled_components = self.enabled_components or {
104
+ 'detection': True,
105
+ 'tracking': True,
106
+ 'anti_spoof': True,
107
+ 'recognition': True,
108
+ 'blink': True,
109
+ 'face_mesh': False,
110
+ 'hand': True,
111
+ 'eye_color': False,
112
+ }
113
+
114
+ def save(self, path: str):
115
+ """Save this config to a pickle file."""
116
+ try:
117
+ os.makedirs(os.path.dirname(path), exist_ok=True)
118
+ with open(path, 'wb') as f:
119
+ pickle.dump(self.__dict__, f)
120
+ logger.info(f"Saved config to {path}")
121
+ logger.debug(f"Config data saved: {self.__dict__}") # Added debug log
122
+ except Exception as e:
123
+ logger.error(f"Config save failed: {str(e)}")
124
+ raise RuntimeError(f"Config save failed: {str(e)}") from e
125
+
126
+ @classmethod
127
+ def load(cls, path: str) -> 'PipelineConfig':
128
+ """Load a config from a pickle file."""
129
+ try:
130
+ if os.path.exists(path):
131
+ with open(path, 'rb') as f:
132
+ data = pickle.load(f)
133
+ logger.info(f"Loaded config from {path}")
134
+ logger.debug(f"Config data loaded: {data}") # Added debug log
135
+ return cls(**data)
136
+ logger.info("No config file found, using default config.") # Added log for default case
137
+ return cls()
138
+ except Exception as e:
139
+ logger.error(f"Config load failed: {str(e)}")
140
+ return cls()
141
+
142
+ def export_config(self) -> bytes:
143
+ """Export your config to bytes."""
144
+ try:
145
+ config_data = self.__dict__
146
+ buf = io.BytesIO()
147
+ pickle.dump(config_data, buf)
148
+ buf.seek(0)
149
+ return buf.read()
150
+ except Exception as e:
151
+ logger.error(f"Export config failed: {str(e)}")
152
+ raise RuntimeError(f"Export config failed: {str(e)}") from e
153
+
154
+ @classmethod
155
+ def import_config(cls, config_bytes: bytes) -> 'PipelineConfig':
156
+ """Import config from bytes."""
157
+ try:
158
+ buf = io.BytesIO(config_bytes)
159
+ data = pickle.load(buf)
160
+ return cls(**data)
161
+ except Exception as e:
162
+ logger.error(f"Import config failed: {str(e)}")
163
+ raise RuntimeError(f"Import config failed: {str(e)}") from e
164
+
165
+ class FaceDatabase:
166
+ def __init__(self, db_path: str = DEFAULT_DB_PATH):
167
+ self.db_path = db_path
168
+ self.embeddings: Dict[str, List[np.ndarray]] = {}
169
+ self._load()
170
+
171
+ def _load(self):
172
+ try:
173
+ if os.path.exists(self.db_path):
174
+ with open(self.db_path, 'rb') as f:
175
+ self.embeddings = pickle.load(f)
176
+ logger.info(f"Loaded database from {self.db_path}")
177
+ except Exception as e:
178
+ logger.error(f"Database load failed: {str(e)}")
179
+ self.embeddings = {}
180
+
181
+ def save(self):
182
+ try:
183
+ os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
184
+ with open(self.db_path, 'wb') as f:
185
+ pickle.dump(self.embeddings, f)
186
+ logger.info(f"Saved database to {self.db_path}")
187
+ except Exception as e:
188
+ logger.error(f"Database save failed: {str(e)}")
189
+ raise RuntimeError(f"Database save failed: {str(e)}") from e
190
+
191
+ def export_database(self) -> bytes:
192
+ """Export the entire face embeddings DB to bytes."""
193
+ try:
194
+ db_data = self.embeddings
195
+ buf = io.BytesIO()
196
+ pickle.dump(db_data, buf)
197
+ buf.seek(0)
198
+ return buf.read()
199
+ except Exception as e:
200
+ logger.error(f"Export database failed: {str(e)}")
201
+ raise RuntimeError(f"Export database failed: {str(e)}") from e
202
+
203
+ def import_database(self, db_bytes: bytes, merge: bool = True):
204
+ """
205
+ Import embeddings from bytes.
206
+ If merge=True, merges with current DB. If False, overwrites.
207
+ """
208
+ try:
209
+ buf = io.BytesIO(db_bytes)
210
+ imported_data = pickle.load(buf)
211
+ if not isinstance(imported_data, dict):
212
+ raise ValueError("Imported data is not a dictionary!")
213
+
214
+ if merge:
215
+ for label, emb_list in imported_data.items():
216
+ if label not in self.embeddings:
217
+ self.embeddings[label] = []
218
+ self.embeddings[label].extend(emb_list)
219
+ else:
220
+ self.embeddings = imported_data
221
+
222
+ self.save()
223
+ logger.info(f"Imported face database, merge={merge}")
224
+ except Exception as e:
225
+ logger.error(f"Import database failed: {str(e)}")
226
+ raise RuntimeError(f"Import database failed: {str(e)}") from e
227
+
228
+ def add_embedding(self, label: str, embedding: np.ndarray):
229
+ try:
230
+ if not isinstance(embedding, np.ndarray) or embedding.ndim != 1:
231
+ raise ValueError("Invalid embedding format")
232
+ if label not in self.embeddings:
233
+ self.embeddings[label] = []
234
+ self.embeddings[label].append(embedding)
235
+ logger.debug(f"Added embedding for {label}")
236
+ except Exception as e:
237
+ logger.error(f"Add embedding failed: {str(e)}")
238
+ raise
239
+
240
+ def remove_label(self, label: str):
241
+ try:
242
+ if label in self.embeddings:
243
+ del self.embeddings[label]
244
+ logger.info(f"Removed {label}")
245
+ else:
246
+ logger.warning(f"Label {label} not found")
247
+ except Exception as e:
248
+ logger.error(f"Remove label failed: {str(e)}")
249
+ raise
250
+
251
+ def list_labels(self) -> List[str]:
252
+ return list(self.embeddings.keys())
253
+
254
+ def get_embeddings_by_label(self, label: str) -> Optional[List[np.ndarray]]:
255
+ return self.embeddings.get(label)
256
+
257
+ def search_by_image(self, query_embedding: np.ndarray, threshold: float = 0.7) -> List[Tuple[str, float]]:
258
+ results = []
259
+ for lbl, embs in self.embeddings.items():
260
+ for db_emb in embs:
261
+ sim = FacePipeline.cosine_similarity(query_embedding, db_emb)
262
+ if sim >= threshold:
263
+ results.append((lbl, sim))
264
+ return sorted(results, key=lambda x: x[1], reverse=True)
265
+
266
+ class YOLOFaceDetector:
267
+ def __init__(self, model_path: str, device: str = 'cpu'):
268
+ self.model = None
269
+ self.device = device
270
+ try:
271
+ if not os.path.exists(model_path):
272
+ logger.info(f"Model not found at {model_path}. Downloading from GitHub...")
273
+ resp = requests.get(DEFAULT_MODEL_URL)
274
+ resp.raise_for_status()
275
+ os.makedirs(os.path.dirname(model_path), exist_ok=True)
276
+ with open(model_path, 'wb') as f:
277
+ f.write(resp.content)
278
+ logger.info(f"Downloaded YOLO model to {model_path}")
279
+
280
+ self.model = YOLO(model_path)
281
+ self.model.to(device)
282
+ logger.info(f"Loaded YOLO model from {model_path}")
283
+ except Exception as e:
284
+ logger.error(f"YOLO init failed: {str(e)}")
285
+ raise
286
+
287
+ def detect(self, image: np.ndarray, conf_thres: float) -> List[Tuple[int, int, int, int, float, int]]:
288
+ try:
289
+ results = self.model.predict(
290
+ source=image, conf=conf_thres, verbose=False, device=self.device
291
+ )
292
+ detections = []
293
+ for result in results:
294
+ for box in result.boxes:
295
+ x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
296
+ conf = float(box.conf[0].cpu().numpy())
297
+ cls = int(box.cls[0].cpu().numpy()) if box.cls is not None else 0
298
+ detections.append((int(x1), int(y1), int(x2), int(y2), conf, cls))
299
+ logger.debug(f"Detected {len(detections)} faces.")
300
+ return detections
301
+ except Exception as e:
302
+ logger.error(f"Detection error: {str(e)}")
303
+ return []
304
+
305
+ class FaceTracker:
306
+ def __init__(self, max_age: int = 30):
307
+ self.tracker = DeepSort(max_age=max_age, embedder='mobilenet')
308
+
309
+ def update(self, detections: List[Tuple], frame: np.ndarray):
310
+ try:
311
+ ds_detections = [
312
+ ([x1, y1, x2 - x1, y2 - y1], conf, cls)
313
+ for (x1, y1, x2, y2, conf, cls) in detections
314
+ ]
315
+ tracks = self.tracker.update_tracks(ds_detections, frame=frame)
316
+ logger.debug(f"Updated tracker with {len(tracks)} tracks.")
317
+ return tracks
318
+ except Exception as e:
319
+ logger.error(f"Tracking error: {str(e)}")
320
+ return []
321
+
322
+ class FaceNetEmbedder:
323
+ def __init__(self, device: str = 'cpu'):
324
+ self.device = device
325
+ self.model = InceptionResnetV1(pretrained='vggface2').eval().to(device)
326
+ self.transform = transforms.Compose([
327
+ transforms.Resize((160, 160)),
328
+ transforms.ToTensor(),
329
+ transforms.Normalize([0.5]*3, [0.5]*3),
330
+ ])
331
+
332
+ def get_embedding(self, face_bgr: np.ndarray) -> Optional[np.ndarray]:
333
+ try:
334
+ face_rgb = cv2.cvtColor(face_bgr, cv2.COLOR_BGR2RGB)
335
+ pil_img = Image.fromarray(face_rgb).convert('RGB')
336
+ tens = self.transform(pil_img).unsqueeze(0).to(self.device)
337
+ with torch.no_grad():
338
+ embedding = self.model(tens)[0].cpu().numpy()
339
+ logger.debug(f"Generated embedding sample: {embedding[:5]}...")
340
+ return embedding
341
+ except Exception as e:
342
+ logger.error(f"Embedding failed: {str(e)}")
343
+ return None
344
+
345
+ def detect_blink(face_roi: np.ndarray, threshold: float = 0.25) -> Tuple[bool, float, float, Optional[np.ndarray], Optional[np.ndarray]]:
346
+ """
347
+ Returns:
348
+ (blink_bool, left_ear, right_ear, left_eye_points, right_eye_points).
349
+ """
350
+ try:
351
+ face_mesh_proc = mp_face_mesh.FaceMesh(
352
+ static_image_mode=True,
353
+ max_num_faces=1,
354
+ refine_landmarks=True,
355
+ min_detection_confidence=0.5
356
+ )
357
+ result = face_mesh_proc.process(cv2.cvtColor(face_roi, cv2.COLOR_BGR2RGB))
358
+ face_mesh_proc.close()
359
+
360
+ if not result.multi_face_landmarks:
361
+ return False, 0.0, 0.0, None, None
362
+
363
+ landmarks = result.multi_face_landmarks[0].landmark
364
+ h, w = face_roi.shape[:2]
365
+
366
+ def eye_aspect_ratio(indices):
367
+ pts = [(landmarks[i].x * w, landmarks[i].y * h) for i in indices]
368
+ vertical = np.linalg.norm(np.array(pts[1]) - np.array(pts[5])) + \
369
+ np.linalg.norm(np.array(pts[2]) - np.array(pts[4]))
370
+ horizontal = np.linalg.norm(np.array(pts[0]) - np.array(pts[3]))
371
+ return vertical / (2.0 * horizontal + 1e-6)
372
+
373
+ left_ear = eye_aspect_ratio(LEFT_EYE_IDX)
374
+ right_ear = eye_aspect_ratio(RIGHT_EYE_IDX)
375
+
376
+ blink = (left_ear < threshold) and (right_ear < threshold)
377
+
378
+ left_eye_pts = np.array([(int(landmarks[i].x * w), int(landmarks[i].y * h)) for i in LEFT_EYE_IDX])
379
+ right_eye_pts = np.array([(int(landmarks[i].x * w), int(landmarks[i].y * h)) for i in RIGHT_EYE_IDX])
380
+
381
+ return blink, left_ear, right_ear, left_eye_pts, right_eye_pts
382
+
383
+ except Exception as e:
384
+ logger.error(f"Blink detection error: {str(e)}")
385
+ return False, 0.0, 0.0, None, None
386
+
387
+ def process_face_mesh(face_roi: np.ndarray):
388
+ try:
389
+ fm_proc = mp_face_mesh.FaceMesh(
390
+ static_image_mode=True,
391
+ max_num_faces=1,
392
+ refine_landmarks=True,
393
+ min_detection_confidence=0.5
394
+ )
395
+ result = fm_proc.process(cv2.cvtColor(face_roi, cv2.COLOR_BGR2RGB))
396
+ fm_proc.close()
397
+ if result.multi_face_landmarks:
398
+ return result.multi_face_landmarks[0]
399
+ return None
400
+ except Exception as e:
401
+ logger.error(f"Face mesh error: {str(e)}")
402
+ return None
403
+
404
+ def draw_face_mesh(image: np.ndarray, face_landmarks, config: Dict, pipeline_config: PipelineConfig):
405
+ mesh_color_bgr = pipeline_config.mesh_color[::-1]
406
+ contour_color_bgr = pipeline_config.contour_color[::-1]
407
+ iris_color_bgr = pipeline_config.iris_color[::-1]
408
+
409
+ if config.get('tesselation'):
410
+ mp_drawing.draw_landmarks(
411
+ image,
412
+ face_landmarks,
413
+ mp_face_mesh.FACEMESH_TESSELATION,
414
+ landmark_drawing_spec=mp_drawing.DrawingSpec(color=mesh_color_bgr, thickness=1, circle_radius=1),
415
+ connection_drawing_spec=mp_drawing.DrawingSpec(color=mesh_color_bgr, thickness=1),
416
+ )
417
+ if config.get('contours'):
418
+ mp_drawing.draw_landmarks(
419
+ image,
420
+ face_landmarks,
421
+ mp_face_mesh.FACEMESH_CONTOURS,
422
+ landmark_drawing_spec=None,
423
+ connection_drawing_spec=mp_drawing.DrawingSpec(color=contour_color_bgr, thickness=2)
424
+ )
425
+ if config.get('irises'):
426
+ mp_drawing.draw_landmarks(
427
+ image,
428
+ face_landmarks,
429
+ mp_face_mesh.FACEMESH_IRISES,
430
+ landmark_drawing_spec=None,
431
+ connection_drawing_spec=mp_drawing.DrawingSpec(color=iris_color_bgr, thickness=2)
432
+ )
433
+
434
+ EYE_COLOR_RANGES = {
435
+ "amber": (255, 191, 0),
436
+ "blue": (0, 0, 255),
437
+ "brown": (139, 69, 19),
438
+ "green": (0, 128, 0),
439
+ "gray": (128, 128, 128),
440
+ "hazel": (102, 51, 0),
441
+ }
442
+
443
+ def classify_eye_color(rgb_color: Tuple[int,int,int]) -> str:
444
+ if rgb_color is None:
445
+ return "Unknown"
446
+ min_dist = float('inf')
447
+ best = "Unknown"
448
+ for color_name, ref_rgb in EYE_COLOR_RANGES.items():
449
+ dist = math.sqrt(sum([(a-b)**2 for a,b in zip(rgb_color, ref_rgb)]))
450
+ if dist < min_dist:
451
+ min_dist = dist
452
+ best = color_name
453
+ return best
454
+
455
+ def get_dominant_color(image_roi, k=3):
456
+ if image_roi.size == 0:
457
+ return None
458
+ pixels = np.float32(image_roi.reshape(-1, 3))
459
+ criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.1)
460
+ _, labels, palette = cv2.kmeans(pixels, k, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
461
+ _, counts = np.unique(labels, return_counts=True)
462
+ dom_color = tuple(palette[np.argmax(counts)].astype(int).tolist())
463
+ return dom_color
464
+
465
+ def detect_eye_color(face_roi: np.ndarray, face_landmarks) -> Optional[str]:
466
+ if face_landmarks is None:
467
+ return None
468
+ h, w = face_roi.shape[:2]
469
+ iris_inds = set()
470
+ for conn in mp_face_mesh.FACEMESH_IRISES:
471
+ iris_inds.update(conn)
472
+
473
+ iris_points = []
474
+ for idx in iris_inds:
475
+ lm = face_landmarks.landmark[idx]
476
+ iris_points.append((int(lm.x * w), int(lm.y * h)))
477
+ if not iris_points:
478
+ return None
479
+
480
+ min_x = min(pt[0] for pt in iris_points)
481
+ max_x = max(pt[0] for pt in iris_points)
482
+ min_y = min(pt[1] for pt in iris_points)
483
+ max_y = max(pt[1] for pt in iris_points)
484
+
485
+ pad = 5
486
+ x1 = max(0, min_x - pad)
487
+ y1 = max(0, min_y - pad)
488
+ x2 = min(w, max_x + pad)
489
+ y2 = min(h, max_y + pad)
490
+
491
+ eye_roi = face_roi[y1:y2, x1:x2]
492
+ eye_roi_resize = cv2.resize(eye_roi, (40, 40), interpolation=cv2.INTER_AREA)
493
+
494
+ if eye_roi_resize.size == 0:
495
+ return None
496
+
497
+ dom_rgb = get_dominant_color(eye_roi_resize)
498
+ if dom_rgb is not None:
499
+ return classify_eye_color(dom_rgb)
500
+ return None
501
+
502
+ class HandTracker:
503
+ def __init__(self, min_detection_confidence=0.5, min_tracking_confidence=0.5):
504
+ self.hands = mp_hands.Hands(
505
+ static_image_mode=True,
506
+ max_num_hands=2,
507
+ min_detection_confidence=min_detection_confidence,
508
+ min_tracking_confidence=min_tracking_confidence,
509
+ )
510
+ logger.info("Initialized Mediapipe HandTracking")
511
+
512
+ def detect_hands(self, image: np.ndarray):
513
+ try:
514
+ img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
515
+ results = self.hands.process(img_rgb)
516
+ return results.multi_hand_landmarks, results.multi_handedness
517
+ except Exception as e:
518
+ logger.error(f"Hand detection error: {str(e)}")
519
+ return None, None
520
+
521
+ def draw_hands(self, image: np.ndarray, hand_landmarks, handedness, config: Dict):
522
+ if not hand_landmarks:
523
+ return image
524
+
525
+ for i, hlms in enumerate(hand_landmarks):
526
+ hl_color = config.hand_landmark_color[::-1]
527
+ hc_color = config.hand_connection_color[::-1]
528
+ mp_drawing.draw_landmarks(
529
+ image,
530
+ hlms,
531
+ mp_hands.HAND_CONNECTIONS,
532
+ mp_drawing.DrawingSpec(color=hl_color, thickness=2, circle_radius=4),
533
+ mp_drawing.DrawingSpec(color=hc_color, thickness=2, circle_radius=2),
534
+ )
535
+ if handedness and i < len(handedness):
536
+ label = handedness[i].classification[0].label
537
+ score = handedness[i].classification[0].score
538
+ text = f"{label}: {score:.2f}"
539
+
540
+ wrist_lm = hlms.landmark[mp_hands.HandLandmark.WRIST]
541
+ h, w_img, _ = image.shape
542
+ cx, cy = int(wrist_lm.x * w_img), int(wrist_lm.y * h)
543
+ ht_color = config.hand_text_color[::-1]
544
+ cv2.putText(image, text, (cx, cy - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, ht_color, 2)
545
+ return image
546
+
547
+ class FacePipeline:
548
+ def __init__(self, config: PipelineConfig):
549
+ self.config = config
550
+ self.detector = None
551
+ self.tracker = None
552
+ self.facenet = None
553
+ self.db = None
554
+ self.hand_tracker = None
555
+ self._initialized = False
556
+
557
+ def initialize(self):
558
+ try:
559
+ self.detector = YOLOFaceDetector(
560
+ model_path=self.config.detector['model_path'],
561
+ device=self.config.detector['device']
562
+ )
563
+ self.tracker = FaceTracker(max_age=self.config.tracker['max_age'])
564
+ self.facenet = FaceNetEmbedder(device=self.config.detector['device'])
565
+ self.db = FaceDatabase()
566
+
567
+ if self.config.hand['enable']:
568
+ self.hand_tracker = HandTracker(
569
+ min_detection_confidence=self.config.hand['min_detection_confidence'],
570
+ min_tracking_confidence=self.config.hand['min_tracking_confidence']
571
+ )
572
+
573
+ self._initialized = True
574
+ logger.info("FacePipeline initialized successfully.")
575
+ except Exception as e:
576
+ logger.error(f"Initialization failed: {str(e)}")
577
+ self._initialized = False
578
+ raise
579
+
580
+ def process_frame(self, frame: np.ndarray) -> Tuple[np.ndarray, List[Dict]]:
581
+ """
582
+ Main pipeline processing: detection, tracking, hand detection, face mesh, blink detection, etc.
583
+ Returns annotated_frame, detection_results.
584
+ """
585
+ if not self._initialized:
586
+ logger.error("Pipeline not initialized.")
587
+ return frame, []
588
+
589
+ try:
590
+ detections = self.detector.detect(frame, self.config.detection_conf_thres)
591
+ tracked_objs = self.tracker.update(detections, frame)
592
+ annotated = frame.copy()
593
+ results = []
594
+
595
+ # Hand detection
596
+ hand_landmarks_list = None
597
+ handedness_list = None
598
+ if self.config.hand['enable'] and self.hand_tracker:
599
+ hand_landmarks_list, handedness_list = self.hand_tracker.detect_hands(annotated)
600
+ annotated = self.hand_tracker.draw_hands(
601
+ annotated, hand_landmarks_list, handedness_list, self.config
602
+ )
603
+
604
+ for obj in tracked_objs:
605
+ if not obj.is_confirmed():
606
+ continue
607
+
608
+ track_id = obj.track_id
609
+ bbox = obj.to_tlbr().astype(int)
610
+ x1, y1, x2, y2 = bbox
611
+ conf = getattr(obj, 'score', 1.0)
612
+ cls = getattr(obj, 'class_id', 0)
613
+
614
+ face_roi = frame[y1:y2, x1:x2]
615
+ if face_roi.size == 0:
616
+ logger.warning(f"Empty face ROI for track={track_id}")
617
+ continue
618
+
619
+ # Anti-spoof
620
+ is_spoofed = False
621
+ if self.config.anti_spoof.get('enable', True):
622
+ is_spoofed = not self.is_real_face(face_roi)
623
+ if is_spoofed:
624
+ cls = 1 # Mark as "spoof"
625
+
626
+ if is_spoofed:
627
+ box_color_bgr = self.config.spoofed_bbox_color[::-1]
628
+ name = "Spoofed"
629
+ similarity = 0.0
630
+ else:
631
+ # Face recognition
632
+ emb = self.facenet.get_embedding(face_roi)
633
+ if emb is not None and self.config.recognition.get('enable', True):
634
+ name, similarity = self.recognize_face(emb, self.config.recognition_conf_thres)
635
+ else:
636
+ name = "Unknown"
637
+ similarity = 0.0
638
+
639
+ box_color_rgb = (self.config.bbox_color if name != "Unknown"
640
+ else self.config.unknown_bbox_color)
641
+ box_color_bgr = box_color_rgb[::-1]
642
+
643
+ label_text = name
644
+ cv2.rectangle(annotated, (x1, y1), (x2, y2), box_color_bgr, 2)
645
+ cv2.putText(annotated, label_text, (x1, y1 - 10),
646
+ cv2.FONT_HERSHEY_SIMPLEX, 0.5, box_color_bgr, 2)
647
+
648
+ # Blink detection
649
+ blink = False
650
+ if self.config.blink.get('enable', False):
651
+ blink, left_ear, right_ear, left_eye_pts, right_eye_pts = detect_blink(
652
+ face_roi, threshold=self.config.blink.get('ear_thresh', 0.25)
653
+ )
654
+ if left_eye_pts is not None and right_eye_pts is not None:
655
+ le_g = left_eye_pts + np.array([x1, y1])
656
+ re_g = right_eye_pts + np.array([x1, y1])
657
+
658
+ eye_outline_bgr = self.config.eye_outline_color[::-1]
659
+ cv2.polylines(annotated, [le_g], True, eye_outline_bgr, 1)
660
+ cv2.polylines(annotated, [re_g], True, eye_outline_bgr, 1)
661
+ if blink:
662
+ blink_msg_color = self.config.blink_text_color[::-1]
663
+ cv2.putText(annotated, "Blink Detected",
664
+ (x1, y2 + 20),
665
+ cv2.FONT_HERSHEY_SIMPLEX, 0.5,
666
+ blink_msg_color, 2)
667
+
668
+ # Face mesh
669
+ face_mesh_landmarks = None
670
+ eye_color_name = None
671
+ if (self.config.face_mesh_options.get('enable') or
672
+ self.config.eye_color.get('enable')):
673
+ face_mesh_landmarks = process_face_mesh(face_roi)
674
+ if face_mesh_landmarks:
675
+ # Draw mesh
676
+ if self.config.face_mesh_options.get('enable', False):
677
+ draw_face_mesh(
678
+ annotated[y1:y2, x1:x2],
679
+ face_mesh_landmarks,
680
+ self.config.face_mesh_options,
681
+ self.config
682
+ )
683
+
684
+ # Eye color detection
685
+ if self.config.eye_color.get('enable', False):
686
+ color_found = detect_eye_color(face_roi, face_mesh_landmarks)
687
+ if color_found:
688
+ eye_color_name = color_found
689
+ text_col_bgr = self.config.eye_color_text_color[::-1]
690
+ cv2.putText(
691
+ annotated, f"Eye Color: {eye_color_name}",
692
+ (x1, y2 + 40),
693
+ cv2.FONT_HERSHEY_SIMPLEX, 0.5,
694
+ text_col_bgr, 2
695
+ )
696
+
697
+ detection_info = {
698
+ "track_id": track_id,
699
+ "bbox": (x1, y1, x2, y2),
700
+ "confidence": float(conf),
701
+ "class_id": cls,
702
+ "name": name,
703
+ "similarity": similarity,
704
+ "blink": blink if self.config.blink.get('enable') else None,
705
+ "face_mesh": bool(face_mesh_landmarks) if self.config.face_mesh_options.get('enable') else False,
706
+ "hands_detected": bool(hand_landmarks_list),
707
+ "hand_count": len(hand_landmarks_list) if hand_landmarks_list else 0,
708
+ "eye_color": eye_color_name if self.config.eye_color.get('enable') else None
709
+ }
710
+ results.append(detection_info)
711
+
712
+ return annotated, results
713
+
714
+ except Exception as e:
715
+ logger.error(f"Frame process error: {str(e)}")
716
+ return frame, []
717
+
718
+ def is_real_face(self, face_roi: np.ndarray) -> bool:
719
+ try:
720
+ gray = cv2.cvtColor(face_roi, cv2.COLOR_BGR2GRAY)
721
+ lapv = cv2.Laplacian(gray, cv2.CV_64F).var()
722
+ return lapv > self.config.anti_spoof.get('lap_thresh', 80.0)
723
+ except Exception as e:
724
+ logger.error(f"Anti-spoof error: {str(e)}")
725
+ return False
726
+
727
+ def recognize_face(self, embedding: np.ndarray, threshold: float) -> Tuple[str, float]:
728
+ try:
729
+ best_name = "Unknown"
730
+ best_sim = 0.0
731
+ for lbl, embs in self.db.embeddings.items():
732
+ for db_emb in embs:
733
+ sim = FacePipeline.cosine_similarity(embedding, db_emb)
734
+ if sim > best_sim:
735
+ best_sim = sim
736
+ best_name = lbl
737
+ if best_sim < threshold:
738
+ best_name = "Unknown"
739
+ return best_name, best_sim
740
+ except Exception as e:
741
+ logger.error(f"Recognition error: {str(e)}")
742
+ return ("Unknown", 0.0)
743
+
744
+ @staticmethod
745
+ def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
746
+ return float(np.dot(a, b) / ((np.linalg.norm(a)*np.linalg.norm(b)) + 1e-6))
747
+
748
+ pipeline = None
749
+ def load_pipeline() -> FacePipeline:
750
+ """Global pipeline loader. Creates if not exists, or returns existing one."""
751
+ global pipeline
752
+ if pipeline is None:
753
+ cfg = PipelineConfig.load(CONFIG_PATH)
754
+ pipeline = FacePipeline(cfg)
755
+ pipeline.initialize()
756
+ return pipeline
757
+
758
+ def hex_to_bgr(hexstr: str) -> Tuple[int,int,int]:
759
+ if not hexstr.startswith('#'):
760
+ hexstr = '#' + hexstr
761
+ h = hexstr.lstrip('#')
762
+ if len(h) != 6:
763
+ return (255, 0, 0)
764
+ r = int(h[0:2], 16)
765
+ g = int(h[2:4], 16)
766
+ b = int(h[4:6], 16)
767
+ return (b,g,r)
768
+
769
+ def bgr_to_hex(bgr: Tuple[int,int,int]) -> str:
770
+ b,g,r = bgr
771
+ return f"#{r:02x}{g:02x}{b:02x}"
772
+
773
+ def update_config(
774
+ enable_recognition, enable_antispoof, enable_blink, enable_hand, enable_eyecolor, enable_facemesh,
775
+ show_tesselation, show_contours, show_irises,
776
+ detection_conf, recognition_thresh, antispoof_thresh, blink_thresh, hand_det_conf, hand_track_conf,
777
+ bbox_hex, spoofed_hex, unknown_hex, eye_hex, blink_hex,
778
+ hand_landmark_hex, hand_connect_hex, hand_text_hex,
779
+ mesh_hex, contour_hex, iris_hex, eye_color_text_hex
780
+ ):
781
+ pl = load_pipeline()
782
+ cfg = pl.config
783
+
784
+ cfg.recognition['enable'] = enable_recognition
785
+ cfg.anti_spoof['enable'] = enable_antispoof
786
+ cfg.blink['enable'] = enable_blink
787
+ cfg.hand['enable'] = enable_hand
788
+ cfg.eye_color['enable'] = enable_eyecolor
789
+ cfg.face_mesh_options['enable'] = enable_facemesh
790
+
791
+ cfg.face_mesh_options['tesselation'] = show_tesselation
792
+ cfg.face_mesh_options['contours'] = show_contours
793
+ cfg.face_mesh_options['irises'] = show_irises
794
+
795
+ cfg.detection_conf_thres = detection_conf
796
+ cfg.recognition_conf_thres = recognition_thresh
797
+ cfg.anti_spoof['lap_thresh'] = antispoof_thresh
798
+ cfg.blink['ear_thresh'] = blink_thresh
799
+ cfg.hand['min_detection_confidence'] = hand_det_conf
800
+ cfg.hand['min_tracking_confidence'] = hand_track_conf
801
+
802
+ cfg.bbox_color = hex_to_bgr(bbox_hex)[::-1]
803
+ cfg.spoofed_bbox_color = hex_to_bgr(spoofed_hex)[::-1]
804
+ cfg.unknown_bbox_color = hex_to_bgr(unknown_hex)[::-1]
805
+ cfg.eye_outline_color = hex_to_bgr(eye_hex)[::-1]
806
+ cfg.blink_text_color = hex_to_bgr(blink_hex)[::-1]
807
+ cfg.hand_landmark_color = hex_to_bgr(hand_landmark_hex)[::-1]
808
+ cfg.hand_connection_color = hex_to_bgr(hand_connect_hex)[::-1]
809
+ cfg.hand_text_color = hex_to_bgr(hand_text_hex)[::-1]
810
+ cfg.mesh_color = hex_to_bgr(mesh_hex)[::-1]
811
+ cfg.contour_color = hex_to_bgr(contour_hex)[::-1]
812
+ cfg.iris_color = hex_to_bgr(iris_hex)[::-1]
813
+ cfg.eye_color_text_color = hex_to_bgr(eye_color_text_hex)[::-1]
814
+
815
+ cfg.save(CONFIG_PATH)
816
+ logger.info("Configuration updated with:") # Added info log
817
+ logger.info(f"Recognition Enabled: {enable_recognition}")
818
+ logger.info(f"Anti-spoof Enabled: {enable_antispoof}")
819
+ logger.info(f"Blink Enabled: {enable_blink}")
820
+ logger.info(f"Face Mesh Enabled: {enable_facemesh}, Tesselation: {show_tesselation}, Contours: {show_contours}, Irises: {show_irises}")
821
+ logger.info(f"Thresholds - Detection Conf: {detection_conf}, Recognition: {recognition_thresh}, Anti-spoof: {antispoof_thresh}, Blink: {blink_thresh}, Hand Det Conf: {hand_det_conf}, Hand Track Conf: {hand_track_conf}")
822
+ logger.info(f"Colors - BBox: {bbox_hex}, Spoofed: {spoofed_hex}, Unknown: {unknown_hex}, Eye Outline: {eye_hex}, Blink Text: {blink_hex}, Hand Landmark: {hand_landmark_hex}, Hand Connect: {hand_connect_hex}, Hand Text: {hand_text_hex}, Mesh: {mesh_hex}, Contour: {contour_hex}, Iris: {iris_hex}, Eye Color Text: {eye_color_text_hex}")
823
+
824
+
825
+ return "Configuration saved successfully!"
826
+
827
+ def enroll_user(label_name: str, files: List[bytes]) -> str:
828
+ """Enrolls a user by name using multiple uploaded image files."""
829
+ pl = load_pipeline()
830
+ if not label_name:
831
+ return "Please provide a user name."
832
+ if not files or len(files) == 0:
833
+ return "No images provided."
834
+
835
+ enrolled_count = 0
836
+ for file_bytes in files:
837
+ if not file_bytes:
838
+ continue
839
+ try:
840
+ img_array = np.frombuffer(file_bytes, np.uint8)
841
+ img_bgr = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
842
+ if img_bgr is None:
843
+ continue
844
+
845
+ dets = pl.detector.detect(img_bgr, pl.config.detection_conf_thres)
846
+ for x1, y1, x2, y2, conf, cls in dets:
847
+ roi = img_bgr[y1:y2, x1:x2]
848
+ if roi.size == 0:
849
+ continue
850
+ emb = pl.facenet.get_embedding(roi)
851
+ if emb is not None:
852
+ pl.db.add_embedding(label_name, emb)
853
+ enrolled_count += 1
854
+ except Exception as e:
855
+ logger.error(f"Error enrolling user from file: {str(e)}")
856
+ continue
857
+
858
+ if enrolled_count > 0:
859
+ pl.db.save()
860
+ return f"Enrolled '{label_name}' with {enrolled_count} face(s)!"
861
+ else:
862
+ return "No faces detected in provided images."
863
+
864
+ def search_by_name(name: str) -> str:
865
+ pl = load_pipeline()
866
+ if not name:
867
+ return "No name entered."
868
+ embs = pl.db.get_embeddings_by_label(name)
869
+ if embs:
870
+ return f"'{name}' found with {len(embs)} embedding(s)."
871
+ else:
872
+ return f"No embeddings found for '{name}'."
873
+
874
+ def search_by_image(img: np.ndarray) -> str:
875
+ pl = load_pipeline()
876
+ if img is None:
877
+ return "No image uploaded."
878
+ img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
879
+ dets = pl.detector.detect(img_bgr, pl.config.detection_conf_thres)
880
+ if not dets:
881
+ return "No faces detected in the uploaded image."
882
+ x1, y1, x2, y2, conf, cls = dets[0]
883
+ roi = img_bgr[y1:y2, x1:x2]
884
+ if roi.size == 0:
885
+ return "Empty face ROI in the uploaded image."
886
+
887
+ emb = pl.facenet.get_embedding(roi)
888
+ if emb is None:
889
+ return "Could not generate embedding from face."
890
+ results = pl.db.search_by_image(emb, pl.config.recognition_conf_thres)
891
+ if not results:
892
+ return "No matches in the database under current threshold."
893
+ lines = [f"- {lbl} (sim={sim:.3f})" for lbl, sim in results]
894
+ return "Search results:\n" + "\n".join(lines)
895
+
896
+ def remove_user(label: str) -> str:
897
+ pl = load_pipeline()
898
+ if not label:
899
+ return "No user label selected."
900
+ pl.db.remove_label(label)
901
+ pl.db.save()
902
+ return f"User '{label}' removed."
903
+
904
+ def list_users() -> str:
905
+ pl = load_pipeline()
906
+ labels = pl.db.list_labels()
907
+ if labels:
908
+ return "Enrolled users:\n" + ", ".join(labels)
909
+ return "No users enrolled."
910
+
911
+ def process_test_image(img: np.ndarray) -> Tuple[np.ndarray, str]:
912
+ """Single-image test: run pipeline and return annotated image + JSON results."""
913
+ if img is None:
914
+ return None, "No image uploaded."
915
+ pl = load_pipeline()
916
+ bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
917
+ processed, detections = pl.process_frame(bgr)
918
+ result_rgb = cv2.cvtColor(processed, cv2.COLOR_BGR2RGB)
919
+ return result_rgb, str(detections)
920
+
921
+ # ===================================
922
+ # Combined Export/Import (Config + DB)
923
+ # ===================================
924
+ def export_all_file() -> str: # Changed return type to str (file path)
925
+ """
926
+ Exports both the pipeline config and database embeddings into a single
927
+ pickle file. Returns the file path for Gradio to handle the download.
928
+ """
929
+ pl = load_pipeline()
930
+ combined_data = {
931
+ "config": pl.config.__dict__,
932
+ "database": pl.db.embeddings
933
+ }
934
+
935
+ # Create an in-memory buffer and pickle the combined data
936
+ buf = io.BytesIO()
937
+ pickle.dump(combined_data, buf)
938
+ buf_bytes = buf.getvalue() # Get bytes from buffer
939
+
940
+ with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as tmp_file:
941
+ tmp_file.write(buf_bytes)
942
+ temp_path = tmp_file.name
943
+ return temp_path # Return the path to the temporary file
944
+
945
+ def import_all_file(file_bytes: bytes, merge_db: bool = True) -> str:
946
+ """
947
+ Imports a single pickle file containing both the config and database.
948
+ If merge_db=False, overwrites the existing DB; otherwise merges.
949
+ """
950
+ if not file_bytes:
951
+ return "No file provided."
952
+
953
+ try:
954
+ # Load the data from the bytes
955
+ buf = io.BytesIO(file_bytes)
956
+ combined_data = pickle.load(buf)
957
+
958
+ if not isinstance(combined_data, dict):
959
+ return "Invalid combined data format."
960
+
961
+ # Rebuild config
962
+ new_cfg_data = combined_data.get("config", {})
963
+ new_cfg = PipelineConfig(**new_cfg_data)
964
+
965
+ # Rebuild DB
966
+ new_db_data = combined_data.get("database", {})
967
+
968
+ # Re-initialize pipeline with new config
969
+ global pipeline
970
+ pipeline = FacePipeline(new_cfg)
971
+ pipeline.initialize()
972
+
973
+ # Merge or overwrite DB
974
+ if merge_db:
975
+ # Merge
976
+ for label, emb_list in new_db_data.items():
977
+ if label not in pipeline.db.embeddings:
978
+ pipeline.db.embeddings[label] = []
979
+ pipeline.db.embeddings[label].extend(emb_list)
980
+ else:
981
+ # Overwrite
982
+ pipeline.db.embeddings = new_db_data
983
+
984
+ pipeline.db.save()
985
+
986
+ return "Config and database imported successfully!"
987
+
988
+ except Exception as e:
989
+ logger.error(f"Import all failed: {str(e)}")
990
+ return f"Import failed: {str(e)}"
991
+
992
+ # ==========================
993
+ # Original Export/Import for
994
+ # Config or DB individually
995
+ # ==========================
996
+
997
+ def export_config_file() -> str: # Changed return type to str (file path)
998
+ """Export the current pipeline config as a downloadable file."""
999
+ pl = load_pipeline()
1000
+ config_bytes = pl.config.export_config()
1001
+ with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as tmp_file:
1002
+ tmp_file.write(config_bytes)
1003
+ temp_path = tmp_file.name
1004
+ return temp_path # Return the path to the temporary file
1005
+
1006
+ def import_config_file(file_bytes: bytes) -> str:
1007
+ """Import a pipeline config from uploaded bytes and re-initialize pipeline."""
1008
+ if not file_bytes:
1009
+ return "No file provided."
1010
+ try:
1011
+ new_cfg = PipelineConfig.import_config(file_bytes)
1012
+ pl = FacePipeline(new_cfg)
1013
+ pl.initialize()
1014
+ global pipeline
1015
+ pipeline = pl
1016
+ return f"Imported config successfully!"
1017
+ except Exception as e:
1018
+ logger.error(f"Import config failed: {str(e)}")
1019
+ return f"Import failed: {str(e)}"
1020
+
1021
+ def export_db_file() -> str: # Changed return type to str (file path)
1022
+ """Export the current face database as a downloadable file."""
1023
+ pl = load_pipeline()
1024
+ db_bytes = pl.db.export_database()
1025
+ with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as tmp_file:
1026
+ tmp_file.write(db_bytes)
1027
+ temp_path = tmp_file.name
1028
+ return temp_path # Return the path to the temporary file
1029
+
1030
+ def import_db_file(db_bytes: bytes, merge: bool=True) -> str:
1031
+ """Import face database from uploaded bytes. Merge or overwrite existing."""
1032
+ if not db_bytes:
1033
+ return "No file provided."
1034
+ try:
1035
+ pl = load_pipeline()
1036
+ pl.db.import_database(db_bytes, merge=merge)
1037
+ return f"Database imported successfully, merge={merge}"
1038
+ except Exception as e:
1039
+ logger.error(f"Import DB failed: {str(e)}")
1040
+ return f"Import DB failed: {str(e)}"
1041
+
1042
+ # Build Gradio App
1043
+ def build_app():
1044
+ with gr.Blocks() as demo:
1045
+ gr.Markdown("# FaceRec: Comprehensive Face Recognition Pipeline")
1046
+ gr.Markdown("**Note:** After downloading, please rename the file to its appropriate extension (e.g., `config_export.pkl`, `database_export.pkl`).")
1047
+
1048
+ with gr.Tab("Image Test"):
1049
+ gr.Markdown("Upload a single image to detect faces, run blink detection, face mesh, hand tracking, etc.")
1050
+ test_in = gr.Image(type="numpy", label="Upload Image")
1051
+ test_out = gr.Image()
1052
+ test_info = gr.Textbox(label="Detections")
1053
+ process_btn = gr.Button("Process Image")
1054
+
1055
+ process_btn.click(
1056
+ fn=process_test_image,
1057
+ inputs=test_in,
1058
+ outputs=[test_out, test_info],
1059
+ )
1060
+
1061
+ with gr.Tab("Configuration"):
1062
+ gr.Markdown("Adjust toggles, thresholds, and colors. Click Save to persist changes.")
1063
+
1064
+ with gr.Row():
1065
+ enable_recognition = gr.Checkbox(label="Enable Recognition", value=True)
1066
+ enable_antispoof = gr.Checkbox(label="Enable Anti-Spoof", value=True)
1067
+ enable_blink = gr.Checkbox(label="Enable Blink Detection", value=True)
1068
+ enable_hand = gr.Checkbox(label="Enable Hand Tracking", value=True)
1069
+ enable_eyecolor = gr.Checkbox(label="Enable Eye Color Detection", value=False)
1070
+ enable_facemesh = gr.Checkbox(label="Enable Face Mesh", value=False)
1071
+
1072
+ gr.Markdown("**Face Mesh Options**")
1073
+ with gr.Row():
1074
+ show_tesselation = gr.Checkbox(label="Tesselation", value=False)
1075
+ show_contours = gr.Checkbox(label="Contours", value=False)
1076
+ show_irises = gr.Checkbox(label="Irises", value=False)
1077
+
1078
+ gr.Markdown("**Thresholds**")
1079
+ detection_conf = gr.Slider(0, 1, 0.4, step=0.01, label="Detection Confidence")
1080
+ recognition_thresh = gr.Slider(0.5, 1.0, 0.85, step=0.01, label="Recognition Threshold")
1081
+ antispoof_thresh = gr.Slider(0, 200, 80, step=1, label="Anti-Spoof Threshold")
1082
+ blink_thresh = gr.Slider(0, 0.5, 0.25, step=0.01, label="Blink EAR Threshold")
1083
+ hand_det_conf = gr.Slider(0, 1, 0.5, step=0.01, label="Hand Detection Confidence")
1084
+ hand_track_conf = gr.Slider(0, 1, 0.5, step=0.01, label="Hand Tracking Confidence")
1085
+
1086
+ gr.Markdown("**Color Options (Hex)**")
1087
+ bbox_hex = gr.Textbox(label="Box Color (Recognized)", value="#00ff00")
1088
+ spoofed_hex = gr.Textbox(label="Box Color (Spoofed)", value="#ff0000")
1089
+ unknown_hex = gr.Textbox(label="Box Color (Unknown)", value="#ff0000")
1090
+ eye_hex = gr.Textbox(label="Eye Outline Color", value="#ffff00")
1091
+ blink_hex = gr.Textbox(label="Blink Text Color", value="#0000ff")
1092
+
1093
+ hand_landmark_hex = gr.Textbox(label="Hand Landmark Color", value="#ffd24d")
1094
+ hand_connect_hex = gr.Textbox(label="Hand Connection Color", value="#cc6600")
1095
+ hand_text_hex = gr.Textbox(label="Hand Text Color", value="#ffffff")
1096
+
1097
+ mesh_hex = gr.Textbox(label="Mesh Color", value="#64ff64")
1098
+ contour_hex = gr.Textbox(label="Contour Color", value="#c8c800")
1099
+ iris_hex = gr.Textbox(label="Iris Color", value="#ff00ff")
1100
+ eye_color_text_hex = gr.Textbox(label="Eye Color Text Color", value="#ffffff")
1101
+
1102
+ save_btn = gr.Button("Save Configuration")
1103
+ save_msg = gr.Textbox(label="", interactive=False)
1104
+
1105
+ save_btn.click(
1106
+ fn=update_config,
1107
+ inputs=[
1108
+ enable_recognition, enable_antispoof, enable_blink, enable_hand, enable_eyecolor, enable_facemesh,
1109
+ show_tesselation, show_contours, show_irises,
1110
+ detection_conf, recognition_thresh, antispoof_thresh, blink_thresh, hand_det_conf, hand_track_conf,
1111
+ bbox_hex, spoofed_hex, unknown_hex, eye_hex, blink_hex,
1112
+ hand_landmark_hex, hand_connect_hex, hand_text_hex,
1113
+ mesh_hex, contour_hex, iris_hex, eye_color_text_hex
1114
+ ],
1115
+ outputs=[save_msg]
1116
+ )
1117
+
1118
+ with gr.Tab("Database Management"):
1119
+ gr.Markdown("Enroll multiple images per user, search by name or image, remove users, list all users.")
1120
+
1121
+ with gr.Accordion("User Enrollment", open=False):
1122
+ enroll_name = gr.Textbox(label="User Name")
1123
+ enroll_paths = gr.File(file_count="multiple", type="binary", label="Upload Multiple Images")
1124
+ enroll_btn = gr.Button("Enroll User")
1125
+ enroll_result = gr.Textbox()
1126
+
1127
+ enroll_btn.click(
1128
+ fn=enroll_user,
1129
+ inputs=[enroll_name, enroll_paths],
1130
+ outputs=[enroll_result]
1131
+ )
1132
+
1133
+ with gr.Accordion("User Search", open=False):
1134
+ search_mode = gr.Radio(["Name", "Image"], label="Search By", value="Name")
1135
+ search_name_box = gr.Dropdown(label="Select User", choices=[], value=None, visible=True)
1136
+ search_image_box = gr.Image(label="Upload Search Image", type="numpy", visible=False)
1137
+ search_btn = gr.Button("Search")
1138
+ search_out = gr.Textbox()
1139
+
1140
+ def toggle_search(mode):
1141
+ if mode == "Name":
1142
+ return gr.update(visible=True), gr.update(visible=False)
1143
+ else:
1144
+ return gr.update(visible=False), gr.update(visible=True)
1145
+
1146
+ search_mode.change(
1147
+ fn=toggle_search,
1148
+ inputs=[search_mode],
1149
+ outputs=[search_name_box, search_image_box]
1150
+ )
1151
+
1152
+ def do_search(mode, uname, img):
1153
+ if mode == "Name":
1154
+ return search_by_name(uname)
1155
+ else:
1156
+ return search_by_image(img)
1157
+
1158
+ search_btn.click(
1159
+ fn=do_search,
1160
+ inputs=[search_mode, search_name_box, search_image_box],
1161
+ outputs=[search_out]
1162
+ )
1163
+
1164
+ with gr.Accordion("User Management Tools", open=False):
1165
+ list_btn = gr.Button("List Enrolled Users")
1166
+ list_out = gr.Textbox()
1167
+ list_btn.click(fn=lambda: list_users(), inputs=[], outputs=[list_out])
1168
+
1169
+ def refresh_choices():
1170
+ pl = load_pipeline()
1171
+ return gr.update(choices=pl.db.list_labels())
1172
+
1173
+ refresh_btn = gr.Button("Refresh User List")
1174
+ refresh_btn.click(fn=refresh_choices, inputs=[], outputs=[search_name_box])
1175
+
1176
+ remove_box = gr.Dropdown(label="Select User to Remove", choices=[])
1177
+ remove_btn = gr.Button("Remove")
1178
+ remove_out = gr.Textbox()
1179
+
1180
+ remove_btn.click(fn=remove_user, inputs=[remove_box], outputs=[remove_out])
1181
+ refresh_btn.click(fn=refresh_choices, inputs=[], outputs=[remove_box])
1182
+
1183
+ with gr.Tab("Export / Import"):
1184
+ gr.Markdown("Export or import pipeline config (thresholds/colors) or face database (embeddings).")
1185
+ gr.Markdown("**Note:** After downloading, please rename the file to its appropriate extension (e.g., `config_export.pkl`, `database_export.pkl`).")
1186
+
1187
+ gr.Markdown("**Export Individually (Download)**")
1188
+ export_config_btn = gr.Button("Export Config")
1189
+ export_config_download = gr.File(label="Download Config Export", type="binary")
1190
+
1191
+ export_db_btn = gr.Button("Export Database")
1192
+ export_db_download = gr.File(label="Download Database Export", type="binary")
1193
+
1194
+ export_config_btn.click(fn=export_config_file, inputs=[], outputs=[export_config_download])
1195
+ export_db_btn.click(fn=export_db_file, inputs=[], outputs=[export_db_download])
1196
+
1197
+ gr.Markdown("**Import Individually (Upload)**")
1198
+ import_config_filebox = gr.File(label="Import Config File", file_count="single", type="binary")
1199
+ import_config_btn = gr.Button("Import Config")
1200
+ import_config_out = gr.Textbox()
1201
+
1202
+ import_db_filebox = gr.File(label="Import Database File", file_count="single", type="binary")
1203
+ merge_db_checkbox = gr.Checkbox(label="Merge instead of overwrite?", value=True)
1204
+ import_db_btn = gr.Button("Import Database")
1205
+ import_db_out = gr.Textbox()
1206
+
1207
+ import_config_btn.click(fn=import_config_file, inputs=[import_config_filebox], outputs=[import_config_out])
1208
+ import_db_btn.click(fn=import_db_file, inputs=[import_db_filebox, merge_db_checkbox], outputs=[import_db_out])
1209
+
1210
+ # =============================
1211
+ # Export/Import All Together
1212
+ # =============================
1213
+ gr.Markdown("---")
1214
+ gr.Markdown("**Export & Import Everything (Config + Database) Together**")
1215
+ gr.Markdown("**Note:** After downloading, please rename the file to `pipeline_export.pkl`.")
1216
+
1217
+ # For exporting: produce a file in-memory
1218
+ export_all_btn = gr.Button("Export All (Config + DB)")
1219
+ export_all_download = gr.File(label="Download Combined Export", type="binary")
1220
+
1221
+ export_all_btn.click(
1222
+ fn=export_all_file, # Now returns file path
1223
+ outputs=[export_all_download],
1224
+ inputs=[]
1225
+ )
1226
+
1227
+ # For importing: user uploads file
1228
+ import_all_in = gr.File(label="Import Combined File (Pickle)", file_count="single", type="binary")
1229
+ import_all_merge_cb = gr.Checkbox(label="Merge DB instead of overwrite?", value=True)
1230
+ import_all_btn = gr.Button("Import All")
1231
+ import_all_out = gr.Textbox()
1232
+
1233
+ import_all_btn.click(
1234
+ fn=import_all_file,
1235
+ inputs=[import_all_in, import_all_merge_cb],
1236
+ outputs=[import_all_out]
1237
+ )
1238
+
1239
+ return demo
1240
+
1241
+ def main():
1242
+ """Entry point to launch the Gradio app."""
1243
+ app = build_app()
1244
+ # We add `.queue()` so that multiple requests can be queued
1245
+ app.queue().launch(server_name="0.0.0.0", server_port=7860)
1246
+
1247
  if __name__ == "__main__":
1248
+ main()