import sensor
import image
import lcd
import KPU as kpu
import time
from Maix import FPIOA, GPIO
import gc
from fpioa_manager import fm
from board import board_info
import utime

from Maix import utils
import machine

#记录的人脸数
face_num = 10

task_fd = kpu.load(0x300000)
task_ld = kpu.load(0x400000)
task_fe = kpu.load(0x500000)

clock = time.clock()
#设置按键的IO
fm.register(2, fm.fpioa.GPIOHS0)
key_gpio = GPIO(GPIO.GPIOHS0, GPIO.IN)
start_processing = False

BOUNCE_PROTECTION = 50

#按键中断
def set_key_state(*_):
    global start_processing
    start_processing = True
    utime.sleep_ms(BOUNCE_PROTECTION)

#设置按键中断
key_gpio.irq(set_key_state, GPIO.IRQ_RISING, GPIO.WAKEUP_NOT_SUPPORT)

lcd.init()
lcd.init(invert = 10)
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QVGA)
sensor.set_hmirror(1)
#sensor.set_vflip(1)
sensor.run(1)
#不可修改
anchor = (1.889, 2.5245, 2.9465, 3.94056, 3.99987, 5.3658, 5.155437,
          6.92275, 6.718375, 9.01025)  # anchor for face detect‘
#人脸关键点
dst_point = [(44, 59), (84, 59), (64, 82), (47, 105),
             (81, 105)]  # standard face key point position
a = kpu.init_yolo2(task_fd, 0.5, 0.3, 5, anchor)
img_lcd = image.Image()
img_face = image.Image(size=(128, 128))
a = img_face.pix_to_ai()

record_ftr = []
record_ftrs = []

#人脸名称，可自定义
names = ['Mr.1', 'Mr.2', 'Mr.3', 'Mr.4', 'Mr.5',
         'Mr.6', 'Mr.7', 'Mr.8', 'Mr.9', 'Mr.10']

#人脸检测的阈值
ACCURACY = 80



while (1):
    img = sensor.snapshot()
    clock.tick()
    #运行人脸检测模型
    code = kpu.run_yolo2(task_fd, img)
    #保存的人脸大于5时退出，这里可自由发挥
    if 	len(record_ftrs) >=face_num:
        break

    #如果存在人脸
    if code:
        for i in code:
            a = img.draw_rectangle(i.rect())
            #裁切人脸的方框并resize至128*128
            face_cut = img.cut(i.x(), i.y(), i.w(), i.h())
            face_cut_128 = face_cut.resize(128, 128)
            #将修改的图片应用至ai缓存区
            a = face_cut_128.pix_to_ai()
            #Landmark for face 5 points
            fmap = kpu.forward(task_ld, face_cut_128)
            plist = fmap[:]
            le = (i.x() + int(plist[0] * i.w() - 10), i.y() + int(plist[1] * i.h()))
            re = (i.x() + int(plist[2] * i.w()), i.y() + int(plist[3] * i.h()))
            nose = (i.x() + int(plist[4] * i.w()), i.y() + int(plist[5] * i.h()))
            lm = (i.x() + int(plist[6] * i.w()), i.y() + int(plist[7] * i.h()))
            rm = (i.x() + int(plist[8] * i.w()), i.y() + int(plist[9] * i.h()))
            # 人脸管件点信息
            src_point = [le, re, nose, lm, rm]
            T = image.get_affine_transform(src_point, dst_point)
            a = image.warp_affine_ai(img, img_face, T)
            a = img_face.ai_to_pix()
            del (face_cut_128)
            #获取人脸特征图
            fmap = kpu.forward(task_fe, img_face)
            #人脸特征
            feature = kpu.face_encode(fmap[:])
            reg_flag = False
            scores = []
            #计算已储存的人脸的特征值
            for j in range(len(record_ftrs)):
                score = kpu.face_compare(record_ftrs[j], feature)
                scores.append(score)
            #找出当前储存的人脸特征与现在检测到的人脸特征最相似的
            max_score = 0
            index = 0
            for k in range(len(scores)):
                if max_score < scores[k]:
                    max_score = scores[k]
                    index = k
            #在显示屏上显示当前人脸的对象
            if max_score > ACCURACY:
                a = img.draw_string(i.x(), i.y(), ("%s :%2.1f" % (
                    names[index], max_score)), lcd.BLUE,2)
            else:
                a = img.draw_string(i.x(), i.y(), ("X :%2.1f" % (
                    max_score)), lcd.RED,2)
            #按键按下，录入人脸
            if start_processing:
                record_ftr = feature
                record_ftrs.append(record_ftr)
                print(record_ftrs)
                start_processing = False
    fps = clock.fps()
    a = lcd.display(img)
    gc.collect()


#保存人脸特征值
with open('/flash/face_tag.dat','wb') as f:
    for i in record_ftrs:
        f.write(i)