16_faceComparison.py 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199
  1. '''
  2. Face recognize demo, download model from maixhub first:
  3. https://maixhub.com/
  4. '''
  5. from maix import nn, display, camera, image
  6. from maix.nn.app import face
  7. from maix.nn.app.face import FaceRecognize
  8. import time
  9. import sys
  10. sys.path.append('/root/')
  11. from CocoPi import BUTTON
  12. import os
  13. ScreenOrientation = False
  14. try:
  15. if os.path.exists("/etc/cameraSize.cfg"):
  16. cameraSize = True
  17. else:
  18. cameraSize = False
  19. except:
  20. cameraSize = False
  21. def getLcdRotation(cameraCapture):
  22. global cameraSize
  23. if cameraSize:
  24. return lcdRotationNew(cameraCapture)
  25. else:
  26. return lcdRotation(cameraCapture)
  27. def lcdRotationNew(inputImg):
  28. global cameraSize,ScreenOrientation
  29. imageRotationBuffer = inputImg.crop(0, 0, 320, 240)
  30. if ScreenOrientation:
  31. imgRotationAim = image.new(size = (240, 320))
  32. rotationAngle = 90
  33. GETROTATION = imageRotationBuffer.rotate(+rotationAngle, adjust=1)
  34. else:
  35. imgRotationAim = image.new(size = (320, 240))
  36. GETROTATION = imageRotationBuffer
  37. GETROTATION = imgRotationAim.draw_image(GETROTATION,0,0,alpha=1)
  38. return GETROTATION
  39. def lcdRotation(inputImg):
  40. global cameraSize,ScreenOrientation
  41. imageRotationBuffer = inputImg.crop(0, 0, 240, 320)
  42. if ScreenOrientation:
  43. imgRotationAim = image.new(size = (240, 320))
  44. rotationAngle = 180
  45. else:
  46. imgRotationAim = image.new(size = (320, 240))
  47. rotationAngle = 90
  48. GETROTATION = imageRotationBuffer.rotate(+rotationAngle, adjust=1)
  49. GETROTATION = imgRotationAim.draw_image(GETROTATION,0,0,alpha=1)
  50. return GETROTATION
  51. key_A = BUTTON(14)
  52. key_B = BUTTON(8)
  53. key_C = BUTTON(13)
  54. key_D = BUTTON(7)
  55. class Face_Recognizer:
  56. def __init__(self, threshold = 0.5, nms = 0.3, max_face_num = 1):
  57. model = "/root/preset/model/retinaface.mud"
  58. model_fe = "/root/preset/model/fe_resnet.mud"
  59. self.input_size = (224, 224, 3)
  60. input_size_fe = (128, 128, 3)
  61. self.feature_len = 256
  62. self.features = []
  63. print("-- load model:", model)
  64. m = nn.load(model)
  65. print("-- load ok")
  66. print("-- load model:", model_fe)
  67. m_fe = nn.load(model_fe)
  68. print("-- load ok")
  69. self.recognizer = FaceRecognize(m, m_fe, self.feature_len, self.input_size, threshold, nms, max_face_num)
  70. print("-- init end")
  71. def get_faces(self, img, std_img = False):
  72. faces = self.recognizer.get_faces(img, std_img)
  73. return faces
  74. def __len__(self):
  75. return len(self.features)
  76. def add_user(self, name, feature):
  77. self.features.append([name, feature])
  78. return True
  79. def remove_user(self, name_del):
  80. rm = None
  81. for name, feature in self.features:
  82. if name_del == name:
  83. rm = [name, feature]
  84. if rm:
  85. self.features.remove(rm)
  86. return True
  87. return False
  88. def recognize(self, feature):
  89. max_score = 0
  90. uid = -1
  91. for i, user in enumerate(self.features):
  92. score = self.recognizer.compare(user[1], feature)
  93. if score > max_score:
  94. max_score = score
  95. uid = i
  96. if uid >= 0:
  97. return self.features[uid][0], max_score
  98. return None, 0
  99. def get_input_size(self):
  100. '''
  101. @return input_size (w, h, c)
  102. '''
  103. return self.input_size
  104. def get_feature_len(self):
  105. return self.feature_len
  106. def darw_info(self, img, box, points, disp_str, bg_color=(255, 0, 0, 255), font_color=(255, 255, 255, 255), font_size=32):
  107. font_wh = image.get_string_size(disp_str)
  108. for p in points:
  109. img.draw_rectangle(p[0] - 1, p[1] -1, p[0] + 1, p[1] + 1, color=bg_color)
  110. img.draw_rectangle(box[0], box[1], box[0] + box[2], box[1] + box[3], color=bg_color, thickness=2)
  111. if disp_str:
  112. img.draw_rectangle(box[0], box[1] - font_wh[1], box[0] + font_wh[0], box[1], color=bg_color, thickness = -1)
  113. img.draw_string(box[0], box[1] - font_wh[1], disp_str, color=font_color)
  114. def darw_title(self, img, dis_size ,key_l = None, key_r =None):
  115. if key_C:
  116. key_l = "| "+ key_l
  117. img.draw_string( 1, 2 ,key_l , scale = 1, color = (255, 255, 255), thickness = 2)
  118. if key_D:
  119. key_r = key_r+" |"
  120. w = int(dis_size[0] - 4 - image.get_string_size(key_r)[0] * 1)
  121. img.draw_string( w, 2 ,key_r , scale = 1, color = (255, 255, 255), thickness = 2)
  122. max_face_num = 4
  123. detect_threshold = 0.5
  124. detect_nms = 0.3
  125. score_threshold = 70
  126. names = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z"]
  127. face_recognizer = Face_Recognizer(detect_threshold, detect_nms, max_face_num = max_face_num)
  128. camera.config(size=face_recognizer.get_input_size()[:2])
  129. canvasImg = image.new(size = (240, 320))
  130. while 1:
  131. canvasImg.clear()
  132. img = getLcdRotation(camera.capture()).crop(0, 0,224, 224)
  133. if not img:
  134. time.sleep(0.02)
  135. continue
  136. faces = face_recognizer.get_faces(img)
  137. face_recognizer.darw_title(img , face_recognizer.get_input_size()[:2] , "rm" ,"add")
  138. if faces:
  139. # for prob, box, landmarks, feature, std_img in faces:
  140. for prob, box, landmarks, feature in faces:
  141. # [ prob, [x,y,w,h], [[x,y], [x,y], [x,y], [x,y], [x,y]], feature ]
  142. if key_C.is_pressed():
  143. if len(face_recognizer) < len(names):
  144. while not (key_C.is_pressed() == False):
  145. time.sleep(0.1)
  146. idx = len(face_recognizer)
  147. print("add user: {}, {}".format(idx, names[idx]))
  148. face_recognizer.add_user(names[idx], feature)
  149. else:
  150. print("user full")
  151. name, score = face_recognizer.recognize(feature)
  152. if name:
  153. if score > score_threshold:
  154. face_recognizer.darw_info(img, box, landmarks, "{}:{:.2f}".format(name, score), font_color=(0, 0, 255, 255), bg_color=(0, 255, 0, 255))
  155. print("user: {}, score: {:.2f}".format(name, score))
  156. else:
  157. face_recognizer.darw_info(img, box, landmarks, "{}:{:.2f}".format(name, score), font_color=(255, 255, 255, 255), bg_color=(255, 0, 0, 255))
  158. print("maybe user: {}, score: {:.2f}".format(name, score))
  159. else:
  160. face_recognizer.darw_info(img, box, landmarks, "", font_color=(255, 255, 255, 255), bg_color=(255, 255, 255, 255))
  161. if key_D.is_pressed():
  162. if len(face_recognizer) > 0:
  163. while not (key_D.is_pressed() == False):
  164. time.sleep(0.1)
  165. idx = len(face_recognizer) - 1
  166. print("remove user:", names[idx])
  167. face_recognizer.remove_user(names[idx])
  168. else:
  169. print("user empty")
  170. canvasImg.draw_image(img,48,8)
  171. canvasImg.draw_image((image.open("/root/preset/img/exit_ff0000_24x24.png")).rotate(0, adjust=0),288,216,alpha=1)
  172. display.show(canvasImg)