Add Face Verification API with FastAPI

.gitignore

@@ -0,0 +1 @@
+__pycache__

README.md

@@ -1,5 +1,3 @@
-
-
 # EdgeFace: Efficient Face Recognition Model for Edge Devices
 
 [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/edgeface-efficient-face-recognition-model-for/lightweight-face-recognition-on-lfw)](https://paperswithcode.com/sota/lightweight-face-recognition-on-lfw?p=edgeface-efficient-face-recognition-model-for)
@@ -97,6 +95,37 @@ model.eval()
 | edgeface_xs_gamma_06|   1.77 | 154.00 | 99.73 ± 0.35 | 95.28 ± 1.37 | 91.58 ± 1.42 | 94.71 ± 1.07 | 96.08 ± 0.95 |
 | edgeface_xxs        |   1.24 |  94.72 | 99.57 ± 0.33 | 94.83 ± 0.98 | 90.27 ± 0.93 | 93.63 ± 0.99 | 94.92 ± 1.15 |
 
+# EdgeFace API
+
+EdgeFace serves an API as well - see api.py for more details. You can run the api with the following command:
+
+```bash
+python api.py --port 8000  # 8000 is the default
+```
+
+## API Endpoints
+
+### GET /
+Welcome page for the API
+
+### POST /verify
+Verify two faces for similarity.
+
+Parameters:
+- `image1`: First image file
+- `image2`: Second image file
+- `metric`: Similarity metric (default: "cosine")
+
+Returns JSON with:
+- `distance`: Similarity distance between faces
+- `threshold`: Threshold for matching
+- `match`: Boolean indicating if faces match
+
+You can test the API endpoints through the interactive Swagger UI documentation at:
+```
+http://127.0.0.1:8000/docs
+```
+
 ## Reference
 If you use this repository, please cite the following paper, which is [published](https://ieeexplore.ieee.org/abstract/document/10388036/) in the IEEE Transactions on Biometrics, Behavior, and Identity Science (IEEE T-BIOM). The PDF version of the paper is available as [pre-print on arxiv](https://arxiv.org/pdf/2307.01838v2.pdf). The complete source code for reproducing all experiments in the paper (including training and evaluation) is also publicly available in the [official repository](https://gitlab.idiap.ch/bob/bob.paper.tbiom2023_edgeface).
 

api.py

@@ -0,0 +1,54 @@
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import JSONResponse, HTMLResponse
+from verification import Verifier
+import numpy as np
+from io import BytesIO
+from PIL import Image
+import argparse
+import uvicorn
+import sys
+
+app = FastAPI()
+verifier = Verifier(model_name="edgeface_s_gamma_05")
+
+def read_image_as_numpy(file: UploadFile) -> np.ndarray:
+    try:
+        image = Image.open(BytesIO(file.file.read())).convert("RGB")
+        return np.array(image)
+    except Exception:
+        raise HTTPException(status_code=400, detail="Invalid image file")
+
+@app.get("/", response_class=HTMLResponse)
+def home():
+    return "<h1> Welcome to EdgeFace API </h1>"
+
+@app.post("/verify")
+async def verify_faces(
+    image1: UploadFile = File(...),
+    image2: UploadFile = File(...),
+    metric: str = "cosine"
+):
+    img1_np = read_image_as_numpy(image1)
+    img2_np = read_image_as_numpy(image2)
+
+    try:
+        distance, threshold, is_same = verifier.verify(img1_np, img2_np, metric)
+        return JSONResponse({
+            "distance": float(distance),
+            "threshold": float(threshold),
+            "match": bool(is_same)
+        })
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run the EdgeFace FastAPI server.")
+    parser.add_argument(
+        "--port", type=int, default=8000,
+        help="Port number to run the server on (default: 8000)"
+    )
+    args = parser.parse_args()
+
+    module_name = sys.argv[0].replace(".py", "")
+    uvicorn.run(f"{module_name}:app", host="0.0.0.0", port=args.port, reload=True)

face_alignment/align.py

@@ -1,13 +1,13 @@
 import sys
 import os
-
+import numpy as np
 from face_alignment import mtcnn
 import argparse
 from PIL import Image
 from tqdm import tqdm
 import random
 from datetime import datetime
-mtcnn_model = mtcnn.MTCNN(device='cuda:0', crop_size=(112, 112))
+mtcnn_model = mtcnn.MTCNN(device='cpu', crop_size=(112, 112))
 
 def add_padding(pil_img, top, right, bottom, left, color=(0,0,0)):
     width, height = pil_img.size
@@ -17,16 +17,35 @@ def add_padding(pil_img, top, right, bottom, left, color=(0,0,0)):
     result.paste(pil_img, (left, top))
     return result
 
-def get_aligned_face(image_path, rgb_pil_image=None):
-    if rgb_pil_image is None:
-        img = Image.open(image_path).convert('RGB')
-    else:
-        assert isinstance(rgb_pil_image, Image.Image), 'Face alignment module requires PIL image or path to the image'
+def get_aligned_face(image_input, rgb_pil_image=None):
+    """Get aligned face from various input types.
+
+    Args:
+        image_input: Can be one of:
+            - str: Path to image file
+            - np.ndarray: RGB image array of shape (H,W,3)
+            - None: When using rgb_pil_image parameter
+        rgb_pil_image (PIL.Image, optional): RGB PIL Image
+
+    Returns:
+        PIL.Image: Aligned face or None if detection fails
+    """
+    if rgb_pil_image is not None:
+        assert isinstance(rgb_pil_image, Image.Image), 'rgb_pil_image must be a PIL Image'
         img = rgb_pil_image
+    elif isinstance(image_input, str):
+        img = Image.open(image_input).convert('RGB')
+    elif isinstance(image_input, np.ndarray):
+        assert len(image_input.shape) == 3 and image_input.shape[2] == 3, \
+            'NumPy array must be RGB with shape (H,W,3)'
+        img = Image.fromarray(image_input)
+    else:
+        raise TypeError("Input must be file path, numpy array, or PIL Image")
+
     # find face
     try:
         bboxes, faces = mtcnn_model.align_multi(img, limit=1)
-        face = faces[0]
+        face = faces[0] if faces else None
     except Exception as e:
         print('Face detection Failed due to error.')
         print(e)

requirement.txt

@@ -1,5 +1,13 @@
-pytorch
+torch
 torchvision
 timm
 mxnet
-opencv-python
+opencv-python
+Pillow
+numpy
+scikit-learn
+mxnet
+
+# Web API
+fastapi
+uvicorn[standard]

verification/__init__.py

@@ -0,0 +1 @@
+from .verifier import Verifier

verification/distance_utils.py

@@ -0,0 +1,40 @@
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
+
+def compute_distance(embedding1, embedding2, metric):
+    embedding1 = np.asarray(embedding1)
+    embedding2 = np.asarray(embedding2)
+
+    if metric == "cosine":
+        return find_cosine_distance(embedding1, embedding2)
+    elif metric == "angular":
+        return find_angular_distance(embedding1, embedding2)
+    elif metric == "euclidean":
+        return find_euclidean_distance(embedding1, embedding2)
+    elif metric == "euclidean_l2":
+        norm_axis = None if embedding1.ndim == 1 else 1
+        embedding1 = l2_normalize(embedding1, axis=norm_axis)
+        embedding2 = l2_normalize(embedding2, axis=norm_axis)
+        return find_euclidean_distance(embedding1, embedding2)
+    else:
+        raise ValueError(f"Invalid distance metric: {metric}")
+
+def get_verification_threshold(metric: str) -> float:
+    default_thresholds = {"cosine": 0.40, "euclidean": 4.15, "euclidean_l2": 0.95, "angular": 0.37}
+
+    return default_thresholds.get(metric, 0.4)
+
+
+def l2_normalize(x, axis=None):
+    norm = np.linalg.norm(x, ord=2, axis=axis, keepdims=True)
+    return x / np.clip(norm, a_min=1e-10, a_max=None)
+
+def find_cosine_distance(a, b):
+    return 1 - cosine_similarity(a.reshape(1, -1), b.reshape(1, -1))[0][0]
+
+def find_angular_distance(a, b):
+    cosine_sim = cosine_similarity(a.reshape(1, -1), b.reshape(1, -1))[0][0]
+    return np.arccos(np.clip(cosine_sim, -1.0, 1.0)) / np.pi
+
+def find_euclidean_distance(a, b):
+    return np.linalg.norm(a - b)

verification/verifier.py

@@ -0,0 +1,50 @@
+import numpy as np
+import torch
+from torchvision import transforms
+from backbones import get_model
+from face_alignment import align
+from .distance_utils import compute_distance, get_verification_threshold
+
+class Verifier():
+    def __init__(self, model_name: str = "edgeface_base"):
+        self.model_name = model_name
+        self.checkpoint_path=f'checkpoints/{self.model_name}.pt'
+        self.model = get_model(self.model_name)
+        self.model.load_state_dict(torch.load(self.checkpoint_path, 
+                                              map_location='cpu'))
+
+        self.model.eval()
+        self.preprocess = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        ])
+
+    def extract_embedding(self, image: np.ndarray) -> np.ndarray:
+        aligned_face = align.get_aligned_face(image)
+        input_tensor = self.preprocess(aligned_face).unsqueeze(0)
+        with torch.no_grad():
+            embedding = self.model(input_tensor).squeeze().numpy()
+        return embedding
+
+    def verify(self, img_1: np.ndarray, img_2: np.ndarray):
+        img_1 = align.get_aligned_face(img_1)
+        img_2 = align.get_aligned_face(img_2)
+        transformed_input_1 = self.transform(img_1).unsqueeze(0)
+        transformed_input_2 = self.transform(img_2).unsqueeze(0)
+
+        with torch.no_grad():
+            embedding_1 = self.model(transformed_input_1).squeeze().numpy()
+            embedding_2 = self.model(transformed_input_2).squeeze().numpy()
+
+        print(embedding_1.shape, embedding_2.shape)
+
+    def verify(self, image1: np.ndarray, image2: np.ndarray, dist_metric: str = "cosine") -> bool:
+        embedding1 = self.extract_embedding(image1)
+        embedding2 = self.extract_embedding(image2)
+
+        distance = compute_distance(embedding1, embedding2, dist_metric)
+        threshold = get_verification_threshold(dist_metric)
+        is_same_person = distance <= threshold
+
+        print(f"Distance ({dist_metric}): {distance:.6f} | Threshold: {threshold} | Match: {is_same_person}")
+        return distance, threshold, is_same_person