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()