face enrollment added

server.py

@@ -18,6 +18,8 @@ import math
 from pydantic import BaseModel, Field
 from pymongo import AsyncMongoClient
 import bcrypt
+import pickle
+from bson import ObjectId
 from jose import JWTError, jwt
 from dotenv import load_dotenv
 from pathlib import Path
@@ -509,6 +511,41 @@ async def verify_token(current_user: UserPublic = Depends(get_current_user)):
     return current_user
 
 
+@app.post("/api/enroll_face")
+async def enroll_face(
+    video: UploadFile = File(...),
+    current_user: UserPublic = Depends(get_current_user)
+):
+    try:
+        temp_path = f"temp_enroll_{uuid.uuid4()}.mp4"
+        with open(temp_path, "wb") as buffer:
+            shutil.copyfileobj(video.file, buffer)
+            
+        logger.info(f"Extracting embeddings for user {current_user.id}")
+        
+        def run_extraction():
+            return face_service.extract_embeddings_from_video(temp_path)
+            
+        embeddings, num_frames = await asyncio.get_event_loop().run_in_executor(
+            executor, run_extraction
+        )
+        
+        pickled_embeddings = pickle.dumps(embeddings)
+        await users_collection.update_one(
+            {"_id": ObjectId(current_user.id)},
+            {"$set": {"embeddings": pickled_embeddings}}
+        )
+        
+        os.remove(temp_path)
+        
+        return {"ok": True, "message": "Face enrolled successfully", "frames_used": num_frames}
+    except Exception as e:
+        logger.error(f"Enrollment failed: {e}")
+        if os.path.exists(temp_path):
+            os.remove(temp_path)
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
     global is_recording, video_writer, recording_filename, latest_obs_frame, is_obs_active, zoom_multiplier
@@ -517,6 +554,7 @@ async def websocket_endpoint(websocket: WebSocket):
     logger.info("New WebSocket connection established.")
 
     current_mode = "single"  # Default mode
+    ws_user_embeddings = None
 
     try:
         while True:
@@ -530,6 +568,21 @@ async def websocket_endpoint(websocket: WebSocket):
                         logger.info(f"Switching mode from {current_mode} to {payload['mode']}")
                         current_mode = payload["mode"]
                         await websocket.send_json({"type": "mode_ack", "mode": current_mode})
+                    elif "type" in payload and payload["type"] == "auth":
+                        token = payload.get("token")
+                        try:
+                            token_data = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+                            user_id = token_data.get("sub")
+                            if user_id:
+                                user = await users_collection.find_one({"_id": ObjectId(user_id)})
+                                if user and "embeddings" in user and user["embeddings"]:
+                                    ws_user_embeddings = pickle.loads(user["embeddings"])
+                                    logger.info(f"Loaded custom face embeddings for user {user_id}")
+                                    await websocket.send_json({"type": "auth_ack", "status": "enrolled"})
+                                else:
+                                    await websocket.send_json({"type": "auth_ack", "status": "no_enrollment"})
+                        except Exception as e:
+                            logger.error(f"WS Auth failed: {e}")
                     elif "zoom_scale" in payload:
                         zoom_multiplier = float(payload["zoom_scale"])
                         logger.info(f"Updated zoom multiplier to {zoom_multiplier}")
@@ -572,9 +625,9 @@ async def websocket_endpoint(websocket: WebSocket):
                     continue
 
                 # Prepare inference function
-                def run_inference(f, mode):
+                def run_inference(f, mode, embeddings=None):
                     if mode == "single":
-                        return single_tracker.process_frame(f)
+                        return single_tracker.process_frame(f, custom_embeddings=embeddings)
                     elif mode == "multi":
                         return multi_tracker.process_frame(f)
                     else:
@@ -584,7 +637,7 @@ async def websocket_endpoint(websocket: WebSocket):
                 response_data = {}
                 try:
                     response_data = await asyncio.get_event_loop().run_in_executor(
-                        executor, run_inference, frame, current_mode
+                        executor, run_inference, frame, current_mode, ws_user_embeddings
                     )
                 except Exception as e:
                     logger.error(f"Error processing frame in {current_mode} mode: {e}")

services/single_tracker.py

@@ -73,7 +73,7 @@ class SingleTracker:
         if not cache_loaded:
             logger.warning(f"Cache invalid or not found at {self.cache_file}. Returning empty embeddings. Please run Model/face_model.py to generate cache.")
         
-    def process_frame(self, frame):
+    def process_frame(self, frame, custom_embeddings=None):
         """
         Process a single BGR image frame for single face tracking.
         Returns a dictionary with tracking results.
@@ -97,14 +97,45 @@ class SingleTracker:
             if results and len(results) > 0 and results[0].boxes.id is not None:
                 boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)
                 track_ids = results[0].boxes.id.cpu().numpy().astype(int)
+                
+                keypoints = None
+                if hasattr(results[0], 'keypoints') and results[0].keypoints is not None:
+                    keypoints = results[0].keypoints.xy.cpu().numpy()
 
-                for box, track_id in zip(boxes, track_ids):
+                for idx, (box, track_id) in enumerate(zip(boxes, track_ids)):
                     x1, y1, x2, y2 = box.tolist()
                     track_id = int(track_id)
                     max_similarity = 0.0
                     
+                    # Compute Head Pose
+                    yaw = 0.0
+                    pitch = 0.0
+                    if keypoints is not None and len(keypoints) > idx:
+                        kpts = keypoints[idx]
+                        if len(kpts) >= 5:
+                            lex, ley = kpts[0]
+                            rex, rey = kpts[1]
+                            nx, ny = kpts[2]
+                            lmx, lmy = kpts[3]
+                            rmx, rmy = kpts[4]
+                            
+                            # Yaw: (-) turned left, (+) turned right
+                            l_nose = abs(nx - lex)
+                            r_nose = abs(nx - rex)
+                            yaw = (l_nose - r_nose) / (l_nose + r_nose + 1e-6)
+                            
+                            # Pitch: (-) looking up, (+) looking down
+                            eye_cy = (ley + rey) / 2
+                            mouth_cy = (lmy + rmy) / 2
+                            n_eye = ny - eye_cy
+                            n_mouth = mouth_cy - ny
+                            pitch = (n_eye - n_mouth) / (n_eye + n_mouth + 1e-6)
+                            
                     # Lock resolution logic
-                    if track_id not in self.known_tracks and len(self.main_user_embeddings) > 0:
+                    
+                    embeddings_to_check = custom_embeddings if custom_embeddings is not None and len(custom_embeddings) > 0 else self.main_user_embeddings
+                    
+                    if track_id not in self.known_tracks and len(embeddings_to_check) > 0:
                         if track_id not in self.track_retries:
                             self.track_retries[track_id] = 0
 
@@ -115,7 +146,9 @@ class SingleTracker:
                             # Strict check
                             current_face = DeepFace.represent(face_crop, model_name=self.model_name, enforce_detection=False)[0]["embedding"]
                             
-                            for user_embedding in self.main_user_embeddings:
+                            embeddings_to_check = custom_embeddings if custom_embeddings is not None and len(custom_embeddings) > 0 else self.main_user_embeddings
+                            
+                            for user_embedding in embeddings_to_check:
                                 sim = np.dot(user_embedding, current_face) / (np.linalg.norm(user_embedding) * np.linalg.norm(current_face))
                                 if sim > max_similarity:
                                     max_similarity = sim
@@ -156,7 +189,9 @@ class SingleTracker:
                             "x2": x2, "y2": y2,
                             "is_target": True,
                             "label": label,
-                            "similarity": max_similarity if 'max_similarity' in locals() else -1.0
+                            "similarity": max_similarity if 'max_similarity' in locals() else -1.0,
+                            "yaw": float(yaw),
+                            "pitch": float(pitch)
                         })
                     elif track_id in self.track_retries:
                         # Draw scanning box
@@ -167,7 +202,9 @@ class SingleTracker:
                             "x2": x2, "y2": y2,
                             "is_target": False,
                             "label": label,
-                            "similarity": max_similarity if 'max_similarity' in locals() else -1.0
+                            "similarity": max_similarity if 'max_similarity' in locals() else -1.0,
+                            "yaw": float(yaw),
+                            "pitch": float(pitch)
                         })
                     
         except Exception as e: