Impl. + tests

.dockerignore

@@ -0,0 +1,4 @@
+.*
+testdata
+test_service.py
+__pycache__

.gitignore

@@ -0,0 +1,4 @@
+.venv
+.env
+__pycache__
+.vscode

Dockerfile

@@ -0,0 +1,11 @@
+FROM python:3.7
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+COPY . .
+
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,65 @@
+import os
+from typing import Union
+
+from fastapi import (
+    Depends,
+    FastAPI,
+    File,
+    HTTPException,
+    Query,
+    Response,
+    UploadFile,
+    status,
+)
+
+from model import Model
+
+app = FastAPI()
+
+model = Model(
+    os.getenv("MODEL_REPO_ID", ""),
+    os.getenv("MODEL_FILENAME", ""),
+    os.getenv("HF_TOKEN", ""),
+)
+
+
+async def validate_token(
+    token: Union[str, None] = Query(default=None),
+):
+    if token is None:
+        raise HTTPException(status.HTTP_401_UNAUTHORIZED, "No token provided")
+    if token != os.getenv("CLIENT_TOKEN"):
+        raise HTTPException(status.HTTP_401_UNAUTHORIZED, "Invalid token")
+    return token
+
+
+@app.post("/embed", status_code=status.HTTP_200_OK)
+async def calculate_embedding(
+    image: UploadFile = File(...), _: str = Depends(validate_token)
+):
+    try:
+        image_content = await image.read()
+        pred = model.predict(model.preprocess(image_content))
+        return {"embedding": pred.tolist()}
+    except Exception as e:
+        return Response(
+            content=str(e), status_code=status.HTTP_500_INTERNAL_SERVER_ERROR
+        )
+
+
+@app.post("/similarity", status_code=status.HTTP_200_OK)
+async def calculate_similarity_score(
+    image1: UploadFile = File(...),
+    image2: UploadFile = File(...),
+    _: str = Depends(validate_token),
+):
+    try:
+        image1_content = await image1.read()
+        image2_content = await image2.read()
+        pred1 = model.predict(model.preprocess(image1_content))
+        pred2 = model.predict(model.preprocess(image2_content))
+        return {"score": float(model.distance(pred1, pred2))}
+    except Exception as e:
+        return Response(
+            content=str(e), status_code=status.HTTP_500_INTERNAL_SERVER_ERROR
+        )

model.py

@@ -0,0 +1,165 @@
+import threading
+
+import cv2
+import mediapipe as mp
+import numpy as np
+import pandas as pd
+from huggingface_hub import hf_hub_download
+from tensorflow import keras
+
+
+class Model:
+    def __init__(self, model_repo_id: str, model_filename: str, hf_token: str):
+        self.landmark_extractor = LandmarkExtractor()
+        self.trained_model = self.create_prod_model()
+        self.trained_model.compile(
+            keras.optimizers.Adam(0.0005),
+            loss=self.triplet_loss_init(0.2),
+            metrics=[self.triplet_accuracy],
+        )
+        custom_objects = {
+            "triplet_loss": self.triplet_loss_init(0.2),
+            "triplet_accuracy": self.triplet_accuracy,
+            "K": keras.backend,
+            "keras": keras,
+        }
+        with keras.utils.custom_object_scope(custom_objects):
+            weights = keras.models.load_model(
+                hf_hub_download(model_repo_id, model_filename, token=hf_token)
+            ).get_weights()
+            self.trained_model.set_weights(weights)
+
+    @staticmethod
+    def fec_net(inputs):
+        x = keras.layers.Dense(1024, activation="relu")(inputs)
+        x = keras.layers.Dense(512, activation="relu")(x)
+        x = keras.layers.Dense(512, activation="relu")(x)
+        x = keras.layers.Dropout(0.5)(x)
+        x = keras.layers.Dense(16)(x)
+        outputs = keras.layers.Lambda(keras.backend.l2_normalize)(x)
+        return outputs
+
+    @classmethod
+    def create_prod_model(cls):
+        inputs = keras.layers.Input(shape=(478 * 3,))
+        outputs = cls.fec_net(inputs)
+        return keras.Model(inputs=inputs, outputs=outputs)
+
+    @staticmethod
+    def triplet_loss_init(alpha):
+        def triplet_loss(y_true, y_pred):
+            dimensions = 16
+            batch_size = y_pred.shape[0]
+            s1 = y_pred[:, :dimensions]
+            s2 = y_pred[:, dimensions : 2 * dimensions]
+            d = y_pred[:, 2 * dimensions :]
+
+            s1_s2 = keras.backend.sum(keras.backend.square(s1 - s2), axis=1)
+            s1_d = keras.backend.sum(keras.backend.square(s1 - d), axis=1)
+            s2_d = keras.backend.sum(keras.backend.square(s2 - d), axis=1)
+
+            loss = keras.backend.maximum(
+                0.0, s1_s2 - s1_d + alpha
+            ) + keras.backend.maximum(0.0, s1_s2 - s2_d + alpha)
+            loss = keras.backend.mean(loss)
+
+            return loss
+
+        return triplet_loss
+
+    @staticmethod
+    def triplet_accuracy(y_true, y_pred):
+        dimensions = 16
+        s1 = y_pred[:, :dimensions]
+        s2 = y_pred[:, dimensions : 2 * dimensions]
+        d = y_pred[:, 2 * dimensions :]
+
+        s1_s2 = keras.backend.sqrt(
+            keras.backend.sum(keras.backend.square(s1 - s2), axis=1)
+        )
+        s1_d = keras.backend.sqrt(
+            keras.backend.sum(keras.backend.square(s1 - d), axis=1)
+        )
+        s2_d = keras.backend.sqrt(
+            keras.backend.sum(keras.backend.square(s2 - d), axis=1)
+        )
+
+        s1_match = keras.backend.less(s1_s2, s1_d)
+        s2_match = keras.backend.less(s1_s2, s2_d)
+        match = keras.backend.cast(
+            keras.backend.all(keras.backend.stack([s1_match, s2_match]), axis=0),
+            "float32",
+        )
+
+        acc = keras.backend.mean(match, axis=0)
+
+        return acc
+
+    @classmethod
+    def load_model(cls, model_path: str):
+        custom_objects = {
+            "triplet_loss": cls.triplet_loss_init(0.2),
+            "triplet_accuracy": cls.triplet_accuracy,
+        }
+        with keras.utils.custom_object_scope(custom_objects):
+            return keras.models.load_model(model_path)
+
+    def predict(self, x: np.ndarray):
+        return self.trained_model.predict(x)
+
+    def preprocess(self, image: bytes) -> np.ndarray:
+        array_repr = np.asarray(bytearray(image), dtype=np.uint8)
+        decoded_img = cv2.imdecode(array_repr, flags=cv2.IMREAD_COLOR)
+        landmarks = self.landmark_extractor.extract_landmarks_flat(decoded_img)
+        return pd.DataFrame([landmarks]).to_numpy()
+
+    @staticmethod
+    def distance(x1: np.ndarray, x2: np.ndarray):
+        return np.linalg.norm(x1 - x2, ord=2)
+
+
+class LandmarkExtractor:
+    def __init__(self):
+        self.lock = threading.Lock()
+        self.face_mesh = mp.solutions.face_mesh.FaceMesh(
+            static_image_mode=True,
+            max_num_faces=1,
+            refine_landmarks=True,
+            min_detection_confidence=0.5,
+        )
+
+    @staticmethod
+    def __normalise_landmarks(landmarks):
+        left_eye_idx = 33
+        right_eye_idx = 263
+        left_eye_v = [landmarks[left_eye_idx].x, landmarks[left_eye_idx].y]
+        right_eye_v = [landmarks[right_eye_idx].x, landmarks[right_eye_idx].y]
+        xy_norm = np.linalg.solve(
+            np.array([left_eye_v, right_eye_v]).T,
+            np.array([[lmk.x, lmk.y] for lmk in landmarks]).T,
+        ).T
+        for i, lmk in enumerate(landmarks):
+            lmk.x = xy_norm[i][0]
+            lmk.y = xy_norm[i][1]
+        return landmarks
+
+    def extract_landmarks(self, img: np.ndarray):
+        results = self.face_mesh.process(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+        if results.multi_face_landmarks is None:
+            return None
+        return self.__normalise_landmarks(results.multi_face_landmarks[0].landmark)
+
+    def extract_landmarks_flat(self, img: np.ndarray):
+        landmarks = self.extract_landmarks(img)
+        if landmarks is None:
+            return None
+        flat_landmarks = {}
+        for i, landmark in enumerate(landmarks):
+            flat_landmarks.update(
+                {
+                    f"{i}_x": landmark.x,
+                    f"{i}_y": landmark.y,
+                    f"{i}_z": landmark.z,
+                }
+            )
+        return flat_landmarks

requirements.txt

@@ -0,0 +1,11 @@
+huggingface_hub
+tensorflow
+numpy
+fastapi==0.74.*
+requests==2.27.*
+uvicorn[standard]==0.17.*
+python-multipart
+mediapipe
+pandas
+pytest
+python-dotenv

test_service.py

@@ -0,0 +1,89 @@
+import os
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+import pytest
+from fastapi.testclient import TestClient
+
+
+@pytest.fixture
+def client():
+    from app import app
+
+    clnt = TestClient(app)
+    return clnt
+
+
+@pytest.mark.parametrize(
+    "image_path",
+    [
+        pytest.param("testdata/face_pic.jpeg", id="happy flow"),
+    ],
+)
+def test_embed(client: TestClient, image_path):
+    with open(image_path, "rb") as image_file:
+        response = client.post(
+            "/embed",
+            files=[("image", (os.path.basename(image_path), image_file))],
+            params={"token": os.getenv("CLIENT_TOKEN")},
+        )
+    assert response.status_code == 200, response.text
+    assert "embedding" in response.json()
+
+
+@pytest.mark.parametrize(
+    "image1_path, image2_path, exp_score",
+    [
+        pytest.param(
+            "testdata/face_pic.jpeg",
+            "testdata/face_pic.jpeg",
+            0,
+            id="same pic",
+        ),
+        pytest.param(
+            "testdata/face_pic.jpeg",
+            "testdata/face_pic3.jpeg",
+            0.00026,
+            id="similar expression",
+        ),
+        pytest.param(
+            "testdata/face_pic.jpeg",
+            "testdata/face_pic2.jpeg",
+            2,
+            id="different expression",
+        ),
+    ],
+)
+def test_similarity(client: TestClient, image1_path, image2_path, exp_score):
+    with open(image1_path, "rb") as image1_file, open(image2_path, "rb") as image2_file:
+        response = client.post(
+            "/similarity",
+            files=[
+                ("image1", (os.path.basename(image1_path), image1_file)),
+                ("image2", (os.path.basename(image2_path), image2_file)),
+            ],
+            params={"token": os.getenv("CLIENT_TOKEN")},
+        )
+    assert response.status_code == 200
+    assert "score" in response.json()
+    if exp_score is not None:
+        assert response.json()["score"] == pytest.approx(exp_score, abs=1e-4)
+
+
+def test_authentication(client: TestClient):
+    response = client.post(
+        "/embed",
+        files=[("image", ("face_pic.jpeg", b""))],
+    )
+    assert response.status_code == 401
+    assert response.json()["detail"] == "No token provided"
+
+    response = client.post(
+        "/embed",
+        files=[("image", ("face_pic.jpeg", b""))],
+        params={"token": "wrong_token"},
+    )
+    assert response.status_code == 401
+    assert response.json()["detail"] == "Invalid token"