#实战项目：智能人脸考勤系统

# 简化版 P-Net 仅展示核心分支（实际使用直接用封装好的库）
import torch
import torch.nn as nn

class PNet(nn.Module):
    """Proposal Network：生成候选窗口，输出分类/框回归/关键点"""
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(3, 10, 3), nn.PReLU(),
            nn.MaxPool2d(2, 2, ceil_mode=True),
            nn.Conv2d(10, 16, 3), nn.PReLU(),
            nn.Conv2d(16, 32, 3), nn.PReLU()
        )
        # 三个并行分支
        self.cls = nn.Conv2d(32, 2, 1)       # 人脸/非人脸分类
        self.box = nn.Conv2d(32, 4, 1)       # 边界框偏移量
        self.landmark = nn.Conv2d(32, 10, 1) # 5个关键点坐标
    
    def forward(self, x):
        x = self.layers(x)
        return self.cls(x), self.box(x), self.landmark(x)

import cv2
from facenet_pytorch import MTCNN

# 初始化轻量级MTCNN
mtcnn = MTCNN(
    image_size=160,                # 输出的对齐后人脸尺寸
    min_face_size=20,              # 最小人脸检测尺寸（像素）
    thresholds=[0.6, 0.7, 0.7],   # 三阶段的置信度阈值
    device="cuda" if torch.cuda.is_available() else "cpu"
)

def detect_and_crop_face(img_path):
    img = cv2.imread(img_path)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    # 检测人脸框、置信度、5个关键点
    boxes, probs, landmarks = mtcnn.detect(img_rgb, landmarks=True)
    
    # 只保留置信度>0.8的人脸
    if boxes is not None:
        for i, box in enumerate(boxes):
            if probs[i] > 0.8:
                x1, y1, x2, y2 = box.astype(int)
                # 画框
                cv2.rectangle(img, (x1, y1), (x2, y2), (0,255,0), 2)
                # 画关键点
                for point in landmarks[i]:
                    cv2.circle(img, tuple(point.astype(int)), 2, (0,0,255), -1)
    return img

# 测试
result = detect_and_crop_face("test_face.jpg")
cv2.imshow("MTCNN检测", result)
cv2.waitKey(0)
cv2.destroyAllWindows()

# 简化版ArcFace损失核心逻辑
import torch
import torch.nn as nn
import math
import torch.nn.functional as F

class ArcMarginProduct(nn.Module):
    def __init__(self, in_features=512, out_features=1000, s=30.0, m=0.50):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.s = s  # 特征缩放因子，让分类器输出更“尖锐”
        self.m = m  # 角度边缘，拉开类别间距离
        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
        nn.init.xavier_uniform_(self.weight)
        
        # 预计算cos(m)和sin(m)
        self.cos_m = math.cos(m)
        self.sin_m = math.sin(m)
        self.th = math.cos(math.pi - m)  # 当θ+m超过π时的边界
        self.mm = math.sin(math.pi - m) * m  # 替代θ+m的线性惩罚

    def forward(self, features, labels):
        # L2归一化特征和权重
        features = F.normalize(features)
        weight = F.normalize(self.weight)
        
        # 计算cosθ = features · weight^T
        cos_theta = F.linear(features, weight)
        # 计算sinθ = sqrt(1 - cos²θ)
        sin_theta = torch.sqrt((1.0 - cos_theta.pow(2)).clamp(0, 1))
        # 计算cos(θ + m) = cosθ*cosm - sinθ*sinm
        cos_theta_m = cos_theta * self.cos_m - sin_theta * self.sin_m
        
        # 当θ+m超出范围时，用cosθ - mm替代，保持梯度稳定
        cos_theta_m = torch.where(cos_theta > self.th, cos_theta_m, cos_theta - self.mm)
        
        # 构建one-hot标签，只对目标类别应用加了边缘的余弦值
        one_hot = torch.zeros_like(cos_theta)
        one_hot.scatter_(1, labels.view(-1, 1).long(), 1)
        
        # 目标类别用cos(θ+m)，其他类别保持原cosθ
        output = (one_hot * cos_theta_m) + ((1.0 - one_hot) * cos_theta)
        # 缩放输出，提高区分度
        output *= self.s
        return output

from facenet_pytorch import InceptionResnetV1
import torchvision.transforms as transforms
from PIL import Image
import torch

# 加载预训练模型（在 VGGFace2 上训练）
model = InceptionResnetV1(pretrained='vggface2').eval().to(
    "cuda" if torch.cuda.is_available() else "cpu"
)

# 预处理必须和训练时一致：Resize → ToTensor → Normalize
transform = transforms.Compose([
    transforms.Resize((160, 160)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

def extract_feature(img_path):
    img = Image.open(img_path).convert('RGB')
    img_tensor = transform(img).unsqueeze(0).to(next(model.parameters()).device)
    with torch.no_grad():
        feature = model(img_tensor).squeeze().cpu().numpy()
    return feature  # 512 维归一化特征向量

import numpy as np

def cosine_sim(feat1, feat2):
    """计算两个归一化特征的余弦相似度"""
    return np.dot(feat1, feat2)  # 已经归一化，直接点积即可

def find_best_match(input_feat, db_feats, db_ids, db_names, threshold=0.6):
    """
    在数据库中找最佳匹配
    - threshold: 相似度阈值（IR-SE50 通常在 0.5-0.7 之间效果最佳）
    """
    best_idx = -1
    best_sim = 0
    for i, db_feat in enumerate(db_feats):
        sim = cosine_sim(input_feat, db_feat)
        if sim > best_sim and sim > threshold:
            best_sim = sim
            best_idx = i
    if best_idx == -1:
        return None, 0
    return (db_ids[best_idx], db_names[best_idx]), best_sim