eval.py

import os
import numpy as np
import pandas as pd
from PIL import Image
from huggingface_model_utils import load_model_from_local_path
from torchvision.transforms import Compose, ToTensor, Normalize
import torch
import inspect
from facenet_pytorch import MTCNN


# Load models
device = 'cuda' if torch.cuda.is_available() else 'cpu'
aligner = load_model_from_local_path('model/minchul/cvlface_DFA_resnet50').to(device)
fr_model = load_model_from_local_path('model/minchul/cvlface_adaface_vit_base_webface4m').to(device)

# Hàm chuyển ảnh PIL sang tensor
def pil_to_input(pil_image, device):
    trans = Compose([ToTensor(), Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
    return trans(pil_image).unsqueeze(0).to(device)

# Hàm trích xuất đặc trưng từ ảnh
def get_feat(pil_image):
    """Trích xuất đặc trưng từ ảnh PIL."""
    input_tensor = pil_to_input(pil_image, device)
    aligned_x, _, aligned_ldmks, _, _, _ = aligner(input_tensor)

    # Kiểm tra xem mô hình có cần keypoints không
    input_signature = inspect.signature(fr_model.model.net.forward)
    if input_signature.parameters.get('keypoints') is not None:
        feat = fr_model(aligned_x, aligned_ldmks)
    else:
        feat = fr_model(aligned_x)
    return feat

# Hàm tính toàn bộ cosine similarity giữa đặc trưng đầu vào và tất cả các đặc trưng trong cơ sở dữ liệu
def compute_cosine_similarity(feat_input, feats_db):
    """Tính toàn bộ cosine similarity giữa đặc trưng đầu vào và tất cả các đặc trưng trong cơ sở dữ liệu."""
    return torch.nn.functional.cosine_similarity(feat_input, feats_db)

# Hàm lấy ID bằng cách so sánh với cơ sở dữ liệu
def get_id(pil_image, feats_db, db_ids, threshold=0.3):
    """So sánh ảnh với cơ sở dữ liệu và trả về ID nếu tìm thấy."""
    # Trích xuất đặc trưng từ ảnh đầu vào
    feat_input = get_feat(pil_image)

    # Tính toán cosine similarity giữa đặc trưng đầu vào và tất cả các đặc trưng trong cơ sở dữ liệu
    similarities = compute_cosine_similarity(feat_input, feats_db)

    # Lấy ra giá trị cosine similarity lớn nhất và index tương ứng
    max_sim, idx_max = similarities.max(0)

    if max_sim.item() >= threshold:
        best_match_id = db_ids[idx_max.item()]
        return {"id": best_match_id, "status": "found"}
    else:
        return {"id": None, "status": "not_found"}

# Hàm chính
def main():
    # Đọc cơ sở dữ liệu enroll một lần
    enroll_csv = 'face_features_vit.csv'
    if not os.path.exists(enroll_csv):
        print(f"Không tìm thấy file cơ sở dữ liệu {enroll_csv}.")
        return

    # Load và chuyển đổi đặc trưng thành tensor một lần duy nhất
    enroll_data = pd.read_csv(enroll_csv, dtype={'id': str})
    
    # Chuyển đặc trưng từ chuỗi thành tensor để tính toán nhanh hơn
    feats_db = torch.stack([torch.tensor(eval(feat), device=device) for feat in enroll_data['feat'].values])
    db_ids = enroll_data['id'].tolist()

    # Khởi tạo các chỉ số
    TP, FP, FN, TN = 0, 0, 0, 0
    threshold = 0.5  # Ngưỡng cosine similarity

    # Đường dẫn tới thư mục inout
    folder_path2 = '/home/quoc14/Code/HeThongNhanDang/_vts_users_v1/vts_in_out_faces_v1'
    processed_folders = 0
    total_folders = len(os.listdir(folder_path2))

    # Duyệt qua các folder trong inout
    for folder in os.listdir(folder_path2):
        folder_path = os.path.join(folder_path2, folder)
        if os.path.isdir(folder_path):
            image_files = [f for f in os.listdir(folder_path) if f.endswith('.jpg')]
            for image_file in image_files:
                image_path = os.path.join(folder_path, image_file)
                pil_image = Image.open(image_path)

                # Nhận diện khuôn mặt bằng cách so sánh đặc trưng
                check_result = get_id(pil_image, feats_db, db_ids, threshold)
                
                if check_result['status'] == 'found':
                    matched_id = check_result['id']
                    if str(folder) == matched_id:
                        TP += 1  # Nhận diện đúng
                    else:
                        FP += 1  # Nhận diện nhầm
                else:
                    if str(folder) in db_ids:
                        FN += 1  # Lẽ ra phải nhận diện nhưng không
                    else:
                        TN += 1  # Không có trong hệ thống và hệ thống không nhận diện

            processed_folders += 1
            print(f"Đã xử lý {processed_folders}/{total_folders} folder")

    # Tính toán FNIR, FPIR, ACC
    FNIR = FN / (TP + FN) if (TP + FN) > 0 else 0
    FPIR = FP / (FP + TN) if (FP + TN) > 0 else 0
    ACC = (TP + TN) / (TP + TN + FP + FN) if (TP + TN + FP + FN) > 0 else 0

    # Lưu kết quả vào file CSV chỉ một lần
    results = [{
        'Model': 'VIT',
        'ACC': ACC,
        'FNIR': FNIR,
        'FPIR': FPIR
    }]
    df_results = pd.DataFrame(results)
    df_results.to_csv('evaluation_results_vts.csv', index=False)
    print(FN, FP, TP, TN)
    print(f"Accuracy: {ACC}, FNIR: {FNIR}, FPIR: {FPIR}")

if __name__ == "__main__":
    main()