Spaces:
Sleeping
Sleeping
| import json | |
| import base64 | |
| import cv2 | |
| import numpy as np | |
| import asyncio | |
| from fastapi import WebSocket, WebSocketDisconnect | |
| from services.vision import process_frame_synchronous, get_face_embedding | |
| from services.attendance import mark_attendance | |
| from services.faiss_manager import db | |
| # Temporary memory for active registration sessions | |
| # { websocket: { "name": "...", "embeddings": [] } } | |
| registration_sessions = {} | |
| async def websocket_endpoint(websocket: WebSocket): | |
| await websocket.accept() | |
| try: | |
| while True: | |
| data = await websocket.receive_text() | |
| payload = json.loads(data) | |
| if payload.get("type") == "heartbeat": | |
| await websocket.send_json({"type": "heartbeat_ack"}) | |
| continue | |
| # --- REGISTRATION FLOW --- | |
| if payload.get("type") == "start_registration": | |
| name = payload.get("name") | |
| registration_sessions[websocket] = {"name": name, "embeddings": []} | |
| await websocket.send_json({"type": "registration_status", "message": f"Started scanning for {name}", "progress": 0}) | |
| continue | |
| if payload.get("type") == "register_frame": | |
| if websocket not in registration_sessions: | |
| continue | |
| encoded_data = payload["image"].split(',')[1] | |
| nparr = np.frombuffer(base64.b64decode(encoded_data), np.uint8) | |
| frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR) | |
| # Extract embedding (will return None if blurry or no face) | |
| emb = await asyncio.to_thread(get_face_embedding, frame) | |
| if emb is not None: | |
| registration_sessions[websocket]["embeddings"].append(emb) | |
| count = len(registration_sessions[websocket]["embeddings"]) | |
| await websocket.send_json({ | |
| "type": "registration_status", | |
| "message": f"Captured {count} high-quality frames", | |
| "progress": count | |
| }) | |
| else: | |
| await websocket.send_json({ | |
| "type": "registration_status", | |
| "message": "Face not found or too blurry. Please hold still.", | |
| "progress": len(registration_sessions[websocket]["embeddings"]) | |
| }) | |
| continue | |
| if payload.get("type") == "finish_registration": | |
| if websocket in registration_sessions: | |
| session = registration_sessions[websocket] | |
| embs = session["embeddings"] | |
| if len(embs) > 0: | |
| # MULTI-ANGLE REGISTRATION: Average all captured embeddings | |
| avg_emb = np.mean(embs, axis=0) | |
| # Re-normalize to ensure Cosine Similarity works correctly | |
| avg_emb = avg_emb / np.linalg.norm(avg_emb) | |
| # Save to FAISS | |
| db.add_identity(session["name"], avg_emb) | |
| await websocket.send_json({"type": "registration_success", "message": f"Successfully registered {session['name']}!"}) | |
| else: | |
| await websocket.send_json({"type": "registration_error", "message": "Failed. No valid frames captured."}) | |
| del registration_sessions[websocket] | |
| continue | |
| # --- STANDARD RECOGNITION FLOW --- | |
| if payload.get("type") == "frame": | |
| encoded_data = payload["image"].split(',')[1] | |
| nparr = np.frombuffer(base64.b64decode(encoded_data), np.uint8) | |
| frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR) | |
| report, stats = await asyncio.to_thread(process_frame_synchronous, frame) | |
| results_summary = [] | |
| client_faces = [] | |
| for face in report: | |
| name = face["name"] | |
| score = face["score"] | |
| status = face["status"] | |
| client_faces.append({ | |
| "name": name, | |
| "score": score, | |
| "status": status, | |
| "box": face["box"], | |
| "crop": face["crop_b64"] | |
| }) | |
| if status == "match": | |
| # --- 🧠 SELF-HEALING AI (Continuous Learning) --- | |
| # If we are highly confident (>95%), we add this new angle/lighting to their profile! | |
| if score >= 95.0 and "embedding" in face: | |
| db.add_identity(name, face["embedding"]) | |
| status_db, time_str = mark_attendance(name) | |
| results_summary.append(f"✅ {name} ({score}%)") | |
| await websocket.send_json({ | |
| "type": "attendance", | |
| "name": name, | |
| "time": time_str or "Just Now", | |
| "status": "success" | |
| }) | |
| elif status == "verifying": | |
| results_summary.append(f"⏳ {name}") | |
| else: | |
| results_summary.append(f"❌ {name} ({score}%)") | |
| debug_msg = f"Faces: {stats['detected']} | " | |
| debug_msg += " | ".join(results_summary) if results_summary else "No faces found." | |
| await websocket.send_json({ | |
| "type": "ready", | |
| "debug": debug_msg, | |
| "faces": client_faces | |
| }) | |
| except WebSocketDisconnect: | |
| if websocket in registration_sessions: | |
| del registration_sessions[websocket] | |
| except Exception as e: | |
| print(f"WebSocket Error: {e}") | |