server.py
3.83 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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import torch
import torch.multiprocessing as mp
import numpy as np
import os
import asyncio
import json
import base64
import websockets
from io import BytesIO
from PIL import Image, ImageDraw
from IPython import display
from models.mtcnn import MTCNN
from models.inception_resnet_v1 import InceptionResnetV1
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Running on device: {}'.format(device))
model = InceptionResnetV1().eval().to(device)
lock = asyncio.Lock()
clients = set()
#processes = []
async def get_embeddings(face_list):
global model
x = torch.Tensor(face_list).to(device)
yhat = model(x)
return yhat
def get_distance(someone, database):
distance = [(someone - data).norm().item() for data in database]
return distance
def get_argmin(someone, database):
distance = get_distance(someone, database)
for i in range(len(distance)):
return np.argmin(distance)
return -1
async def register(websocket):
global lock
global clients
async with lock:
clients.add(websocket)
remote_ip = websocket.remote_address[0]
msg='[{ip}] connected'.format(ip=remote_ip)
print(msg)
async def unregister(websocket):
global lock
global clients
async with lock:
clients.remove(websocket)
remote_ip = websocket.remote_address[0]
msg='[{ip}] disconnected'.format(ip=remote_ip)
print(msg)
async def thread(websocket, path):
# register(websocket) sends user_event() to websocket
await register(websocket)
try:
# await websocket.send(state_event())
async for message in websocket:
data = json.loads(message)
if data['action'] == 'register':
# data['id']
face = np.asarray(data['MTCNN'], dtype = np.float32)
face = face.reshape((1,3,160,160))
remote_ip = websocket.remote_address[0]
msg='[{ip}] register face'.format(ip=remote_ip)
print(msg)
embedding = await get_embeddings(face)
await websocket.send('registered')
#await notify_state()
elif data['action'] == "verify":
face = np.asarray(data['MTCNN'], dtype = np.float32)
print(face.shape)
face = face.reshape((1,3,160,160))
remote_ip = websocket.remote_address[0]
msg='[{ip}] verify face'.format(ip=remote_ip)
print(msg)
embedding = await get_embeddings(face)
# Todo: 아래 embedding.numpy()를 데이터베이스에 저장해야함.
# embedding.numpy()
# [1, 512] numpy()임
# np.bytes() 명령으로 바꾼 뒤 np.frombuffer()로 불러오는 것이 좋을 듯.
await websocket.send('정해갑')
elif data['action'] == "save_image":
# 출석이 제대로 이뤄지지 않으면 이미지를 저장하여
# 나중에 교강사가 출석을 확인할 수 있도록 한다
arr = np.asarray(data['image'], dtype = np.uint8)
# 이 데이터는 데이터베이스(과목명/일자/undefined)에 저장하는 것이 좋을듯
# image = Image.fromarray(arr)
# image.save('face.jpg')# storage에 데이터 저장
remote_ip = websocket.remote_address[0]
msg='[{ip}] save image'.format(ip=remote_ip)
print(msg)
###
await websocket.send('정해갑')
else:
print("unsupported event: {}", data)
finally:
await unregister(websocket)
print('run verification server')
start_server = websockets.serve(thread, '0.0.0.0', 8765)
asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()