Graduate

Connect DB

import os
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
import pymysql
import datetime
from PIL import Image, ImageDraw
from IPython import display
......@@ -18,6 +20,13 @@ device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Running on device: {}'.format(device))
model = InceptionResnetV1().eval().to(device)
attendance_db = pymysql.connect(
user='root',
passwd='5978',
host='localhost',
db='attendance',
charset='utf8'
)
lock = asyncio.Lock()
clients = set()
......@@ -29,15 +38,15 @@ async def get_embeddings(face_list):
yhat = model(x)
return yhat
def get_distance(someone, database):
distance = [(someone - data).norm().item() for data in database]
return distance
async def get_distance(arr1, arr2):
distance = (arr1 - arr2).norm().item()
return distance
def get_argmin(someone, database):
distance = get_distance(someone, database)
for i in range(len(distance)):
return np.argmin(distance)
return -1
distance = get_distance(someone, database)
for i in range(len(distance)):
return np.argmin(distance)
return -1
async def register(websocket):
global lock
......@@ -64,29 +73,83 @@ async def thread(websocket, path):
# await websocket.send(state_event())
async for message in websocket:
data = json.loads(message)
remote_ip = websocket.remote_address[0]
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]
# log
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":
# load json
student_id = data['student_id']
student_name = data['student_name']
face = np.asarray(data['MTCNN'], dtype = np.float32)
print(face.shape)
face = face.reshape((1,3,160,160))
remote_ip = websocket.remote_address[0]
# DB에 연결
cursor = attendance_db.cursor(pymysql.cursors.DictCursor)
# 학생을 찾음
sql = "SELECT student_id FROM student WHERE student_id = %s;"
cursor.execute(sql, (student_id))
# DB에 학생이 없으면 등록
if not cursor.fetchone():
sql = "insert into student(student_id, student_name) values (%s, %s)"
cursor.execute(sql, (student_id, student_name))
attendance_db.commit()
# student_embedding Table에 등록
embedding = await get_embeddings(face)
embedding = embedding.detach().numpy().tobytes()
embedding_date = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
sql = "insert into student_embedding(student_id, embedding_date, embedding) values (%s, %s, %s)"
cursor.execute(sql, (student_id, embedding_date, embedding))
attendance_db.commit()
await websocket.send('{id} registered'.format(id=student_id))
elif data['action'] == "verify":
# log
msg='[{ip}] verify face'.format(ip=remote_ip)
print(msg)
###############
# load json
face = np.asarray(data['MTCNN'], dtype = np.float32)
face = face.reshape((1,3,160,160))
# embedding 구하기
embedding = await get_embeddings(face)
# Todo: 아래 embedding.numpy()를 데이터베이스에 저장해야함.
embedding = embedding.detach().numpy()
# embedding.numpy()
# [1, 512] numpy()임
# np.bytes() 명령으로 바꾼 뒤 np.frombuffer()로 불러오는 것이 좋을 듯.
await websocket.send('정해갑')
# np.frombuffer()로 불러오는 것이 좋을 듯.
# DB에 연결
cursor = attendance_db.cursor(pymysql.cursors.DictCursor)
# 학생을 찾음
sql = "SELECT student_id, embedding FROM student_embedding;"
cursor.execute(sql)
result = cursor.fetchall()
verified_id = '0000000000'
distance_min = 1
for row_data in result:
db_embedding = np.frombuffer(row_data['embedding'], dtype=np.float32)
db_embedding = db_embedding.reshape((1,512))
distance = get_distance(embedding, db_embedding)
if (distance < distance_min):
verified_id = row_data['student_id']
distance_min = distance
# 출석 데이터 전송
data = ''
if distance_min >= 0.6:
# 해당하는 사람 DB에 없음
print('verification failed: not in DB')
data = json.dumps({'state': 'fail'})
else:
# 해당하는 사람 DB에 있음
print('verification success:', verified_id)
data = json.dumps({'state': 'success', 'id': verified_id})
await websocket.send(data)
elif data['action'] == "save_image":
# 출석이 제대로 이뤄지지 않으면 이미지를 저장하여
# 나중에 교강사가 출석을 확인할 수 있도록 한다
......@@ -107,4 +170,4 @@ async def thread(websocket, path):
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()
\ No newline at end of file
asyncio.get_event_loop().run_forever()
......