Take a look at the code structure
faceRecognition_knn
- There are three image collections under knn_example:
2.model
- test: test picture
- train: training set pictures (the picture sets are downloaded from the Internet)
- train1: also the training set picture (the train training set picture is split, the set is relatively small
- trained_knn_model.clf (Save the model after knn classifier training, mainly the coding features of the pictures in the picture collection)
Code directly
# -*- coding: utf-8 -*-
# !/usr/bin/env python
# @Time : 2019/1/10 15:50
# @Author : xhh
# @Desc : 利用knn分类器来进行人脸识别
# @File : face_recognition_knn1.py
# @Software: PyCharm
import math
from sklearn import neighbors
import os
import os.path
import pickle
from PIL import Image, ImageDraw
import face_recognition
from face_recognition.face_recognition_cli import image_files_in_folder
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
def train(train_dir, model_save_path=None, n_neighbors=None, knn_algo='ball_tree', verbose=False):
"""
Structure:
<train_dir>/
├── <person1>/
│ ├── <somename1>.jpeg
│ ├── <somename2>.jpeg
│ ├── ...
├── <person2>/
│ ├── <somename1>.jpeg
│ └── <somename2>.jpeg
└── ...
:param train_dir: 训练集地址
:param model_save_path: 模型所保存的地址
:param n_neighbors: 在训练时如果没有指定权重，则自动选择权重
:param knn_algo:
:param verbose:
:return: knn_clf返回训练后的模型
"""
X = []
y = []
# 循环遍历训练集中的每一个人
for class_dir in os.listdir(train_dir):
if not os.path.isdir(os.path.join(train_dir, class_dir)):
continue
# 循环遍历当前训练集中的每个人
for img_path in image_files_in_folder(os.path.join(train_dir, class_dir)):
image = face_recognition.load_image_file(img_path)
face_bounding_boxes = face_recognition.face_locations(image)
if len(face_bounding_boxes) != 1:
# 如果该训练集中没有人或者有很多人，则跳过该图像
if verbose:
print("Image {} not suitable for training: {}".format(img_path,
"Didn't find a face" if len(
face_bounding_boxes) < 1 else "Found more than one face"))
else:
# 将图片中的人脸的编码加入到训练集中
X.append(face_recognition.face_encodings(image, known_face_locations=face_bounding_boxes)[0])
y.append(class_dir)
# 确定KNN分类器中的权重
if n_neighbors is None:
n_neighbors = int(round(math.sqrt(len(X))))
if verbose:
print("Chose n_neighbors automatically:", n_neighbors)
# 建立并训练KNN训练集
knn_clf = neighbors.KNeighborsClassifier(n_neighbors=n_neighbors, algorithm=knn_algo, weights='distance')
knn_clf.fit(X, y)
# 保存KNN分类器
if model_save_path is not None:
with open(model_save_path, 'wb') as f:
pickle.dump(knn_clf, f)
return knn_clf
def predict(X_img_path, knn_clf=None, model_path=None, distance_threshold=0.6):
"""
:param X_img_path: 测试集的图片地址
:param knn_clf: 训练好的模型
:param model_path: 模型地址
:param distance_threshold: 给出当前测试图片中的人脸向量与模型中的距离
:return:
"""
if not os.path.isfile(X_img_path) or os.path.splitext(X_img_path)[1][1:] not in ALLOWED_EXTENSIONS:
raise Exception("Invalid image path: {}".format(X_img_path))
if knn_clf is None and model_path is None:
raise Exception("Must supply knn classifier either thourgh knn_clf or model_path")
# 加载KNN模型
if knn_clf is None:
with open(model_path, 'rb') as f:
knn_clf = pickle.load(f)
# 加载图片文件夹以及人脸
X_img = face_recognition.load_image_file(X_img_path)
X_face_locations = face_recognition.face_locations(X_img)
# 如过图片中没有人脸，则返回空的结果集
if len(X_face_locations) == 0:
return []
# 找出测试集中的人脸编码
faces_encodings = face_recognition.face_encodings(X_img, known_face_locations=X_face_locations)
# 利用KNN模型找出测试集中最匹配的人脸
closest_distances = knn_clf.kneighbors(faces_encodings, n_neighbors=1)
are_matches = []
for i in range(len(X_face_locations)):
are_matches.append(closest_distances[0][i][0] <= distance_threshold)
# 预测类并删除不在阈值范围内的分类
return [(pred, loc) if rec else ("unknown", loc) for pred, loc, rec in
zip(knn_clf.predict(faces_encodings), X_face_locations, are_matches)]
def show_prediction_labels_on_image(img_path, predictions):
"""
在图片给出标签并展示人脸
:param img_path: path to image to be recognized
:param predictions: results of the predict function
:return:
"""
pil_image = Image.open(img_path).convert("RGB")
draw = ImageDraw.Draw(pil_image)
for name, (top, right, bottom, left) in predictions:
# 将人脸框出来利用pillow进行标注
draw.rectangle(((left, top), (right, bottom)), outline=(0, 0, 255))
name = name.encode("UTF-8")
# 在人脸下面进行标配注
text_width, text_height = draw.textsize(name)
draw.rectangle(((left, bottom - text_height - 10), (right, bottom)), fill=(0, 0, 255), outline=(0, 0, 255))
draw.text((left + 6, bottom - text_height - 5), name, fill=(255, 255, 255, 255))
# 根据文档删除标注
del draw
# 显示图片结果
pil_image.show()
if __name__ == "__main__":
# 1:训练KNN分类器，并保存
# Once the model is trained and saved, you can skip this step next time.
# print("Training KNN classifier...")
# classifier = train("knn_examples/train1", model_save_path="model/trained_knn_model.clf", n_neighbors=2)
# print("Training complete!")
# STEP 2: 利用分类器，来预测该图片是否为已知的
for image_file in os.listdir("knn_examples/test"):
full_file_path = os.path.join("knn_examples/test", image_file)
print("Looking for faces in {}".format(image_file))
# Find all people in the image using a trained classifier model
# Note: You can pass in either a classifier file name or a classifier model instance
predictions = predict(full_file_path, model_path="model/trained_knn_model.clf")
# 在控制台打印结果
for name, (top, right, bottom, left) in predictions:
print("- Found {} at ({}, {})".format(name, left, top))
# 在结果集中进行显示
show_prediction_labels_on_image(os.path.join("knn_examples/test", image_file), predictions)
operation result:
Training KNN classifier...
Training complete!
Looking for faces in Aaron_Eckhart_0001.jpg
- Found Aaron_Eckhart at (67, 80)
Looking for faces in Aaron_Peirsol_0003.jpg
- Found Aaron_Peirsol at (76, 86)
Looking for faces in Abba_Eban_0001.jpg
- Found Abba_Eban at (67, 80)
Look at the screenshot:
Annotation after recognition of some pictures:
Not in the training set, here is the unknow, like the little girl in the picture below. . .
github address: https://github.com/XHHz/faceRecognition_knn
You can pay attention to the official account of my friend and me~~~ Here are some python technical information that my friend and I update from time to time! ! You can also leave a message to discuss technical issues. I hope you can support and pay attention to it. Thank you~~
