Files
mixly3/boards/default/micropython_k210_mixgoai/build/lib/ai_face.py

149 lines
3.7 KiB
Python

import KPU as kpu
import gc, image, time
import board
try:
kpu.deinit(task_fe)
kpu.deinit(task_ld)
kpu.deinit(task_fd)
del task_fe
del task_ld
del task_fd
except Exception:
pass
gc.collect()
record_ftr = []
record_ftrs = []
img_face = image.Image(size=(128, 128))
a = img_face.pix_to_ai()
dst_point = [(44, 59), (84, 59), (64, 82), (47, 105), (81, 105)]
start_processing = False
tim2 = time.ticks_ms()
task_fd = None
task_ld = None
task_fe = None
info = None
bb = 1
def set_key_state(*_):
global start_processing
global tim2
if (time.ticks_ms() - tim2) > 4000:
start_processing = True
tim2 = time.ticks_ms()
def init(FD, LD, FE):
global task_fd
global task_ld
global task_fe
# task_fd = kpu.load(0x200000)
# task_ld = kpu.load(0x300000)
# task_fe = kpu.load(0x400000)
task_fd = kpu.load(FD)
task_ld = kpu.load(LD)
task_fe = kpu.load(FE)
gc.collect()
key_gpio = board.pin(9, board.GPIO.IN, board.GPIO.PULL_UP)
key_gpio.irq(set_key_state, board.GPIO.IRQ_RISING, board.GPIO.WAKEUP_NOT_SUPPORT)
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
kpu.init_yolo2(task_fd, 0.5, 0.3, 5, anchor)
def train(img, names, threshold):
global task_fd
global task_ld
global task_fe
global start_processing
global info
global bb
code = kpu.run_yolo2(task_fd, img)
if code:
for i in code:
face_cut = img.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()))
lb = i.rect()
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 start_processing:
record_ftr = feature
record_ftrs.append(record_ftr)
start_processing = False
if max_score > threshold:
info = [names[index], max_score, lb, src_point]
else:
if bb == 1:
print("Please press BOOT key to enter the face")
bb = 0
info = [None, max_score, lb, src_point]
return True
break
else:
info = None
bb = 1
return False
gc.collect()
def info_name():
gc.collect()
return info[0]
def info_score():
return info[1]
def info_face():
return info[2]
def info_organs():
return info[3]