Upload main.py

main.py

@@ -0,0 +1,122 @@
+import numpy as np
+import onnxruntime as rt
+import mediapipe as mp
+import cv2
+import os
+import time
+from skimage.transform import SimilarityTransform
+
+
+# ---------------------------------------------------------------------------------------------------------------------
+# 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,
+)
+
+# Initialize Face Detector (For Example Mediapipe)
+FACE_DETECTOR = mp.solutions.face_mesh.FaceMesh(
+    refine_landmarks=True, min_detection_confidence=0.5, min_tracking_confidence=0.5, max_num_faces=1
+)
+
+# Initialize the Face Recognition Model (FaceTransformerOctupletLoss)
+FACE_RECOGNIZER = rt.InferenceSession("FaceTransformerOctupletLoss.onnx", providers=rt.get_available_providers())
+
+
+# ---------------------------------------------------------------------------------------------------------------------
+# FACE CAPTURE
+
+# Capture a frame with your Webcam and store it on disk
+if not os.path.exists("img.jpg"):
+    cap = cv2.VideoCapture(1) # open webcam
+    time.sleep(2) # wait for camera to warm up
+    
+    if not cap.isOpened():
+        raise IOError("Cannot open webcam")
+    
+    ret, img = cap.read() # capture a frame
+    if ret:
+        cv2.imwrite("img.jpg", img) # save the frame
+else:
+    img = cv2.imread("img.jpg") # read the frame from disk
+
+
+# ---------------------------------------------------------------------------------------------------------------------
+# FACE DETECTION
+
+# Process the image with the face detector
+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]
+    )
+
+    # Extract the x and y coordinates of all landmarks
+    all_x_coords = [landmark.x * img.shape[1] for landmark in result.multi_face_landmarks[0].landmark]
+    all_y_coords = [landmark.y * img.shape[0] for landmark in result.multi_face_landmarks[0].landmark]
+
+    # Compute the bounding box of the face
+    x_min, x_max = int(min(all_x_coords)), int(max(all_x_coords))
+    y_min, y_max = int(min(all_y_coords)), int(max(all_y_coords))
+    bbox = [[x_min, y_min], [x_max, y_max]]
+
+else:
+    print("No faces detected")
+    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("img2_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)
+embedding = FACE_RECOGNIZER.run(None, {"input_image": input_image})[0][0]
+
+print("Embedding:", embedding)
+
+# If you have embeddings for several facial images - you can then compute the cosine distance between them and distinguish
+# between different or same people based on a threshold. For example, if the cosine distance is less than 0.5, then the
+# two images are of the same person, otherwise they are of different people. The lower the cosine distance, the more similar
+# the two images are. The cosine distance is a value between 0 and 2, where 0 means the two images are identical and 2 means 
+# the two images are completely different. 
+
+# ---------------------------------------------------------------------------------------------------------------------
+# VISUALIZATION
+
+# Draw Boundingbox on a copy of image
+img_draw = img.copy()
+cv2.rectangle(img_draw, (bbox[0][0], bbox[0][1]), (bbox[1][0], bbox[1][1]), (255, 0, 0), 2)
+
+# Show the detected face on the image
+cv2.imshow("img", img_draw)
+cv2.waitKey(0)
+
+# Show the aligned image
+cv2.imshow("img", img_aligned)
+cv2.waitKey(0)