Add vector_search function for pipeline aggregation (#30)

* Add vector_search function for pipeline aggregation

* Added Vector search

* chore: Add error handling for no match found in face recognition

* Refactor face recognition code to use Facenet512 model for better accuracy

* chore: Add tests for recognizing faces with no match found

* feat: Add test for vector search functionality

API/database.py

@@ -22,3 +22,29 @@ def find_one_and_delete(self, collection, query):
 
     def update_one(self, collection, query, update):
         return self.db[collection].update_one(query, update)
+
+    # add a function for pipeline aggregation vector search
+    def vector_search(self, collection, embedding):
+
+        result = self.db[collection].aggregate([
+            {
+                "$vectorSearch": {
+                "index": "vector_index",
+                "path": "face_embedding",
+                "queryVector": embedding,
+                "numCandidates": 20,
+                "limit": 20
+                }
+            }, {
+                '$project': {
+                '_id': 0, 
+                'Name': 1,
+                'Image': 1,
+                'score': {
+                    '$meta': 'vectorSearchScore'
+                    }
+                }
+            }
+            ])
+        result_arr = [i for i in result]
+        return result_arr

API/route.py

@@ -8,22 +8,27 @@
 
 from bson import ObjectId
 from deepface import DeepFace
-from fastapi import APIRouter, HTTPException, Response
+from fastapi import APIRouter, HTTPException, Response, UploadFile, File
 from matplotlib import pyplot as plt
 from PIL import Image
 from pydantic import BaseModel
 
 from API.database import Database
 from API.utils import init_logging_config
+from dotenv import load_dotenv
 
+load_dotenv()
 init_logging_config()
 
+MONGO_URI = os.getenv("MONGO_URL1")
 router = APIRouter()
 
 
 client = Database()
+client2 = Database(MONGO_URI, "FaceRec")
 
 collection = "faceEntries"
+collection2 = "ImageDB"
 
 
 # Models  for the data to be sent and received by the server
@@ -79,16 +84,16 @@ async def create_new_faceEntry(Employee: Employee):
         plt.imsave(f"Images/Faces/{Name}.jpg", face_image_data[0]["face"])
         logging.info(f"Face saved {Name}")
         embedding = DeepFace.represent(
-            image_filename, model_name="Facenet", detector_backend="mtcnn"
+            image_filename, model_name="Facenet512", detector_backend="mtcnn"
         )
         embeddings.append(embedding)
         logging.info(f"Embedding created Embeddings for {Name}")
         os.remove(image_filename)
 
     logging.debug(f"About to insert Embeddings: {embeddings}")
     # Store the data in the database
-    client.insert_one(
-        collection,
+    client2.insert_one(
+        collection2,
         {
             "EmployeeCode": EmployeeCode,
             "Name": Name,
@@ -267,3 +272,39 @@ async def delete_employees(EmployeeCode: int):
     client.find_one_and_delete(collection, {"EmployeeCode": EmployeeCode})
 
     return {"Message": "Successfully Deleted"}
+
+
+@router.post("/recognize_face", response_class=Response)
+async def recognize_face(Face: UploadFile = File(...)):
+    """
+    Recognize a face from the provided image.
+
+    Args:
+        Face (UploadFile): The image file to be recognized.
+
+    Returns:
+        Response: A response object containing the recognized employee information in JSON format.
+
+    Raises:
+        HTTPException: If an internal server error occurs.
+    """
+    logging.info("Recognizing Face")
+    try:
+        img_data = await Face.read()
+        with open("temp.png", "wb") as f:
+            f.write(img_data)
+
+        embedding = DeepFace.represent(img_path="temp.png", model_name="Facenet512", detector_backend="mtcnn")
+        result = client2.vector_search(collection2, embedding[0]['embedding'])
+        logging.info(f"Result: {result[0]['Name']}, {result[0]['score']}")
+        os.remove("temp.png")
+        if result[0]['score'] < 0.5:
+            return Response(status_code=404, content=json.dumps({"message": "No match found"}))
+    except Exception as e:
+        logging.error(f"Error: {e}")
+        os.remove("temp.png")
+        raise HTTPException(status_code=500, detail="Internal server error")
+    return Response(
+        content=bytes(json.dumps(result[0], default=str), "utf-8"),
+        media_type="application/json",
+    )

CHANGELOG.md

@@ -49,4 +49,17 @@
 - Resolved various bugs and issues identified during the testing process.
 
 ### Removed
-- Removed deprecated code and unused dependencies from the project.
+- Removed deprecated code and unused dependencies from the project.
+
+## [0.1.4] - 2024-03-16 - 23:00
+
+### Added
+- Implemented a new `recognize_Face` endpoint in [`route.py`](API/route.py). This endpoint accepts a base64 string as input, converts it into embeddings, and performs a vector search query on the MongoDB Atlas database. Changes made by @Devasy23.
+- Added a new `vector_search` function in [`database.py`](API/database.py). This function performs a vector similarity search on the MongoDB Atlas database using Euclidean distance as the similarity measure. Changes made by @Devasy23.
+- Updated [`index.ipynb`](index.ipynb) to include examples and usage of the new `recognize_Face` endpoint and `vector_search` function. Changes made by @Devasy23.
+
+### Changed
+- Updated the `Database` class in [`database.py`](API/database.py) to include the new `vector_search` function. Changes made by @Devasy23.
+
+### Fixed
+- Resolved various bugs and issues identified during the implementation and testing of the new features. Fixes made by @Devasy23.

testing/test_database.py

@@ -0,0 +1,36 @@
+import base64
+import logging
+from unittest.mock import MagicMock, patch
+
+
+import pytest
+from fastapi.testclient import TestClient
+
+from API.database import Database
+from API.route import router
+from API.utils import init_logging_config
+
+init_logging_config()
+
+def test_vector_search():
+    mock_result = [
+        {
+            "Name": "Test1",
+            "Image": "encoded_string1",
+            "score": 0.8
+        },
+        {
+            "Name": "Test2",
+            "Image": "encoded_string2",
+            "score": 0.7
+        }
+    ]
+
+    mock_vector_search = MagicMock(return_value=mock_result)
+
+    with patch("API.database.Database.vector_search", mock_vector_search):
+        embedding = [0.1, 0.2, 0.3]
+        result = Database.vector_search("collection_name", embedding)
+
+        assert result == mock_result
+        mock_vector_search.assert_called_once_with("collection_name", embedding)

testing/test_face_endpoints.py

@@ -162,3 +162,37 @@ def test_delete_face():
     response = client.delete("/delete/1")
     assert response.status_code == 200
     assert response.json() == {"Message": "Successfully Deleted"}
+
+@pytest.mark.run(order=6)
+def test_recognize_face_fail():
+    mock_doc = {
+        "Image": "encoded_string2",
+        "Name": "Test2",
+        "score": 0.0,   
+    }
+    with patch("API.database.Database.vector_search", return_value=[mock_doc]):
+
+        with open("./test-faces/devansh.jpg", "rb") as image_file:
+            response = client.post(
+                "/recognize_face",
+                files={"Face":image_file},
+            )
+        assert response.status_code == 404
+        assert response.json() == {'message': 'No match found'}
+
+@pytest.mark.run(order=7)
+def test_recognize_face_success():
+    mock_doc = {
+        "Image": "encoded_string2",
+        "Name": "Test2",
+        "score": 1.0,   
+    }
+    with patch("API.database.Database.vector_search", return_value=[mock_doc]):
+
+        with open("./test-faces/devansh.jpg", "rb") as image_file:
+            response = client.post(
+                "/recognize_face",
+                files={"Face":image_file},
+            )
+        assert response.status_code == 200
+        assert response.json() == {'Name': 'Test2', 'Image':'encoded_string2','score': 1.0}