import os import sys import math import requests import numpy as np import cv2 import torch import pickle import logging from PIL import Image from typing import Optional, Dict, List, Tuple from dataclasses import dataclass, field from collections import Counter import io import tempfile # Import tempfile import gradio as gr from ultralytics import YOLO from facenet_pytorch import InceptionResnetV1 from torchvision import transforms from deep_sort_realtime.deepsort_tracker import DeepSort import mediapipe as mp # Configure logging logging.basicConfig( level=logging.DEBUG, # Changed to DEBUG for more detailed logs format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', handlers=[logging.FileHandler('face_pipeline.log'), logging.StreamHandler()], ) logger = logging.getLogger(__name__) # Suppress verbose logs from libraries logging.getLogger('torch').setLevel(logging.ERROR) logging.getLogger('mediapipe').setLevel(logging.ERROR) logging.getLogger('deep_sort_realtime').setLevel(logging.ERROR) # Constants and default paths 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") # Mediapipe indices for eye landmarks LEFT_EYE_IDX = [33, 160, 158, 133, 153, 144] RIGHT_EYE_IDX = [263, 387, 385, 362, 380, 373] # Initialize Mediapipe drawing utilities mp_drawing = mp.solutions.drawing_utils mp_face_mesh = mp.solutions.face_mesh mp_hands = mp.solutions.hands @dataclass class PipelineConfig: detector: Dict = field(default_factory=dict) tracker: Dict = field(default_factory=dict) recognition: Dict = field(default_factory=dict) anti_spoof: Dict = field(default_factory=dict) blink: Dict = field(default_factory=dict) face_mesh_options: Dict = field(default_factory=dict) hand: Dict = field(default_factory=dict) eye_color: Dict = field(default_factory=dict) enabled_components: Dict = field(default_factory=dict) detection_conf_thres: float = 0.4 recognition_conf_thres: float = 0.85 bbox_color: Tuple[int, int, int] = (0, 255, 0) spoofed_bbox_color: Tuple[int, int, int] = (0, 0, 255) unknown_bbox_color: Tuple[int, int, int] = (0, 0, 255) eye_outline_color: Tuple[int, int, int] = (255, 255, 0) blink_text_color: Tuple[int, int, int] = (0, 0, 255) hand_landmark_color: Tuple[int, int, int] = (255, 210, 77) hand_connection_color: Tuple[int, int, int] = (204, 102, 0) hand_text_color: Tuple[int, int, int] = (255, 255, 255) mesh_color: Tuple[int, int, int] = (100, 255, 100) contour_color: Tuple[int, int, int] = (200, 200, 0) iris_color: Tuple[int, int, int] = (255, 0, 255) eye_color_text_color: Tuple[int, int, int] = (255, 255, 255) def __post_init__(self): self.detector = self.detector or { 'model_path': os.path.join(MODEL_DIR, "face2.pt"), 'device': 'cuda' if torch.cuda.is_available() else 'cpu', } self.tracker = self.tracker or {'max_age': 30} self.recognition = self.recognition or {'enable': True} self.anti_spoof = self.anti_spoof or {'enable': True, 'lap_thresh': 80.0} self.blink = self.blink or {'enable': True, 'ear_thresh': 0.25} self.face_mesh_options = self.face_mesh_options or { 'enable': False, 'tesselation': False, 'contours': False, 'irises': False, } self.hand = self.hand or { 'enable': True, 'min_detection_confidence': 0.5, 'min_tracking_confidence': 0.5, } self.eye_color = self.eye_color or {'enable': False} self.enabled_components = self.enabled_components or { 'detection': True, 'tracking': True, 'anti_spoof': True, 'recognition': True, 'blink': True, 'face_mesh': False, 'hand': True, 'eye_color': False, } def save(self, path: str): """Save this config to a pickle file.""" try: os.makedirs(os.path.dirname(path), exist_ok=True) with open(path, 'wb') as f: pickle.dump(self.__dict__, f) logger.info(f"Saved config to {path}") logger.debug(f"Config data saved: {self.__dict__}") # Added debug log except Exception as e: logger.error(f"Config save failed: {str(e)}") raise RuntimeError(f"Config save failed: {str(e)}") from e @classmethod def load(cls, path: str) -> 'PipelineConfig': """Load a config from a pickle file.""" try: if os.path.exists(path): with open(path, 'rb') as f: data = pickle.load(f) logger.info(f"Loaded config from {path}") logger.debug(f"Config data loaded: {data}") # Added debug log return cls(**data) logger.info("No config file found, using default config.") # Added log for default case return cls() except Exception as e: logger.error(f"Config load failed: {str(e)}") return cls() def export_config(self) -> bytes: """Export your config to bytes.""" try: config_data = self.__dict__ buf = io.BytesIO() pickle.dump(config_data, buf) buf.seek(0) return buf.read() except Exception as e: logger.error(f"Export config failed: {str(e)}") raise RuntimeError(f"Export config failed: {str(e)}") from e @classmethod def import_config(cls, config_bytes: bytes) -> 'PipelineConfig': """Import config from bytes.""" try: buf = io.BytesIO(config_bytes) data = pickle.load(buf) return cls(**data) except Exception as e: logger.error(f"Import config failed: {str(e)}") raise RuntimeError(f"Import config failed: {str(e)}") from e class FaceDatabase: def __init__(self, db_path: str = DEFAULT_DB_PATH): self.db_path = db_path self.embeddings: Dict[str, List[np.ndarray]] = {} self._load() def _load(self): try: if os.path.exists(self.db_path): with open(self.db_path, 'rb') as f: self.embeddings = pickle.load(f) logger.info(f"Loaded database from {self.db_path}") except Exception as e: logger.error(f"Database load failed: {str(e)}") self.embeddings = {} def save(self): try: os.makedirs(os.path.dirname(self.db_path), exist_ok=True) with open(self.db_path, 'wb') as f: pickle.dump(self.embeddings, f) logger.info(f"Saved database to {self.db_path}") except Exception as e: logger.error(f"Database save failed: {str(e)}") raise RuntimeError(f"Database save failed: {str(e)}") from e def export_database(self) -> bytes: """Export the entire face embeddings DB to bytes.""" try: db_data = self.embeddings buf = io.BytesIO() pickle.dump(db_data, buf) buf.seek(0) return buf.read() except Exception as e: logger.error(f"Export database failed: {str(e)}") raise RuntimeError(f"Export database failed: {str(e)}") from e def import_database(self, db_bytes: bytes, merge: bool = True): """ Import embeddings from bytes. If merge=True, merges with current DB. If False, overwrites. """ try: buf = io.BytesIO(db_bytes) imported_data = pickle.load(buf) if not isinstance(imported_data, dict): raise ValueError("Imported data is not a dictionary!") if merge: for label, emb_list in imported_data.items(): if label not in self.embeddings: self.embeddings[label] = [] self.embeddings[label].extend(emb_list) else: self.embeddings = imported_data self.save() logger.info(f"Imported face database, merge={merge}") except Exception as e: logger.error(f"Import database failed: {str(e)}") raise RuntimeError(f"Import database failed: {str(e)}") from e def add_embedding(self, label: str, embedding: np.ndarray): try: if not isinstance(embedding, np.ndarray) or embedding.ndim != 1: raise ValueError("Invalid embedding format") if label not in self.embeddings: self.embeddings[label] = [] self.embeddings[label].append(embedding) logger.debug(f"Added embedding for {label}") except Exception as e: logger.error(f"Add embedding failed: {str(e)}") raise def remove_label(self, label: str): try: if label in self.embeddings: del self.embeddings[label] logger.info(f"Removed {label}") else: logger.warning(f"Label {label} not found") except Exception as e: logger.error(f"Remove label failed: {str(e)}") raise def list_labels(self) -> List[str]: return list(self.embeddings.keys()) def get_embeddings_by_label(self, label: str) -> Optional[List[np.ndarray]]: return self.embeddings.get(label) def search_by_image(self, query_embedding: np.ndarray, threshold: float = 0.7) -> List[Tuple[str, float]]: results = [] for lbl, embs in self.embeddings.items(): for db_emb in embs: sim = FacePipeline.cosine_similarity(query_embedding, db_emb) if sim >= threshold: results.append((lbl, sim)) return sorted(results, key=lambda x: x[1], reverse=True) class YOLOFaceDetector: def __init__(self, model_path: str, device: str = 'cpu'): self.model = None self.device = device try: if not os.path.exists(model_path): logger.info(f"Model not found at {model_path}. Downloading from GitHub...") resp = requests.get(DEFAULT_MODEL_URL) resp.raise_for_status() os.makedirs(os.path.dirname(model_path), exist_ok=True) with open(model_path, 'wb') as f: f.write(resp.content) logger.info(f"Downloaded YOLO model to {model_path}") self.model = YOLO(model_path) self.model.to(device) logger.info(f"Loaded YOLO model from {model_path}") except Exception as e: logger.error(f"YOLO init failed: {str(e)}") raise def detect(self, image: np.ndarray, conf_thres: float) -> List[Tuple[int, int, int, int, float, int]]: try: results = self.model.predict( source=image, conf=conf_thres, verbose=False, device=self.device ) detections = [] for result in results: for box in result.boxes: x1, y1, x2, y2 = box.xyxy[0].cpu().numpy() conf = float(box.conf[0].cpu().numpy()) cls = int(box.cls[0].cpu().numpy()) if box.cls is not None else 0 detections.append((int(x1), int(y1), int(x2), int(y2), conf, cls)) logger.debug(f"Detected {len(detections)} faces.") return detections except Exception as e: logger.error(f"Detection error: {str(e)}") return [] class FaceTracker: def __init__(self, max_age: int = 30): self.tracker = DeepSort(max_age=max_age, embedder='mobilenet') def update(self, detections: List[Tuple], frame: np.ndarray): try: ds_detections = [ ([x1, y1, x2 - x1, y2 - y1], conf, cls) for (x1, y1, x2, y2, conf, cls) in detections ] tracks = self.tracker.update_tracks(ds_detections, frame=frame) logger.debug(f"Updated tracker with {len(tracks)} tracks.") return tracks except Exception as e: logger.error(f"Tracking error: {str(e)}") return [] class FaceNetEmbedder: def __init__(self, device: str = 'cpu'): self.device = device self.model = InceptionResnetV1(pretrained='vggface2').eval().to(device) self.transform = transforms.Compose([ transforms.Resize((160, 160)), transforms.ToTensor(), transforms.Normalize([0.5]*3, [0.5]*3), ]) def get_embedding(self, face_bgr: np.ndarray) -> Optional[np.ndarray]: try: face_rgb = cv2.cvtColor(face_bgr, cv2.COLOR_BGR2RGB) pil_img = Image.fromarray(face_rgb).convert('RGB') tens = self.transform(pil_img).unsqueeze(0).to(self.device) with torch.no_grad(): embedding = self.model(tens)[0].cpu().numpy() logger.debug(f"Generated embedding sample: {embedding[:5]}...") return embedding except Exception as e: logger.error(f"Embedding failed: {str(e)}") return None def detect_blink(face_roi: np.ndarray, threshold: float = 0.25) -> Tuple[bool, float, float, Optional[np.ndarray], Optional[np.ndarray]]: """ Returns: (blink_bool, left_ear, right_ear, left_eye_points, right_eye_points). """ try: face_mesh_proc = mp_face_mesh.FaceMesh( static_image_mode=True, max_num_faces=1, refine_landmarks=True, min_detection_confidence=0.5 ) result = face_mesh_proc.process(cv2.cvtColor(face_roi, cv2.COLOR_BGR2RGB)) face_mesh_proc.close() if not result.multi_face_landmarks: return False, 0.0, 0.0, None, None landmarks = result.multi_face_landmarks[0].landmark h, w = face_roi.shape[:2] def eye_aspect_ratio(indices): pts = [(landmarks[i].x * w, landmarks[i].y * h) for i in indices] vertical = np.linalg.norm(np.array(pts[1]) - np.array(pts[5])) + \ np.linalg.norm(np.array(pts[2]) - np.array(pts[4])) horizontal = np.linalg.norm(np.array(pts[0]) - np.array(pts[3])) return vertical / (2.0 * horizontal + 1e-6) left_ear = eye_aspect_ratio(LEFT_EYE_IDX) right_ear = eye_aspect_ratio(RIGHT_EYE_IDX) blink = (left_ear < threshold) and (right_ear < threshold) left_eye_pts = np.array([(int(landmarks[i].x * w), int(landmarks[i].y * h)) for i in LEFT_EYE_IDX]) right_eye_pts = np.array([(int(landmarks[i].x * w), int(landmarks[i].y * h)) for i in RIGHT_EYE_IDX]) return blink, left_ear, right_ear, left_eye_pts, right_eye_pts except Exception as e: logger.error(f"Blink detection error: {str(e)}") return False, 0.0, 0.0, None, None def process_face_mesh(face_roi: np.ndarray): try: fm_proc = mp_face_mesh.FaceMesh( static_image_mode=True, max_num_faces=1, refine_landmarks=True, min_detection_confidence=0.5 ) result = fm_proc.process(cv2.cvtColor(face_roi, cv2.COLOR_BGR2RGB)) fm_proc.close() if result.multi_face_landmarks: return result.multi_face_landmarks[0] return None except Exception as e: logger.error(f"Face mesh error: {str(e)}") return None def draw_face_mesh(image: np.ndarray, face_landmarks, config: Dict, pipeline_config: PipelineConfig): mesh_color_bgr = pipeline_config.mesh_color[::-1] contour_color_bgr = pipeline_config.contour_color[::-1] iris_color_bgr = pipeline_config.iris_color[::-1] if config.get('tesselation'): mp_drawing.draw_landmarks( image, face_landmarks, mp_face_mesh.FACEMESH_TESSELATION, landmark_drawing_spec=mp_drawing.DrawingSpec(color=mesh_color_bgr, thickness=1, circle_radius=1), connection_drawing_spec=mp_drawing.DrawingSpec(color=mesh_color_bgr, thickness=1), ) if config.get('contours'): mp_drawing.draw_landmarks( image, face_landmarks, mp_face_mesh.FACEMESH_CONTOURS, landmark_drawing_spec=None, connection_drawing_spec=mp_drawing.DrawingSpec(color=contour_color_bgr, thickness=2) ) if config.get('irises'): mp_drawing.draw_landmarks( image, face_landmarks, mp_face_mesh.FACEMESH_IRISES, landmark_drawing_spec=None, connection_drawing_spec=mp_drawing.DrawingSpec(color=iris_color_bgr, thickness=2) ) EYE_COLOR_RANGES = { "amber": (255, 191, 0), "blue": (0, 0, 255), "brown": (139, 69, 19), "green": (0, 128, 0), "gray": (128, 128, 128), "hazel": (102, 51, 0), } def classify_eye_color(rgb_color: Tuple[int,int,int]) -> str: if rgb_color is None: return "Unknown" min_dist = float('inf') best = "Unknown" for color_name, ref_rgb in EYE_COLOR_RANGES.items(): dist = math.sqrt(sum([(a-b)**2 for a,b in zip(rgb_color, ref_rgb)])) if dist < min_dist: min_dist = dist best = color_name return best def get_dominant_color(image_roi, k=3): if image_roi.size == 0: return None pixels = np.float32(image_roi.reshape(-1, 3)) criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.1) _, labels, palette = cv2.kmeans(pixels, k, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS) _, counts = np.unique(labels, return_counts=True) dom_color = tuple(palette[np.argmax(counts)].astype(int).tolist()) return dom_color def detect_eye_color(face_roi: np.ndarray, face_landmarks) -> Optional[str]: if face_landmarks is None: return None h, w = face_roi.shape[:2] iris_inds = set() for conn in mp_face_mesh.FACEMESH_IRISES: iris_inds.update(conn) iris_points = [] for idx in iris_inds: lm = face_landmarks.landmark[idx] iris_points.append((int(lm.x * w), int(lm.y * h))) if not iris_points: return None min_x = min(pt[0] for pt in iris_points) max_x = max(pt[0] for pt in iris_points) min_y = min(pt[1] for pt in iris_points) max_y = max(pt[1] for pt in iris_points) pad = 5 x1 = max(0, min_x - pad) y1 = max(0, min_y - pad) x2 = min(w, max_x + pad) y2 = min(h, max_y + pad) eye_roi = face_roi[y1:y2, x1:x2] eye_roi_resize = cv2.resize(eye_roi, (40, 40), interpolation=cv2.INTER_AREA) if eye_roi_resize.size == 0: return None dom_rgb = get_dominant_color(eye_roi_resize) if dom_rgb is not None: return classify_eye_color(dom_rgb) return None class HandTracker: def __init__(self, min_detection_confidence=0.5, min_tracking_confidence=0.5): self.hands = mp_hands.Hands( static_image_mode=True, max_num_hands=2, min_detection_confidence=min_detection_confidence, min_tracking_confidence=min_tracking_confidence, ) logger.info("Initialized Mediapipe HandTracking") def detect_hands(self, image: np.ndarray): try: img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) results = self.hands.process(img_rgb) return results.multi_hand_landmarks, results.multi_handedness except Exception as e: logger.error(f"Hand detection error: {str(e)}") return None, None def draw_hands(self, image: np.ndarray, hand_landmarks, handedness, config: Dict): if not hand_landmarks: return image for i, hlms in enumerate(hand_landmarks): hl_color = config.hand_landmark_color[::-1] hc_color = config.hand_connection_color[::-1] mp_drawing.draw_landmarks( image, hlms, mp_hands.HAND_CONNECTIONS, mp_drawing.DrawingSpec(color=hl_color, thickness=2, circle_radius=4), mp_drawing.DrawingSpec(color=hc_color, thickness=2, circle_radius=2), ) if handedness and i < len(handedness): label = handedness[i].classification[0].label score = handedness[i].classification[0].score text = f"{label}: {score:.2f}" wrist_lm = hlms.landmark[mp_hands.HandLandmark.WRIST] h, w_img, _ = image.shape cx, cy = int(wrist_lm.x * w_img), int(wrist_lm.y * h) ht_color = config.hand_text_color[::-1] cv2.putText(image, text, (cx, cy - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, ht_color, 2) return image class FacePipeline: def __init__(self, config: PipelineConfig): self.config = config self.detector = None self.tracker = None self.facenet = None self.db = None self.hand_tracker = None self._initialized = False def initialize(self): try: self.detector = YOLOFaceDetector( model_path=self.config.detector['model_path'], device=self.config.detector['device'] ) self.tracker = FaceTracker(max_age=self.config.tracker['max_age']) self.facenet = FaceNetEmbedder(device=self.config.detector['device']) self.db = FaceDatabase() if self.config.hand['enable']: self.hand_tracker = HandTracker( min_detection_confidence=self.config.hand['min_detection_confidence'], min_tracking_confidence=self.config.hand['min_tracking_confidence'] ) self._initialized = True logger.info("FacePipeline initialized successfully.") except Exception as e: logger.error(f"Initialization failed: {str(e)}") self._initialized = False raise def process_frame(self, frame: np.ndarray) -> Tuple[np.ndarray, List[Dict]]: """ Main pipeline processing: detection, tracking, hand detection, face mesh, blink detection, etc. Returns annotated_frame, detection_results. """ if not self._initialized: logger.error("Pipeline not initialized.") return frame, [] try: detections = self.detector.detect(frame, self.config.detection_conf_thres) tracked_objs = self.tracker.update(detections, frame) annotated = frame.copy() results = [] # Hand detection hand_landmarks_list = None handedness_list = None if self.config.hand['enable'] and self.hand_tracker: hand_landmarks_list, handedness_list = self.hand_tracker.detect_hands(annotated) annotated = self.hand_tracker.draw_hands( annotated, hand_landmarks_list, handedness_list, self.config ) for obj in tracked_objs: if not obj.is_confirmed(): continue track_id = obj.track_id bbox = obj.to_tlbr().astype(int) x1, y1, x2, y2 = bbox conf = getattr(obj, 'score', 1.0) cls = getattr(obj, 'class_id', 0) face_roi = frame[y1:y2, x1:x2] if face_roi.size == 0: logger.warning(f"Empty face ROI for track={track_id}") continue # Anti-spoof is_spoofed = False if self.config.anti_spoof.get('enable', True): is_spoofed = not self.is_real_face(face_roi) if is_spoofed: cls = 1 # Mark as "spoof" if is_spoofed: box_color_bgr = self.config.spoofed_bbox_color[::-1] name = "Spoofed" similarity = 0.0 else: # Face recognition emb = self.facenet.get_embedding(face_roi) if emb is not None and self.config.recognition.get('enable', True): name, similarity = self.recognize_face(emb, self.config.recognition_conf_thres) else: name = "Unknown" similarity = 0.0 box_color_rgb = (self.config.bbox_color if name != "Unknown" else self.config.unknown_bbox_color) box_color_bgr = box_color_rgb[::-1] label_text = name cv2.rectangle(annotated, (x1, y1), (x2, y2), box_color_bgr, 2) cv2.putText(annotated, label_text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, box_color_bgr, 2) # Blink detection blink = False if self.config.blink.get('enable', False): blink, left_ear, right_ear, left_eye_pts, right_eye_pts = detect_blink( face_roi, threshold=self.config.blink.get('ear_thresh', 0.25) ) if left_eye_pts is not None and right_eye_pts is not None: le_g = left_eye_pts + np.array([x1, y1]) re_g = right_eye_pts + np.array([x1, y1]) eye_outline_bgr = self.config.eye_outline_color[::-1] cv2.polylines(annotated, [le_g], True, eye_outline_bgr, 1) cv2.polylines(annotated, [re_g], True, eye_outline_bgr, 1) if blink: blink_msg_color = self.config.blink_text_color[::-1] cv2.putText(annotated, "Blink Detected", (x1, y2 + 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, blink_msg_color, 2) # Face mesh face_mesh_landmarks = None eye_color_name = None if (self.config.face_mesh_options.get('enable') or self.config.eye_color.get('enable')): face_mesh_landmarks = process_face_mesh(face_roi) if face_mesh_landmarks: # Draw mesh if self.config.face_mesh_options.get('enable', False): draw_face_mesh( annotated[y1:y2, x1:x2], face_mesh_landmarks, self.config.face_mesh_options, self.config ) # Eye color detection if self.config.eye_color.get('enable', False): color_found = detect_eye_color(face_roi, face_mesh_landmarks) if color_found: eye_color_name = color_found text_col_bgr = self.config.eye_color_text_color[::-1] cv2.putText( annotated, f"Eye Color: {eye_color_name}", (x1, y2 + 40), cv2.FONT_HERSHEY_SIMPLEX, 0.5, text_col_bgr, 2 ) detection_info = { "track_id": track_id, "bbox": (x1, y1, x2, y2), "confidence": float(conf), "class_id": cls, "name": name, "similarity": similarity, "blink": blink if self.config.blink.get('enable') else None, "face_mesh": bool(face_mesh_landmarks) if self.config.face_mesh_options.get('enable') else False, "hands_detected": bool(hand_landmarks_list), "hand_count": len(hand_landmarks_list) if hand_landmarks_list else 0, "eye_color": eye_color_name if self.config.eye_color.get('enable') else None } results.append(detection_info) return annotated, results except Exception as e: logger.error(f"Frame process error: {str(e)}") return frame, [] def is_real_face(self, face_roi: np.ndarray) -> bool: try: gray = cv2.cvtColor(face_roi, cv2.COLOR_BGR2GRAY) lapv = cv2.Laplacian(gray, cv2.CV_64F).var() return lapv > self.config.anti_spoof.get('lap_thresh', 80.0) except Exception as e: logger.error(f"Anti-spoof error: {str(e)}") return False def recognize_face(self, embedding: np.ndarray, threshold: float) -> Tuple[str, float]: try: best_name = "Unknown" best_sim = 0.0 for lbl, embs in self.db.embeddings.items(): for db_emb in embs: sim = FacePipeline.cosine_similarity(embedding, db_emb) if sim > best_sim: best_sim = sim best_name = lbl if best_sim < threshold: best_name = "Unknown" return best_name, best_sim except Exception as e: logger.error(f"Recognition error: {str(e)}") return ("Unknown", 0.0) @staticmethod def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float: return float(np.dot(a, b) / ((np.linalg.norm(a)*np.linalg.norm(b)) + 1e-6)) pipeline = None def load_pipeline() -> FacePipeline: """Global pipeline loader. Creates if not exists, or returns existing one.""" global pipeline if pipeline is None: cfg = PipelineConfig.load(CONFIG_PATH) pipeline = FacePipeline(cfg) pipeline.initialize() return pipeline def hex_to_bgr(hexstr: str) -> Tuple[int,int,int]: if not hexstr.startswith('#'): hexstr = '#' + hexstr h = hexstr.lstrip('#') if len(h) != 6: return (255, 0, 0) r = int(h[0:2], 16) g = int(h[2:4], 16) b = int(h[4:6], 16) return (b,g,r) def bgr_to_hex(bgr: Tuple[int,int,int]) -> str: b,g,r = bgr return f"#{r:02x}{g:02x}{b:02x}" def update_config( 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, bbox_hex, spoofed_hex, unknown_hex, eye_hex, blink_hex, hand_landmark_hex, hand_connect_hex, hand_text_hex, mesh_hex, contour_hex, iris_hex, eye_color_text_hex ): pl = load_pipeline() cfg = pl.config cfg.recognition['enable'] = enable_recognition cfg.anti_spoof['enable'] = enable_antispoof cfg.blink['enable'] = enable_blink cfg.hand['enable'] = enable_hand cfg.eye_color['enable'] = enable_eyecolor cfg.face_mesh_options['enable'] = enable_facemesh cfg.face_mesh_options['tesselation'] = show_tesselation cfg.face_mesh_options['contours'] = show_contours cfg.face_mesh_options['irises'] = show_irises cfg.detection_conf_thres = detection_conf cfg.recognition_conf_thres = recognition_thresh cfg.anti_spoof['lap_thresh'] = antispoof_thresh cfg.blink['ear_thresh'] = blink_thresh cfg.hand['min_detection_confidence'] = hand_det_conf cfg.hand['min_tracking_confidence'] = hand_track_conf 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] cfg.blink_text_color = hex_to_bgr(blink_hex)[::-1] cfg.hand_landmark_color = hex_to_bgr(hand_landmark_hex)[::-1] cfg.hand_connection_color = hex_to_bgr(hand_connect_hex)[::-1] cfg.hand_text_color = hex_to_bgr(hand_text_hex)[::-1] cfg.mesh_color = hex_to_bgr(mesh_hex)[::-1] cfg.contour_color = hex_to_bgr(contour_hex)[::-1] cfg.iris_color = hex_to_bgr(iris_hex)[::-1] cfg.eye_color_text_color = hex_to_bgr(eye_color_text_hex)[::-1] cfg.save(CONFIG_PATH) logger.info("Configuration updated with:") # Added info log logger.info(f"Recognition Enabled: {enable_recognition}") logger.info(f"Anti-spoof Enabled: {enable_antispoof}") logger.info(f"Blink Enabled: {enable_blink}") logger.info(f"Face Mesh Enabled: {enable_facemesh}, Tesselation: {show_tesselation}, Contours: {show_contours}, Irises: {show_irises}") 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}") 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}") return "Configuration saved successfully!" def enroll_user(label_name: str, files: List[bytes]) -> str: """Enrolls a user by name using multiple uploaded image files.""" pl = load_pipeline() if not label_name: return "Please provide a user name." if not files or len(files) == 0: return "No images provided." enrolled_count = 0 for file_bytes in files: if not file_bytes: continue try: img_array = np.frombuffer(file_bytes, np.uint8) img_bgr = cv2.imdecode(img_array, cv2.IMREAD_COLOR) if img_bgr is None: continue dets = pl.detector.detect(img_bgr, pl.config.detection_conf_thres) for x1, y1, x2, y2, conf, cls in dets: roi = img_bgr[y1:y2, x1:x2] if roi.size == 0: continue emb = pl.facenet.get_embedding(roi) if emb is not None: pl.db.add_embedding(label_name, emb) enrolled_count += 1 except Exception as e: logger.error(f"Error enrolling user from file: {str(e)}") continue if enrolled_count > 0: pl.db.save() return f"Enrolled '{label_name}' with {enrolled_count} face(s)!" else: return "No faces detected in provided images." def search_by_name(name: str) -> str: pl = load_pipeline() if not name: return "No name entered." embs = pl.db.get_embeddings_by_label(name) if embs: return f"'{name}' found with {len(embs)} embedding(s)." else: return f"No embeddings found for '{name}'." def search_by_image(img: np.ndarray) -> str: pl = load_pipeline() if img is None: return "No image uploaded." img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) dets = pl.detector.detect(img_bgr, pl.config.detection_conf_thres) if not dets: return "No faces detected in the uploaded image." x1, y1, x2, y2, conf, cls = dets[0] roi = img_bgr[y1:y2, x1:x2] if roi.size == 0: return "Empty face ROI in the uploaded image." emb = pl.facenet.get_embedding(roi) if emb is None: return "Could not generate embedding from face." results = pl.db.search_by_image(emb, pl.config.recognition_conf_thres) if not results: return "No matches in the database under current threshold." lines = [f"- {lbl} (sim={sim:.3f})" for lbl, sim in results] return "Search results:\n" + "\n".join(lines) def remove_user(label: str) -> str: pl = load_pipeline() if not label: return "No user label selected." pl.db.remove_label(label) pl.db.save() return f"User '{label}' removed." def list_users() -> str: pl = load_pipeline() labels = pl.db.list_labels() if labels: return "Enrolled users:\n" + ", ".join(labels) return "No users enrolled." def process_test_image(img: np.ndarray) -> Tuple[np.ndarray, str]: """Single-image test: run pipeline and return annotated image + JSON results.""" if img is None: return None, "No image uploaded." pl = load_pipeline() bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) processed, detections = pl.process_frame(bgr) result_rgb = cv2.cvtColor(processed, cv2.COLOR_BGR2RGB) return result_rgb, str(detections) # =================================== # Combined Export/Import (Config + DB) # =================================== def export_all_file() -> str: # Changed return type to str (file path) """ Exports both the pipeline config and database embeddings into a single pickle file. Returns the file path for Gradio to handle the download. """ pl = load_pipeline() combined_data = { "config": pl.config.__dict__, "database": pl.db.embeddings } # Create an in-memory buffer and pickle the combined data buf = io.BytesIO() pickle.dump(combined_data, buf) buf_bytes = buf.getvalue() # Get bytes from buffer with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as tmp_file: tmp_file.write(buf_bytes) temp_path = tmp_file.name return temp_path # Return the path to the temporary file def import_all_file(file_bytes: bytes, merge_db: bool = True) -> str: """ Imports a single pickle file containing both the config and database. If merge_db=False, overwrites the existing DB; otherwise merges. """ if not file_bytes: return "No file provided." try: # Load the data from the bytes buf = io.BytesIO(file_bytes) combined_data = pickle.load(buf) if not isinstance(combined_data, dict): return "Invalid combined data format." # Rebuild config new_cfg_data = combined_data.get("config", {}) new_cfg = PipelineConfig(**new_cfg_data) # Rebuild DB new_db_data = combined_data.get("database", {}) # Re-initialize pipeline with new config global pipeline pipeline = FacePipeline(new_cfg) pipeline.initialize() # Merge or overwrite DB if merge_db: # Merge for label, emb_list in new_db_data.items(): if label not in pipeline.db.embeddings: pipeline.db.embeddings[label] = [] pipeline.db.embeddings[label].extend(emb_list) else: # Overwrite pipeline.db.embeddings = new_db_data pipeline.db.save() return "Config and database imported successfully!" except Exception as e: logger.error(f"Import all failed: {str(e)}") return f"Import failed: {str(e)}" # ========================== # Original Export/Import for # Config or DB individually # ========================== def export_config_file() -> str: # Changed return type to str (file path) """Export the current pipeline config as a downloadable file.""" pl = load_pipeline() config_bytes = pl.config.export_config() with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as tmp_file: tmp_file.write(config_bytes) temp_path = tmp_file.name return temp_path # Return the path to the temporary file def import_config_file(file_bytes: bytes) -> str: """Import a pipeline config from uploaded bytes and re-initialize pipeline.""" if not file_bytes: return "No file provided." try: new_cfg = PipelineConfig.import_config(file_bytes) pl = FacePipeline(new_cfg) pl.initialize() global pipeline pipeline = pl return f"Imported config successfully!" except Exception as e: logger.error(f"Import config failed: {str(e)}") return f"Import failed: {str(e)}" def export_db_file() -> str: # Changed return type to str (file path) """Export the current face database as a downloadable file.""" pl = load_pipeline() db_bytes = pl.db.export_database() with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as tmp_file: tmp_file.write(db_bytes) temp_path = tmp_file.name return temp_path # Return the path to the temporary file def import_db_file(db_bytes: bytes, merge: bool=True) -> str: """Import face database from uploaded bytes. Merge or overwrite existing.""" if not db_bytes: return "No file provided." try: pl = load_pipeline() pl.db.import_database(db_bytes, merge=merge) return f"Database imported successfully, merge={merge}" except Exception as e: logger.error(f"Import DB failed: {str(e)}") return f"Import DB failed: {str(e)}" # Build Gradio App def build_app(): with gr.Blocks() as demo: gr.Markdown("# FaceRec: Comprehensive Face Recognition Pipeline") gr.Markdown("**Note:** After downloading, please rename the file to its appropriate extension (e.g., `config_export.pkl`, `database_export.pkl`).") with gr.Tab("Image Test"): gr.Markdown("Upload a single image to detect faces, run blink detection, face mesh, hand tracking, etc.") test_in = gr.Image(type="numpy", label="Upload Image") test_out = gr.Image() test_info = gr.Textbox(label="Detections") process_btn = gr.Button("Process Image") process_btn.click( fn=process_test_image, inputs=test_in, outputs=[test_out, test_info], ) with gr.Tab("Configuration"): gr.Markdown("Adjust toggles, thresholds, and colors. Click Save to persist changes.") with gr.Row(): enable_recognition = gr.Checkbox(label="Enable Recognition", value=True) enable_antispoof = gr.Checkbox(label="Enable Anti-Spoof", value=True) enable_blink = gr.Checkbox(label="Enable Blink Detection", value=True) enable_hand = gr.Checkbox(label="Enable Hand Tracking", value=True) 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**") with gr.Row(): show_tesselation = gr.Checkbox(label="Tesselation", value=False) show_contours = gr.Checkbox(label="Contours", value=False) show_irises = gr.Checkbox(label="Irises", value=False) gr.Markdown("**Thresholds**") detection_conf = gr.Slider(0, 1, 0.4, step=0.01, label="Detection Confidence") recognition_thresh = gr.Slider(0.5, 1.0, 0.85, step=0.01, label="Recognition Threshold") antispoof_thresh = gr.Slider(0, 200, 80, step=1, label="Anti-Spoof Threshold") blink_thresh = gr.Slider(0, 0.5, 0.25, step=0.01, label="Blink EAR Threshold") hand_det_conf = gr.Slider(0, 1, 0.5, step=0.01, label="Hand Detection Confidence") hand_track_conf = gr.Slider(0, 1, 0.5, step=0.01, label="Hand Tracking Confidence") gr.Markdown("**Color Options (Hex)**") bbox_hex = gr.Textbox(label="Box Color (Recognized)", value="#00ff00") spoofed_hex = gr.Textbox(label="Box Color (Spoofed)", value="#ff0000") unknown_hex = gr.Textbox(label="Box Color (Unknown)", value="#ff0000") eye_hex = gr.Textbox(label="Eye Outline Color", value="#ffff00") blink_hex = gr.Textbox(label="Blink Text Color", value="#0000ff") hand_landmark_hex = gr.Textbox(label="Hand Landmark Color", value="#ffd24d") hand_connect_hex = gr.Textbox(label="Hand Connection Color", value="#cc6600") hand_text_hex = gr.Textbox(label="Hand Text Color", value="#ffffff") 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") save_msg = gr.Textbox(label="", interactive=False) save_btn.click( fn=update_config, inputs=[ 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, bbox_hex, spoofed_hex, unknown_hex, eye_hex, blink_hex, hand_landmark_hex, hand_connect_hex, hand_text_hex, mesh_hex, contour_hex, iris_hex, eye_color_text_hex ], outputs=[save_msg] ) with gr.Tab("Database Management"): gr.Markdown("Enroll multiple images per user, search by name or image, remove users, list all users.") with gr.Accordion("User Enrollment", open=False): enroll_name = gr.Textbox(label="User Name") enroll_paths = gr.File(file_count="multiple", type="binary", label="Upload Multiple Images") enroll_btn = gr.Button("Enroll User") enroll_result = gr.Textbox() enroll_btn.click( fn=enroll_user, inputs=[enroll_name, enroll_paths], outputs=[enroll_result] ) with gr.Accordion("User Search", open=False): search_mode = gr.Radio(["Name", "Image"], label="Search By", value="Name") search_name_box = gr.Dropdown(label="Select User", choices=[], value=None, visible=True) search_image_box = gr.Image(label="Upload Search Image", type="numpy", visible=False) search_btn = gr.Button("Search") search_out = gr.Textbox() def toggle_search(mode): 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=toggle_search, inputs=[search_mode], outputs=[search_name_box, search_image_box] ) def do_search(mode, uname, img): if mode == "Name": return search_by_name(uname) else: return search_by_image(img) search_btn.click( fn=do_search, inputs=[search_mode, search_name_box, search_image_box], outputs=[search_out] ) with gr.Accordion("User Management Tools", open=False): list_btn = gr.Button("List Enrolled Users") list_out = gr.Textbox() list_btn.click(fn=lambda: list_users(), inputs=[], outputs=[list_out]) def refresh_choices(): pl = load_pipeline() return gr.update(choices=pl.db.list_labels()) refresh_btn = gr.Button("Refresh User List") refresh_btn.click(fn=refresh_choices, inputs=[], outputs=[search_name_box]) remove_box = gr.Dropdown(label="Select User to Remove", choices=[]) remove_btn = gr.Button("Remove") remove_out = gr.Textbox() remove_btn.click(fn=remove_user, inputs=[remove_box], outputs=[remove_out]) refresh_btn.click(fn=refresh_choices, inputs=[], outputs=[remove_box]) with gr.Tab("Export / Import"): gr.Markdown("Export or import pipeline config (thresholds/colors) or face database (embeddings).") gr.Markdown("**Note:** After downloading, please rename the file to its appropriate extension (e.g., `config_export.pkl`, `database_export.pkl`).") gr.Markdown("**Export Individually (Download)**") export_config_btn = gr.Button("Export Config") export_config_download = gr.File(label="Download Config Export", type="binary") export_db_btn = gr.Button("Export Database") export_db_download = gr.File(label="Download Database Export", type="binary") export_config_btn.click(fn=export_config_file, inputs=[], outputs=[export_config_download]) export_db_btn.click(fn=export_db_file, inputs=[], outputs=[export_db_download]) gr.Markdown("**Import Individually (Upload)**") import_config_filebox = gr.File(label="Import Config File", file_count="single", type="binary") import_config_btn = gr.Button("Import Config") import_config_out = gr.Textbox() import_db_filebox = gr.File(label="Import Database File", file_count="single", type="binary") merge_db_checkbox = gr.Checkbox(label="Merge instead of overwrite?", value=True) import_db_btn = gr.Button("Import Database") import_db_out = gr.Textbox() import_config_btn.click(fn=import_config_file, inputs=[import_config_filebox], outputs=[import_config_out]) import_db_btn.click(fn=import_db_file, inputs=[import_db_filebox, merge_db_checkbox], outputs=[import_db_out]) # ============================= # Export/Import All Together # ============================= gr.Markdown("---") gr.Markdown("**Export & Import Everything (Config + Database) Together**") gr.Markdown("**Note:** After downloading, please rename the file to `pipeline_export.pkl`.") # For exporting: produce a file in-memory export_all_btn = gr.Button("Export All (Config + DB)") export_all_download = gr.File(label="Download Combined Export", type="binary") export_all_btn.click( fn=export_all_file, # Now returns file path outputs=[export_all_download], inputs=[] ) # For importing: user uploads file import_all_in = gr.File(label="Import Combined File (Pickle)", file_count="single", type="binary") import_all_merge_cb = gr.Checkbox(label="Merge DB instead of overwrite?", value=True) import_all_btn = gr.Button("Import All") import_all_out = gr.Textbox() import_all_btn.click( fn=import_all_file, inputs=[import_all_in, import_all_merge_cb], outputs=[import_all_out] ) return demo def main(): """Entry point to launch the Gradio app.""" app = build_app() # We add `.queue()` so that multiple requests can be queued app.queue().launch(server_name="0.0.0.0", server_port=7860) if __name__ == "__main__": main()