import dlib
import numpy as np
import cv2
from PIL import Image, ImageOps
import warnings

# 忽略 DeprecationWarning
warnings.filterwarnings("ignore", category=DeprecationWarning)

# 模型路径
shape_predictor_path = "shape_predictor_68_face_landmarks.dat"
face_rec_model_path = "dlib_face_recognition_resnet_model_v1.dat"

# 加载模型
detector = dlib.get_frontal_face_detector()
shape_predictor = dlib.shape_predictor(shape_predictor_path)
face_rec_model = dlib.face_recognition_model_v1(face_rec_model_path)

# 图像路径
image_path = "/home/linus/图片/pic1.jpg"
output_path = "/home/linus/图片/pic1_with_landmarks.jpg"

# 使用 PIL 加载图像并修正旋转
image_pil = Image.open(image_path)

# 自动修正图像旋转
image_pil = ImageOps.exif_transpose(image_pil)

# 将图像转换为 numpy 数组，并强制创建新副本
image = np.array(image_pil, copy=True)

# 转换为 BGR 以便使用 OpenCV 进行处理
output_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

# 如果图片太大，可以进行缩小，但这里我们避免缩小图片
# small_image = cv2.resize(image, (image.shape[1] // 2, image.shape[0] // 2))

# 使用一次上采样来提高检测精度
faces = detector(image, 1)

if len(faces) == 0:
    print("未检测到人脸")
else:
    for i, face in enumerate(faces):
        # 获取关键点
        shape = shape_predictor(image, face)
        landmarks = np.array([[p.x, p.y] for p in shape.parts()])

        # 在原图像上绘制68个关键点
        for (x, y) in landmarks:
            cv2.circle(output_image, (x, y), 3, (0, 255, 0), -1)  # 绿色小圆点

        # 计算嵌入特征
        face_descriptor = face_rec_model.compute_face_descriptor(image, shape)
        face_embedding = np.array(face_descriptor)
        
        print(f"人脸 {i + 1} 的嵌入特征向量:")
        print(face_embedding)

    # 保存图像文件
    cv2.imwrite(output_path, output_image)
    print(f"结果图像已保存到 {output_path}")