Untitled

mail@pastecode.io avatar
unknown
python
8 months ago
3.6 kB
12
Indexable
Never
import sys
sys.path.append("../..")

import cv2
import torch
from fastapi_camera_detection.face_detection.face_detection import FaceDetector
from libs.utils.utils import preprocess_image, load_facenet_model
import libs.utils.utils as utils
from libs.dbs_connection.dbs_connection import DatabaseOperations
import time
from libs.utils.schemas import Embedding, Match
import numpy as np
import faiss
import psycopg2

class FaceRecognition:
    def __init__(self, db_connection):
        self.db_connection = db_connection
        self.person_names = []
        self.image_paths = []
        self.logger = utils.Logger("FACE_RECOGNITION")
        self.model = load_facenet_model()

    def face_recognition(self, aligned_face) -> Embedding:
        start_time = time.time()
        preprocess_face = preprocess_image(aligned_face)
        with torch.no_grad():
            embedding = self.model(preprocess_face).cpu().detach().numpy()

        end_time = time.time()

        elapsed_time = end_time - start_time
        # Time is second and milisecond eg: 3.123s
        readable_time = "{:.3f} seconds".format(elapsed_time)
        self.logger.info(f"Recognition Process Time: {readable_time}")
        return Embedding(vector=embedding)

    def similarity_searching(self, embedding) -> int:
        try:
            start_time = time.time()
            faiss.normalize_L2(embedding)
            # Build Faiss index if it doesn't exist
            ids = self.db_connection.build_faiss_index()
            # Perform a search in the Faiss index
            D, I = self.db_connection.index.search(embedding, 1)

            # I contains the nearest neighbor index for each query embedding (here, we only have one query)
            nearest_neighbor_index = I[0][0]
            cosine_sim = D[0][0]

            # Retrieve the corresponding ID from the ids list
            nearest_neighbor_id = ids[nearest_neighbor_index]
            person_name, person_image_path = self.db_connection.get_infor_by_id(nearest_neighbor_id)
            end_time = time.time()

            # Convert time to a readable format
            elapsed_time = end_time - start_time
            readable_time = "{:.3f} seconds".format(elapsed_time)

            self.logger.info(f"Search Time: {readable_time} - Name: {person_name} - ID: {nearest_neighbor_id} - (Distance: {cosine_sim}) - Image Path: {person_image_path}")
            # Return the ID of the nearest embedding
            print(nearest_neighbor_id)
            print(type(nearest_neighbor_id))
            match = Match(id=nearest_neighbor_id,
                          name=person_name,
                          image_path=person_image_path,
                          distance=cosine_sim)

            return match,embedding
        except Exception as e:
            self.logger.error(f"Error in similarity searching: {str(e)}")
            raise e


if __name__ == "__main__":
    db_connection = DatabaseOperations()
    # Face detector
    detector = FaceDetector()

    # Get image
    img = cv2.imread('/Users/trietlethongminh/Downloads/INT01S02_Code/ivs_api/IVS_CameraAttendance/fastapi_camera_registration/images/obama/230616115528-barack-obama-2022-file.jpg')

    # Detect face
    face = detector.detect(img)
    aligned_face = detector.crop_align(img, face).face
    # Face recognition
    recognizer = FaceRecognition(db_connection)

    # Get embeddings
    embedding = recognizer.face_recognition(aligned_face)

    # Similarity search
    match = recognizer.similarity_searching(embedding.vector)

    print(f"Person: {match.name}")
Leave a Comment