75 lines
2.3 KiB
Python
75 lines
2.3 KiB
Python
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 = []
|
|
face_bboxes = []
|
|
for rect in detections:
|
|
shape = predictor(img, rect)
|
|
face_descriptor = face_rec_model.compute_face_descriptor(img, shape)
|
|
face_features.append(face_descriptor)
|
|
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 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))
|
|
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')
|
|
|
|
# 获取图片
|
|
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')
|