for name in dir():
    if not name.startswith('_'):
        del globals()[name]
import gc
import KPU as kpu
from Maix import utils
utils.gc_heap_size(0xAF000)
import sensor
import image
import lcd
import time
from Maix import FPIOA, GPIO
import gc
from fpioa_manager import fm
from board import board_info
try:from cocorobo import display_cjk_string
except:pass
def lcd_draw_string(canvas, x, y, text, color=(255,255,255), font_size=1, scale=1, mono_space=False, auto_wrap=True):
    try:
        display_cjk_string(canvas, x, y, text, font_size=font_size, color=color)
        return canvas
    except: return canvas.draw_string(x, y, text, color=color, scale=scale, mono_space=mono_space)

import lcd
import image

try:from cocorobo import display_cjk_string
except:pass

def lcd_draw_string(canvas, x, y, text, color=(255,255,255), font_size=1, scale=1, mono_space=False, auto_wrap=True):
    try:
        display_cjk_string(canvas, x, y, text, font_size=font_size, color=color)
        return canvas
    except: return canvas.draw_string(x, y, text, color=color, scale=scale, mono_space=mono_space)

_canvas_x, _canvas_y = 0, 0

def set_key_state(*_):
    global start_processing
    start_processing = True

gc.enable()
gc.collect()

# kpu.memtest()
task_fd = kpu.load("/sd/preset/models/face_reocgnition/FD.emodel")
task_ld = kpu.load("/sd/preset/models/face_reocgnition/KP_chwise.emodel")
task_fe = kpu.load("/sd/preset/models/face_reocgnition/face_recognition_v4_2020_09_07.emodel")
#kpu.set_outputs(task, 0, 1,1,1) # 配置输出形状
kpu.set_outputs(task_fe, 0, 1,1,512)

# kpu.memtest()
kpu.set_outputs(task_fe, 0,1,1, 512)
fm.register(10, fm.fpioa.GPIOHS0)
key_gpio = GPIO(GPIO.GPIOHS0, GPIO.IN)
start_processing = False

key_gpio.irq(set_key_state, GPIO.IRQ_RISING, GPIO.WAKEUP_NOT_SUPPORT)

lcd.init(type=2)
lcd.rotation(1)
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QVGA)
sensor.set_vflip(0)
sensor.set_hmirror(0)
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))
index = -1
a = img_face.pix_to_ai()
record_ftr = []
record_ftrs = []
gc.collect()

names = ["李晓明", "王晓丽"]




lcd.init(type=2,freq=15000000,width=240,height=240,color=(0,0,0))
lcd.rotation(1)
lcd.clear(lcd.BLACK)
while True:
    canvas = sensor.snapshot()
    code = kpu.run_yolo2(task_fd, canvas)
    index = -1

    total_faces = 0

    if code:
        for i in code:
            # print(code)
            canvas.draw_rectangle(i.rect())
            # print(code[0])

            face_cut = canvas.cut(i.x(), i.y(), i.w(), i.h())
            face_cut_128 = face_cut.resize(128, 128)
            a = face_cut_128.pix_to_ai()
            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()))
            a = canvas.draw_circle(le[0], le[1], 4)
            a = canvas.draw_circle(re[0], re[1], 4)
            a = canvas.draw_circle(nose[0], nose[1], 4)
            a = canvas.draw_circle(lm[0], lm[1], 4)
            a = canvas.draw_circle(rm[0], rm[1], 4)
            src_point = [le, re, nose, lm, rm]
            T = image.get_affine_transform(src_point, dst_point)
            a = image.warp_affine_ai(canvas, 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 = -1

            for k in range(len(scores)):
                if max_score < scores[k]:
                    max_score = scores[k]
                    index = k

            print(len(names), index)
            
            if(len(names) > index):
                if max_score > 85:
                    a = lcd_draw_string(canvas, i.x(), i.y(), ("%s :%2.1f" % (names[index], max_score)), color=(0, 255, 0), scale=2, mono_space=False)
                else:
                    a = lcd_draw_string(canvas, i.x(), i.y(), ("X :%2.1f" % (max_score)), color=(255, 0, 0), scale=2, mono_space=False)
                if start_processing:
                    record_ftr = feature
                    record_ftrs.append(record_ftr)
                    start_processing = False
            break

    canvas = canvas.cut(40,0,240,240)
    lcd.display(canvas, oft=(0,0))
    gc.collect()
    # kpu.memtest()
    if "李晓明" == (names[index] if (index != -1 and len(names) > index) else ""):
        lcd_draw_string(canvas, 0,0, "李晓明", color=(255,0,0), scale=5, mono_space=False)
    elif "王晓丽" == (names[index] if (index != -1 and len(names) > index) else ""):
        lcd_draw_string(canvas, 0,0, "王晓丽", color=(255,0,0), scale=5, mono_space=False)
    lcd.display(canvas, oft=(_canvas_x,_canvas_y))