Untitled
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