train!

实验八/3.train.py

@@ -1,31 +1,74 @@
 import dlib
 import cv2 as cv
+import joblib
 import numpy as np
 import os
 
+from sklearn.ensemble import RandomForestClassifier
+from sklearn.model_selection import train_test_split
+from sklearn.multioutput import MultiOutputClassifier
+from tqdm import tqdm
+
 # 提取人脸特征
 def extract_face_features(image_path):
-    img=cv.imread(image_path)
-    detections=detector(img,1)
-    face_features=[]
+    img = cv.imread(image_path)
+    detections = detector(img, 1)
+    face_features = []
+    face_bboxes = []
     for rect in detections:
-        shape=predictor(img,rect)
-        face_descriptor=face_rec_model.compute_face_descriptor(img,shape)
+        shape = predictor(img, rect)
+        face_descriptor = face_rec_model.compute_face_descriptor(img, shape)
         face_features.append(face_descriptor)
-    return face_features
+        bbox = (rect.left(), rect.top(), rect.right(), rect.bottom())
+        face_bboxes.append(bbox)
+    return face_features, face_bboxes
 
 # 获取图片路径
-def get_img_path(directory,extension=None):
-    if extension==None:
-        extension=['.jpg','.jpeg','.png']
-    files=[]
+def get_img_path(directory, extension=None):
+    if extension is None:
+        extension = ['.jpg', '.jpeg', '.png']
+    files = []
     for root, dirs, file_names in os.walk(directory):
         for file_name in file_names:
             if any(file_name.lower().endswith(ext) for ext in extension):
-                files.append(os.path.join(root,file_name))
+                files.append(os.path.join(root, file_name))
     return files
 
+
 # 加载dlib模型
 detector = dlib.get_frontal_face_detector()
-predictor=dlib.shape_predictor('models/shape_predictor_68_face_landmarks.dat')
-face_rec_model=dlib.face_recognition_model_v1('实验八/models/dlib_face_recognition_resnet_model_v1.dat')
+predictor = dlib.shape_predictor('models/shape_predictor_68_face_landmarks.dat')
+face_rec_model = dlib.face_recognition_model_v1('models/dlib_face_recognition_resnet_model_v1.dat')
+
+# 获取图片
+img_directory = 'cache/pretrained/'
+images = get_img_path(img_directory)
+
+# 提取特征
+features = []
+labels = []
+
+for image in tqdm(images, desc='提取图片特征中：'):
+    extracted_features, face_bboxes = extract_face_features(image)
+    for feature,bbox in zip(extracted_features, face_bboxes):
+        features.append(feature)
+        labels.append(bbox)
+
+X_train = np.array(features)
+Y_train = np.array(labels)
+
+# 分割测试集
+X_train, X_test, Y_train, Y_test = train_test_split(X_train, Y_train, test_size=0.2, random_state=42)
+
+# 训练SVM模型
+print('训练模型中')
+clf=MultiOutputClassifier(RandomForestClassifier(n_estimators=100,random_state=42))
+clf.fit(X_train, Y_train)
+
+# 评估训练数据
+predictions = clf.predict(X_test)
+accuracy=(predictions==Y_test).mean()
+print(f'分类器准确度：{accuracy * 100:.2f}%')
+
+os.makedirs('models/', exist_ok=True)
+joblib.dump(clf, 'models/my_classifier.pkl')