CodeAnhGao

import gradio as gr
from face_engine.engine import FaceEngine
from PIL import Image
import torch
from torchvision.transforms import Compose, ToTensor, Normalize
import glob
from face_engine.huggingface_model_utils import load_model_from_local_path
import cv2
import numpy as np

model_extract = load_model_from_local_path("face_engine/models/minchul/cvlface_adaface_ir101_webface4m")
model_aligner = load_model_from_local_path("face_engine/models/minchul/cvlface_DFA_resnet50")

def pil_to_input(pil_image):
    trans = Compose([ToTensor(), Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
    return trans(pil_image).unsqueeze(0)

def get_feature(pil_image):
    input_tensor = pil_to_input(pil_image)
    aligned_x, _, aligned_ldmks, _, _, _ = model_aligner(input_tensor)
    print(aligned_x.shape)
    img = aligned_x.numpy()

    feature = model_extract(aligned_x)
    feature_normm = torch.nn.functional.normalize(feature)
    print(feature_normm.shape)
    return feature_normm, img
  

#imglink1 = "vts\\vts_in_out_faces_v1\\vts_in_out_faces_v1\\001067012926\\66a1b96dd53d525db9f90eb9.jpg"
#imglink2 = "vts\\vts_in_out_faces_v1\\vts_in_out_faces_v1\\001067012926\\66a1b277d53d525db9e959d8.jpg"

imglink1 = "D:\\_vts_users_v1\\vts_enroll_images_v1\\0a532a40-c8a7-4ad2-aad4-4368c65dec18\\0.jpg"
imglink2 = "D:\\_vts_users_v1\\vts_in_out_faces_v1\\040089018059\\66a9f5eed53d525db936c902.jpg"



img1 = Image.open(imglink1)
img2 = Image.open(imglink2)


f1, im1 = get_feature(img1)
f2, im2 = get_feature(img2)

im1 = np.transpose(im1[0], (1, 2, 0))
im2 = np.transpose(im2[0], (1, 2, 0))       

im1 = (im1+0.5)
im2 = (im2+0.5)
im1 = cv2.cvtColor(im1, cv2.COLOR_RGB2BGR)
im2 = cv2.cvtColor(im2, cv2.COLOR_RGB2BGR)
cv2.imshow('im1.jpg', im1)


distEclid = (f1 - f2).pow(2).sum().sqrt().item()
cossim = torch.nn.functional.cosine_similarity(f1, f2).item()

list_image = glob.glob("D:\\_vts_users_v1\\vts_in_out_faces_v1\\040089018059\\*.jpg")
for imglink in list_image:
    img2 = Image.open(imglink)
    f2, im2 = get_feature(img2)
    distEclid = (f1 - f2).pow(2).sum().sqrt().item()
    cossim = torch.nn.functional.cosine_similarity(f1, f2).item() #0.5
    print(imglink.split("\\")[-1].split('/')[-1], cossim, distEclid)

    im2 = np.transpose(im2[0], (1, 2, 0))
    im2 = cv2.cvtColor(im2, cv2.COLOR_RGB2BGR)
    cv2.imshow('im2.jpg', im2)
    cv2.waitKey()