File size: 2,464 Bytes
2f18963
b87deef
 
 
cd7b3bb
790a68a
515d778
b87deef
 
790a68a
b87deef
067d3bb
 
 
b87deef
515d778
2f18963
515d778
b87deef
 
 
 
 
 
 
 
515d778
b87deef
 
 
2f18963
b87deef
 
 
 
515d778
 
 
 
790a68a
515d778
b87deef
515d778
b87deef
 
 
2f18963
 
b87deef
cd7b3bb
b87deef
 
 
 
 
 
515d778
 
 
 
790a68a
515d778
b87deef
 
 
 
 
 
2f18963
 
b87deef
2f18963
 
cd7b3bb
 
 
 
 
 
b87deef
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
import hmac
import os
from typing import Union

import numpy as np
from fastapi import Depends, FastAPI, File, HTTPException, Response, UploadFile, status
from fastapi.security import APIKeyQuery

from model import Model
from schema import EmbeddingResponse, SimilarityResponse

app = FastAPI(
    title="Facial Expression Embedding Service",
)

api_key = APIKeyQuery(name="token", auto_error=False)
client_token: str = os.getenv("CLIENT_TOKEN", "")

model = Model(
    os.getenv("MODEL_REPO_ID", ""),
    os.getenv("MODEL_FILENAME", ""),
    os.getenv("HF_TOKEN", ""),
)


async def validate_token(
    token: Union[str, None] = Depends(api_key),
):
    if token is None:
        raise HTTPException(status.HTTP_401_UNAUTHORIZED, "No token provided")
    if not hmac.compare_digest(token, client_token):
        raise HTTPException(status.HTTP_401_UNAUTHORIZED, "Invalid token")
    return token


@app.post(
    "/embed",
    status_code=status.HTTP_200_OK,
    dependencies=[Depends(validate_token)],
    response_model=EmbeddingResponse,
)
async def calculate_embedding(
    image: UploadFile = File(...),
):
    try:
        image_content = await image.read()
        if isinstance(image_content, str):
            image_content = image_content.encode()
        pred = model.predict(model.preprocess(image_content))
        return EmbeddingResponse(embedding=pred[0].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,
    dependencies=[Depends(validate_token)],
    response_model=SimilarityResponse,
)
async def calculate_similarity_score(
    image1: UploadFile = File(...),
    image2: UploadFile = File(...),
):
    try:
        image1_content = await image1.read()
        if isinstance(image1_content, str):
            image1_content = image1_content.encode()
        image2_content = await image2.read()
        if isinstance(image2_content, str):
            image2_content = image2_content.encode()
        pred = model.predict(
            np.vstack(
                [model.preprocess(image1_content), model.preprocess(image2_content)]
            )
        )
        return SimilarityResponse(score=float(model.distance(pred[0], pred[1])))
    except Exception as e:
        return Response(
            content=str(e), status_code=status.HTTP_500_INTERNAL_SERVER_ERROR
        )