face_detect.py 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159
  1. import argparse
  2. import numpy as np
  3. import cv2 as cv
  4. def str2bool(v):
  5. if v.lower() in ['on', 'yes', 'true', 'y', 't']:
  6. return True
  7. elif v.lower() in ['off', 'no', 'false', 'n', 'f']:
  8. return False
  9. else:
  10. raise NotImplementedError
  11. parser = argparse.ArgumentParser()
  12. parser.add_argument('--image1', '-i1', type=str, help='Path to the input image1. Omit for detecting on default camera.')
  13. parser.add_argument('--image2', '-i2', type=str, help='Path to the input image2. When image1 and image2 parameters given then the program try to find a face on both images and runs face recognition algorithm.')
  14. parser.add_argument('--video', '-v', type=str, help='Path to the input video.')
  15. parser.add_argument('--scale', '-sc', type=float, default=1.0, help='Scale factor used to resize input video frames.')
  16. parser.add_argument('--face_detection_model', '-fd', type=str, default='face_detection_yunet_2021dec.onnx', help='Path to the face detection model. Download the model at https://github.com/opencv/opencv_zoo/tree/master/models/face_detection_yunet')
  17. parser.add_argument('--face_recognition_model', '-fr', type=str, default='face_recognition_sface_2021dec.onnx', help='Path to the face recognition model. Download the model at https://github.com/opencv/opencv_zoo/tree/master/models/face_recognition_sface')
  18. parser.add_argument('--score_threshold', type=float, default=0.9, help='Filtering out faces of score < score_threshold.')
  19. parser.add_argument('--nms_threshold', type=float, default=0.3, help='Suppress bounding boxes of iou >= nms_threshold.')
  20. parser.add_argument('--top_k', type=int, default=5000, help='Keep top_k bounding boxes before NMS.')
  21. parser.add_argument('--save', '-s', type=str2bool, default=False, help='Set true to save results. This flag is invalid when using camera.')
  22. args = parser.parse_args()
  23. def visualize(input, faces, fps, thickness=2):
  24. if faces[1] is not None:
  25. for idx, face in enumerate(faces[1]):
  26. print('Face {}, top-left coordinates: ({:.0f}, {:.0f}), box width: {:.0f}, box height {:.0f}, score: {:.2f}'.format(idx, face[0], face[1], face[2], face[3], face[-1]))
  27. coords = face[:-1].astype(np.int32)
  28. cv.rectangle(input, (coords[0], coords[1]), (coords[0]+coords[2], coords[1]+coords[3]), (0, 255, 0), thickness)
  29. cv.circle(input, (coords[4], coords[5]), 2, (255, 0, 0), thickness)
  30. cv.circle(input, (coords[6], coords[7]), 2, (0, 0, 255), thickness)
  31. cv.circle(input, (coords[8], coords[9]), 2, (0, 255, 0), thickness)
  32. cv.circle(input, (coords[10], coords[11]), 2, (255, 0, 255), thickness)
  33. cv.circle(input, (coords[12], coords[13]), 2, (0, 255, 255), thickness)
  34. cv.putText(input, 'FPS: {:.2f}'.format(fps), (1, 16), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
  35. if __name__ == '__main__':
  36. ## [initialize_FaceDetectorYN]
  37. detector = cv.FaceDetectorYN.create(
  38. args.face_detection_model,
  39. "",
  40. (320, 320),
  41. args.score_threshold,
  42. args.nms_threshold,
  43. args.top_k
  44. )
  45. ## [initialize_FaceDetectorYN]
  46. tm = cv.TickMeter()
  47. # If input is an image
  48. if args.image1 is not None:
  49. img1 = cv.imread(cv.samples.findFile(args.image1))
  50. img1Width = int(img1.shape[1]*args.scale)
  51. img1Height = int(img1.shape[0]*args.scale)
  52. img1 = cv.resize(img1, (img1Width, img1Height))
  53. tm.start()
  54. ## [inference]
  55. # Set input size before inference
  56. detector.setInputSize((img1Width, img1Height))
  57. faces1 = detector.detect(img1)
  58. ## [inference]
  59. tm.stop()
  60. assert faces1[1] is not None, 'Cannot find a face in {}'.format(args.image1)
  61. # Draw results on the input image
  62. visualize(img1, faces1, tm.getFPS())
  63. # Save results if save is true
  64. if args.save:
  65. print('Results saved to result.jpg\n')
  66. cv.imwrite('result.jpg', img1)
  67. # Visualize results in a new window
  68. cv.imshow("image1", img1)
  69. if args.image2 is not None:
  70. img2 = cv.imread(cv.samples.findFile(args.image2))
  71. tm.reset()
  72. tm.start()
  73. detector.setInputSize((img2.shape[1], img2.shape[0]))
  74. faces2 = detector.detect(img2)
  75. tm.stop()
  76. assert faces2[1] is not None, 'Cannot find a face in {}'.format(args.image2)
  77. visualize(img2, faces2, tm.getFPS())
  78. cv.imshow("image2", img2)
  79. ## [initialize_FaceRecognizerSF]
  80. recognizer = cv.FaceRecognizerSF.create(
  81. args.face_recognition_model,"")
  82. ## [initialize_FaceRecognizerSF]
  83. ## [facerecognizer]
  84. # Align faces
  85. face1_align = recognizer.alignCrop(img1, faces1[1][0])
  86. face2_align = recognizer.alignCrop(img2, faces2[1][0])
  87. # Extract features
  88. face1_feature = recognizer.feature(face1_align)
  89. face2_feature = recognizer.feature(face2_align)
  90. ## [facerecognizer]
  91. cosine_similarity_threshold = 0.363
  92. l2_similarity_threshold = 1.128
  93. ## [match]
  94. cosine_score = recognizer.match(face1_feature, face2_feature, cv.FaceRecognizerSF_FR_COSINE)
  95. l2_score = recognizer.match(face1_feature, face2_feature, cv.FaceRecognizerSF_FR_NORM_L2)
  96. ## [match]
  97. msg = 'different identities'
  98. if cosine_score >= cosine_similarity_threshold:
  99. msg = 'the same identity'
  100. print('They have {}. Cosine Similarity: {}, threshold: {} (higher value means higher similarity, max 1.0).'.format(msg, cosine_score, cosine_similarity_threshold))
  101. msg = 'different identities'
  102. if l2_score <= l2_similarity_threshold:
  103. msg = 'the same identity'
  104. print('They have {}. NormL2 Distance: {}, threshold: {} (lower value means higher similarity, min 0.0).'.format(msg, l2_score, l2_similarity_threshold))
  105. cv.waitKey(0)
  106. else: # Omit input to call default camera
  107. if args.video is not None:
  108. deviceId = args.video
  109. else:
  110. deviceId = 0
  111. cap = cv.VideoCapture(deviceId)
  112. frameWidth = int(cap.get(cv.CAP_PROP_FRAME_WIDTH)*args.scale)
  113. frameHeight = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT)*args.scale)
  114. detector.setInputSize([frameWidth, frameHeight])
  115. while cv.waitKey(1) < 0:
  116. hasFrame, frame = cap.read()
  117. if not hasFrame:
  118. print('No frames grabbed!')
  119. break
  120. frame = cv.resize(frame, (frameWidth, frameHeight))
  121. # Inference
  122. tm.start()
  123. faces = detector.detect(frame) # faces is a tuple
  124. tm.stop()
  125. # Draw results on the input image
  126. visualize(frame, faces, tm.getFPS())
  127. # Visualize results
  128. cv.imshow('Live', frame)
  129. cv.destroyAllWindows()