face_detection.py 2.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970
  1. import cv2
  2. import dlib
  3. import numpy as np
  4. ## Face detection
  5. def face_detection(img,upsample_times=1):
  6. # Ask the detector to find the bounding boxes of each face. The 1 in the
  7. # second argument indicates that we should upsample the image 1 time. This
  8. # will make everything bigger and allow us to detect more faces.
  9. detector = dlib.get_frontal_face_detector()
  10. faces = detector(img, upsample_times)
  11. return faces
  12. PREDICTOR_PATH = 'FaceSwap/models/shape_predictor_68_face_landmarks.dat'
  13. predictor = dlib.shape_predictor(PREDICTOR_PATH)
  14. ## Face and points detection
  15. def face_points_detection(img, bbox:dlib.rectangle):
  16. # Get the landmarks/parts for the face in box d.
  17. shape = predictor(img, bbox)
  18. # loop over the 68 facial landmarks and convert them
  19. # to a 2-tuple of (x, y)-coordinates
  20. coords = np.asarray(list([p.x, p.y] for p in shape.parts()), dtype=np.int)
  21. # return the array of (x, y)-coordinates
  22. return coords
  23. def select_face(im, r=10, choose=True):
  24. faces = face_detection(im)
  25. if len(faces) == 0:
  26. return None, None, None
  27. if len(faces) == 1 or not choose:
  28. idx = np.argmax([(face.right() - face.left()) * (face.bottom() - face.top()) for face in faces])
  29. bbox = faces[idx]
  30. else:
  31. bbox = []
  32. def click_on_face(event, x, y, flags, params):
  33. if event != cv2.EVENT_LBUTTONDOWN:
  34. return
  35. for face in faces:
  36. if face.left() < x < face.right() and face.top() < y < face.bottom():
  37. bbox.append(face)
  38. break
  39. im_copy = im.copy()
  40. for face in faces:
  41. # draw the face bounding box
  42. cv2.rectangle(im_copy, (face.left(), face.top()), (face.right(), face.bottom()), (0, 0, 255), 1)
  43. cv2.imshow('Click the Face:', im_copy)
  44. cv2.setMouseCallback('Click the Face:', click_on_face)
  45. while len(bbox) == 0:
  46. cv2.waitKey(1)
  47. cv2.destroyAllWindows()
  48. bbox = bbox[0]
  49. points = np.asarray(face_points_detection(im, bbox))
  50. im_w, im_h = im.shape[:2]
  51. left, top = np.min(points, 0)
  52. right, bottom = np.max(points, 0)
  53. x, y = max(0, left - r), max(0, top - r)
  54. w, h = min(right + r, im_h) - x, min(bottom + r, im_w) - y
  55. return points - np.asarray([[x, y]]), (x, y, w, h), im[y:y + h, x:x + w]