You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
104 lines
3.3 KiB
104 lines
3.3 KiB
# -*- coding:utf-8 -*-
|
|
"""
|
|
@Time : 2023/5/4 4:03 PM
|
|
@Author: hxl
|
|
@Des: websocket
|
|
"""
|
|
import asyncio
|
|
import json
|
|
|
|
import cv2
|
|
from fastapi import APIRouter, WebSocketDisconnect
|
|
from starlette.websockets import WebSocket
|
|
|
|
from arcsoft.arcface import FaceRecognitionEngine
|
|
from helper.websocket_manage import ConnectionManager
|
|
from models.user import User
|
|
|
|
router = APIRouter()
|
|
|
|
manager = ConnectionManager()
|
|
|
|
|
|
class FaceCameraEngine:
|
|
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
|
|
frame = None
|
|
|
|
# 描绘人脸
|
|
async def detect_faces(self, img):
|
|
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)
|
|
if len(faces) == 0:
|
|
return img
|
|
for (x, y, w, h) in faces:
|
|
cv2.rectangle(img, (x, y), (x + w, y + h), (255, 0, 0), 2)
|
|
return img
|
|
|
|
# 开始推流
|
|
async def face_recoWork(self):
|
|
cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
|
|
while len(manager.active_connections):
|
|
_, frame = cap.read()
|
|
self.frame = frame
|
|
frame = await self.detect_faces(frame)
|
|
_, jpeg = cv2.imencode('.jpg', frame)
|
|
await manager.broadcast_stream(jpeg.tobytes())
|
|
await asyncio.sleep(0.1)
|
|
cap.release()
|
|
|
|
async def register(self, user_id: int):
|
|
while len(manager.active_connections):
|
|
with FaceRecognitionEngine() as frc:
|
|
dat = frc.face_registration(self.frame)
|
|
if dat:
|
|
user = await User.get_or_none(id=user_id)
|
|
if user:
|
|
user.face = dat
|
|
await user.save()
|
|
result = {"code": 200, "desc": "success"}
|
|
await manager.broadcast_json("face", result)
|
|
await asyncio.sleep(2)
|
|
|
|
async def login(self):
|
|
dat_list = await User.filter(face__isnull=False, locked=False).values_list('id', 'face')
|
|
while len(manager.active_connections):
|
|
with FaceRecognitionEngine() as frc:
|
|
user_id = frc.face_login(self.frame, dat_list)
|
|
if user_id:
|
|
user = await User.get_or_none(id=user_id)
|
|
result = {"code": 200, "data": user.profile}
|
|
await manager.broadcast_json("face", json.dumps(result))
|
|
await asyncio.sleep(2)
|
|
|
|
|
|
@router.websocket("/login")
|
|
async def face_login(websocket: WebSocket):
|
|
await manager.connect(websocket)
|
|
if len(manager.active_connections) == 1:
|
|
fc = FaceCameraEngine()
|
|
task1 = asyncio.create_task(fc.face_recoWork())
|
|
task2 = asyncio.create_task(fc.login())
|
|
await task1, task2
|
|
try:
|
|
while True:
|
|
await websocket.receive_bytes()
|
|
except WebSocketDisconnect:
|
|
manager.disconnect(websocket)
|
|
|
|
|
|
@router.websocket("/register/{id}")
|
|
async def face_login(websocket: WebSocket):
|
|
await manager.connect(websocket)
|
|
if len(manager.active_connections) == 1:
|
|
user_id = websocket.path_params.get('id')
|
|
fc = FaceCameraEngine()
|
|
task1 = asyncio.create_task(fc.face_recoWork())
|
|
task2 = asyncio.create_task(fc.register(user_id))
|
|
await task1, task2
|
|
|
|
try:
|
|
while True:
|
|
await websocket.receive_bytes()
|
|
except WebSocketDisconnect:
|
|
manager.disconnect(websocket)
|