routers/fecnet_router.py

import os

import numpy as np
from fastapi import APIRouter, Depends, File, UploadFile, status

from auth import validate_token
from model.fecnet import FECNetModel
from schema import EmbeddingResponse, SimilarityResponse

router = APIRouter(
    prefix="/fecnet",
    tags=["fecnet"],
    dependencies=[Depends(validate_token)],
)

model = FECNetModel(os.getenv("HF_TOKEN", ""))


@router.get(
    "/embed",
    status_code=status.HTTP_200_OK,
    response_model=EmbeddingResponse,
)
async def calculate_embedding(
    image: UploadFile = File(...),
    should_extract_face: bool = False,
):
    image_arr = np.asarray(bytearray(await image.read()), dtype=np.uint8) # type: ignore
    rep = model.embed_image(image_arr, should_extract_face)
    return EmbeddingResponse(embedding=rep.tolist())


@router.get(
    "/similarity",
    status_code=status.HTTP_200_OK,
    response_model=SimilarityResponse,
)
async def calculate_similarity_score(
    image1: UploadFile = File(...),
    image2: UploadFile = File(...),
    should_extract_face: bool = False,
):
    image1_arr = np.asarray(bytearray(await image1.read()), dtype=np.uint8) # type: ignore
    image2_arr = np.asarray(bytearray(await image2.read()), dtype=np.uint8) # type: ignore
    rep1 = model.embed_image(image1_arr, should_extract_face)
    rep2 = model.embed_image(image2_arr, should_extract_face)
    return SimilarityResponse(score=np.linalg.norm(rep1, rep2))