facenet-pytorch篇

import cv2
from PIL import Image
import os
from facenet_pytorch import MTCNN
from torchvision.transforms import ToPILImage

folder_path = r'E:\Code\Python\QQh5\sourceLib'
output_folder = r'testPersonFace'

# 如果输出文件夹不存在，则创建
os.makedirs(output_folder, exist_ok=True)

# 加载MTCNN模型用于人脸检测
mtcnn = MTCNN(keep_all=True)

# 转换Tensor到PIL Image
to_pil = ToPILImage()


# 人脸文件集合
persons_name = []

# 遍历文件夹中的图像
for file in os.listdir(folder_path):
    if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
        img_path = os.path.join(folder_path, file)

        # 打开图像
        img_obj = Image.open(img_path)

        # 使用MTCNN检测人脸
        boxes, _ = mtcnn.detect(img_obj)

        # 检查是否检测到人脸
        if boxes is not None:
            persons_name.append(img_path)
            print(f'{img_path} 人脸√')


            # # 保存带有检测到人脸的图像（选择第一个人脸）
            # img_with_faces = mtcnn(img_obj)[0]
            #
            # # 转换为PIL Image对象
            # img_with_faces_pil = to_pil(img_with_faces)
            #
            # # 保存图像
            # img_with_faces_pil.save(os.path.join(output_folder, f'output_{file}'))
        else:
            print(f'{img_path} 人脸×')

yolov5 篇

YOLO地址：https://github.com/ultralytics/yolov5/blob/master/README.zh-CN.md

import torch
import cv2
from PIL import Image
import os

pt_path = r'E:\Code\Python\YoLov5\yolov5\yolov5s.pt'

# 模型
model = torch.hub.load(r"E:\Code\Python\QQh5\yolov5", 'custom', path=pt_path,
                       source='local')  # 或 yolov5n - yolov5x6, custom

# 图像
# img_path = r"E:\Code\Python\QQh5\sourceLib\test1.jpg"  # 或文件、路径、PIL、OpenCV、numpy、列表

folder_path = r'E:\Code\Python\QQh5\sourceLib'

# 使用列表推导式获取文件夹中的所有图片路径
image_paths = [os.path.join(folder_path, file) for file in os.listdir(folder_path) if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp'))]

# 打印所有图片路径
for img_path in image_paths:
    # print(img_path)



    img_obj = Image.open(img_path)

    # 推理
    results = model(img_obj)

    # 检查图像中是否存在人物
    for item in results.xyxy[0]:
        # 类别 0 对应 COCO 数据集中的人物类别，根据你的数据集可能需要调整这个
        if int(item[5]) == 0:
            # 检测到人物，保存图像
            print(f'{img_path} 人脸√')
            im = results.crop(save=False)[0]['im']
            a = cv2.imencode(".jpg", im)
            success, enc_img = cv2.imencode(".jpg", im)
            img_data = enc_img.tobytes()

            # 从原始路径获取文件名
            img_filename = img_path.split("\\")[-1]
            with open(f'testPersonFace/output_{img_filename}', 'wb') as f:
                f.write(img_data)
                # print(img_data)

            # # 显示图像
            # cv2.imshow('win', im)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()
        else:
            print(f'{img_path} 人脸×')