martlgap / octuplet-loss Goto Github PK

import numpy as np
import onnxruntime as rt
import mediapipe as mp
import cv2
import time
from skimage.transform import SimilarityTransform
from scipy.spatial.distance import cosine

# ---------------------------------------------------------------------------------------------------------------------
# INITIALIZATIONS

# Target landmark coordinates for alignment (used in training)
LANDMARKS_TARGET = np.array(
   [
       [38.2946, 51.6963],
       [73.5318, 51.5014],
       [56.0252, 71.7366],
       [41.5493, 92.3655],
       [70.7299, 92.2041],
   ],
   dtype=np.float32,
)

def compare(img_path1, img_path2):
   img1_embedding = infer(img_path1)
   img2_embedding = infer(img_path2)

   cos_dist = cosine(img1_embedding, img2_embedding)

   if cos_dist <= 0.5:
       print(f'{img_path1} and {img_path2} are the same')
   else:
       print(f'{img_path1} and {img_path2} are different')

   print(f'cosine distance = {cos_dist}')

def infer(img_path):
   img = cv2.imread(img_path)

   # Process the image with the face detector
   FACE_DETECTOR = mp.solutions.face_mesh.FaceMesh(
       refine_landmarks=True, min_detection_confidence=0.5, min_tracking_confidence=0.5, max_num_faces=1
   )
   result = FACE_DETECTOR.process(img)

   if result.multi_face_landmarks:
       # Select 5 Landmarks (Eye Centers, Nose Tip, Left Mouth Corner, Right Mouth Corner)
       five_landmarks = np.asarray(result.multi_face_landmarks[0].landmark)[[470, 475, 1, 57, 287]]

       # Extract the x and y coordinates of the landmarks of interest
       landmarks = np.asarray(
           [[landmark.x * img.shape[1], landmark.y * img.shape[0]] for landmark in five_landmarks]
       )

   else:
       print(f"No faces detected in {img_path}")
       exit()

   # ---------------------------------------------------------------------------------------------------------------------
   # FACE ALIGNMENT

   # Align Image with the 5 Landmarks
   tform = SimilarityTransform()
   tform.estimate(landmarks, LANDMARKS_TARGET)
   tmatrix = tform.params[0:2, :]
   img_aligned = cv2.warpAffine(img, tmatrix, (112, 112), borderValue=0.0)
   # safe to disk
   cv2.imwrite(f"{img_path}_aligned.jpg", img_aligned)

   # ---------------------------------------------------------------------------------------------------------------------
   # FACE RECOGNITION

   # Inference face embeddings with onnxruntime
   input_image = (np.asarray([img_aligned]).astype(np.float32)).clip(0.0, 255.0).transpose(0, 3, 1, 2)
   FACE_RECOGNIZER = rt.InferenceSession("FaceTransformerOctupletLoss.onnx", providers=rt.get_available_providers())
   embedding = FACE_RECOGNIZER.run(None, {"input_image": input_image})[0][0]

   return embedding

if __name__ == "__main__":
   ds = "David_Schwimmer.jpg"
   ja = "jennifer_aniston.jpg"

   compare(ja, ds)