Untitled
unknown
python
2 years ago
3.6 kB
20
Indexable
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}")Editor is loading...
Leave a Comment