src/model.py

import cv2
import insightface
from insightface.app.common import Face
from insightface.model_zoo import model_zoo

import numpy as np

class ImageEmbedding:
    def __init__(self):

        det_model_path = 'models/buffalo_l/det_10g.onnx'
        rec_model_path = 'models/antelopev2/glintr100.onnx'

        self.det_model = model_zoo.get_model(f'{det_model_path}')
        self.rec_model = model_zoo.get_model(f'{rec_model_path}')

        self.det_model.prepare(ctx_id=0, input_size=(640, 640), det_thres=0.5)

    def cosine_similarity(self, embeddings1, embeddings2):
        return np.dot(embeddings1, embeddings2) / (np.linalg.norm(embeddings1) * np.linalg.norm(embeddings2))

    def get_face_embedding(self, img):
        bboxes, kpss = self.det_model.detect(img, max_num=0, metric='defualt')

        bbox = bboxes[0, :4]
        det_score = bboxes[0, 4]
        kps = kpss[0]
        face = Face(bbox=bbox, kps=kps, det_score=det_score)

        self.rec_model.get(img, face)
        embeddings = face.normed_embedding

        bbox = face.bbox
        # Convert float coordinates to integers
        bbox = [int(coord) for coord in bbox]

        # Draw the rectangle on the image
        cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2)

        return embeddings, img