register.py
2.64 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
##################################################
#1. webcam에서 얼굴을 인식합니다
#2. 인식한 얼굴을 등록합니다
##################################################
import torch
import numpy as np
import cv2
import asyncio
import websockets
import json
import os
import timeit
import base64
from PIL import Image
from io import BytesIO
import requests
from models.mtcnn import MTCNN
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Running on device: {}'.format(device))
mtcnn = MTCNN(keep_all=True, device=device)
uri = 'ws://169.56.95.131:8765'
async def send_face(face_list, image_list):
global uri
async with websockets.connect(uri) as websocket:
for face, image in zip(face_list, image_list):
#type: np.float32
send = json.dumps({'action': 'register', 'student_id':'2014104149', 'student_name':'정해갑', 'MTCNN': face.tolist()})
await websocket.send(send)
recv = await websocket.recv()
data = json.loads(recv)
if data['status'] == 'success':
# 성공
print(data['student_id'], 'is registered')
def detect_face(frame):
# If required, create a face detection pipeline using MTCNN:
global mtcnn
results = mtcnn.detect(frame)
image_list = []
if results[1][0] == None:
return []
for box, prob in zip(results[0], results[1]):
if prob < 0.95:
continue
print('face detected. prob:', prob)
x1, y1, x2, y2 = box
image = frame[int(y1-10):int(y2+10), int(x1-10):int(x2+10)]
image_list.append(image)
return image_list
def detect_face(frame):
results = mtcnn.detect(frame)
faces = mtcnn(frame, return_prob = False)
image_list = []
face_list = []
if results[1][0] == None:
return [], []
for box, face, prob in zip(results[0], faces, results[1]):
if prob < 0.97:
continue
print('face detected. prob:', prob)
x1, y1, x2, y2 = box
if (x2-x1) * (y2-y1) < 15000:
# 얼굴 해상도가 너무 낮으면 무시
continue
# 얼굴 주변 ±3 영역 저장
image = frame[int(y1-3):int(y2+3), int(x1-3):int(x2+3)]
image_list.append(image)
# MTCNN 데이터 저장
face_list.append(face.numpy())
return image_list, face_list
cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
cap.set(3, 720)
cap.set(4, 480)
ret, frame = cap.read()
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
face_list, image_list = detect_face(frame)
if face_list:
asyncio.get_event_loop().run_until_complete(send_face(face_list, image_list))