环境:
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia -y
pip install onnx==1.15 onnxruntime-gpu==1.17
pip install insightface==0.7.3
pip install opencv-python
pip install gradio
图片清洗,只留下图像中只有一张人脸的,而且人脸是全的。
import os
import shutil
import cv2
import numpy as np
from insightface.app import FaceAnalysis
def is_full_face(facedata, img_shape, threshold=0.5):
"""
判断是否为整张脸
:param facedata: 人脸数据
:param img_shape: 图片尺寸
:param threshold: 阈值
:return: 布尔值,True 表示整张脸,False 表示部分脸
"""
img_width, img_height = img_shape[1], img_shape[0]
# 检查人脸关键点是否在图片内部
kps = facedata['kps']
if np.all(kps >= 10) and np.all(kps[:, 0] <= img_width - 10) and np.all(kps[:, 1] <= img_height - 10):
keypoints_inside = True
else:
keypoints_inside = False
# 满足阈值并且关键点在图片内部
return keypoints_inside
def listPathAllfiles(dirname):
result = []
for maindir, subdir, file_name_list in os.walk(dirname):
for filename in file_name_list:
apath = os.path.join(maindir, filename)
result.append(apath)
return result
# 使用的检测模型名为 buffalo_sc
app = FaceAnalysis(name='buffalo_sc', providers=['CUDAExecutionProvider'])
app.prepare(ctx_id=0, det_size=(640, 640)) # ctx_id 小于0 表示用 CPU 预测,det_size 表示 resize 后的图片分辨率
src = r"C:\Users\Administrator\Pictures\girl_no_train\mangguo_dst"
dst = r"C:\Users\Administrator\Pictures\girl_no_train\mangguo_dst2"
os.makedirs(dst, exist_ok=True)
files = listPathAllfiles(src)
for file in files:
img = cv2.imread(file) # 读取图片
faces = app.get(img) # 得到人脸信息
if len(faces) == 0:
continue
if len(faces) > 1:
continue
for facedata in faces:
if is_full_face(facedata, img.shape):
# print("This is a full face.")
shutil.copy(file, dst)