face_align.py 2.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192
  1. import face_alignment
  2. from skimage import io
  3. import joblib
  4. import os
  5. from sklearn import svm
  6. import concurrent.futures
  7. import numpy as np
  8. import time
  9. fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=False)
  10. # Training the SVC classifier
  11. # The training data would be all the face encodings from all the known images and the labels are their names
  12. encodings = []
  13. names = []
  14. tuples = []
  15. train_list = []
  16. lock = None
  17. is_train = False
  18. if not is_train:
  19. try:
  20. clf = joblib.load('saved_model_fa.pkl')
  21. except:
  22. clf = None
  23. if clf is None:
  24. is_train = True
  25. def train_image(image, person):
  26. img = io.imread(image)
  27. preds = fa.get_landmarks(img)
  28. # If training image contains exactly one face
  29. if len(preds) == 1:
  30. tuples.append((preds[0].reshape(136), person))
  31. else:
  32. print(image + " was skipped and can't be used for training")
  33. if is_train:
  34. # Training directory
  35. train_str = os.path.join('data', 'peeps', 'anchor')
  36. train_dir = os.listdir(os.path.join('data', 'peeps', 'anchor'))
  37. # Loop through each person in the training directory
  38. with concurrent.futures.ThreadPoolExecutor() as executor:
  39. for person in train_dir:
  40. pix_str = os.path.join(train_str, person)
  41. pix = os.listdir(os.path.join(train_str, person))
  42. # Loop through each training image for the current person
  43. for person_img in pix:
  44. # Get the face encodings for the face in each image file
  45. image = os.path.join(pix_str, person_img)
  46. executor.submit(train_image, image, person)
  47. # Create and train the SVC classifier
  48. encodings = [x for x, _ in tuples]
  49. names = [y for _, y in tuples]
  50. clf = svm.SVC(gamma='scale', probability=True)
  51. clf.fit(encodings, names)
  52. joblib.dump(clf, 'saved_model_fa.pkl')
  53. # Load the test image with unknown faces into a numpy array
  54. test_image = io.imread(os.path.join('application_data', 'input_image', 'input_image4.jpg'))
  55. # Find all the faces in the test image using the default HOG-based model
  56. face_locations = fa.get_landmarks(test_image)
  57. no = len(face_locations)
  58. print("Number of faces detected: ", no)
  59. # Predict all the faces in the test image using the trained classifier
  60. print("Found:")
  61. for i in range(no):
  62. test_image_enc = face_locations[i].reshape(136)
  63. start_time = time.perf_counter_ns()
  64. proba = clf.predict_proba([test_image_enc])
  65. end_time = time.perf_counter_ns()
  66. process_time = end_time - start_time
  67. classes = clf.classes_
  68. print(classes)
  69. i = np.argmax(proba)
  70. proba = list(*proba)
  71. name = classes[i]
  72. print(name, "{:.2f}".format(proba[i]), proba, process_time)
  73. # input = io.imread(os.path.join('application_data','input_image','input_image.jpg'))
  74. # preds = fa.get_landmarks(input)
  75. # print(preds)