Merge pull request #36 from bassamadnan/main

Face recognition

Author: Pranav0-0Aggarwal

backend/app/config/settings.py

@@ -12,3 +12,4 @@
 ALBUM_DATABASE_PATH = "app/database/albums.db"
 MAPPINGS_DATABASE_PATH = "app/database/yolo_mapping.db"
 FACES_DATABASE_PATH = "app/database/faces.db"
+CLUSTERS_DATABASE_PATH = "app/database/clusters.db"

backend/app/database/albums.py

@@ -3,22 +3,26 @@
 
 from app.config.settings import ALBUM_DATABASE_PATH
 from app.utils.wrappers import image_exists, album_exists
-from app.database.images import get_id_from_path, get_path_from_id
+from app.utils.path_id_mapping import get_id_from_path, get_path_from_id
+
 
 def create_albums_table():
     conn = sqlite3.connect(ALBUM_DATABASE_PATH)
     cursor = conn.cursor()
-    cursor.execute("""
+    cursor.execute(
+        """
         CREATE TABLE IF NOT EXISTS albums (
             album_name TEXT PRIMARY KEY,
             image_ids TEXT,
             description TEXT,
             date_created TIMESTAMP DEFAULT CURRENT_TIMESTAMP
         )
-    """)
+    """
+    )
     conn.commit()
     conn.close()
 
+
 def create_album(album_name, description=None):
     conn = sqlite3.connect(ALBUM_DATABASE_PATH)
     cursor = conn.cursor()
@@ -30,11 +34,14 @@ def create_album(album_name, description=None):
         conn.close()
         raise ValueError(f"Album '{album_name}' already exists")
 
-    cursor.execute("INSERT INTO albums (album_name, image_ids, description) VALUES (?, ?, ?)",
-                   (album_name, json.dumps([]), description))
+    cursor.execute(
+        "INSERT INTO albums (album_name, image_ids, description) VALUES (?, ?, ?)",
+        (album_name, json.dumps([]), description),
+    )
     conn.commit()
     conn.close()
 
+
 @album_exists
 def delete_album(album_name):
     conn = sqlite3.connect(ALBUM_DATABASE_PATH)
@@ -43,13 +50,14 @@ def delete_album(album_name):
     conn.commit()
     conn.close()
 
+
 @album_exists
-@image_exists
 def add_photo_to_album(album_name, image_path):
     conn = sqlite3.connect(ALBUM_DATABASE_PATH)
     cursor = conn.cursor()
 
     image_id = get_id_from_path(image_path)
+    # print("GOT IMAGE ID", image_id, flush=True)
     if image_id is None:
         conn.close()
         raise ValueError(f"Image '{image_path}' not found in the database")
@@ -59,28 +67,13 @@ def add_photo_to_album(album_name, image_path):
     image_ids = json.loads(result[0])
     if image_id not in image_ids:
         image_ids.append(image_id)
-        cursor.execute("UPDATE albums SET image_ids = ? WHERE album_name = ?",
-                       (json.dumps(image_ids), album_name))
+        cursor.execute(
+            "UPDATE albums SET image_ids = ? WHERE album_name = ?",
+            (json.dumps(image_ids), album_name),
+        )
         conn.commit()
     conn.close()
 
-@album_exists
-def add_photos_to_album(album_name, image_paths):
-    conn = sqlite3.connect(ALBUM_DATABASE_PATH)
-    cursor = conn.cursor()
-
-    cursor.execute("SELECT image_ids FROM albums WHERE album_name = ?", (album_name,))
-    result = cursor.fetchone()
-    if result:
-        existing_ids = json.loads(result[0])
-        new_ids = [get_id_from_path(path) for path in image_paths if get_id_from_path(path) is not None]
-        updated_ids = list(set(existing_ids + new_ids))
-
-        cursor.execute("UPDATE albums SET image_ids = ? WHERE album_name = ?",
-                       (json.dumps(updated_ids), album_name))
-
-    conn.commit()
-    conn.close()
 
 @album_exists
 def get_album_photos(album_name):
@@ -96,6 +89,7 @@ def get_album_photos(album_name):
         return [get_path_from_id(image_id) for image_id in image_ids]
     return None
 
+
 @album_exists
 @image_exists
 def remove_photo_from_album(album_name, image_path):
@@ -112,28 +106,60 @@ def remove_photo_from_album(album_name, image_path):
     image_ids = json.loads(result[0])
     if image_id in image_ids:
         image_ids.remove(image_id)
-        cursor.execute("UPDATE albums SET image_ids = ? WHERE album_name = ?",
-                       (json.dumps(image_ids), album_name))
+        cursor.execute(
+            "UPDATE albums SET image_ids = ? WHERE album_name = ?",
+            (json.dumps(image_ids), album_name),
+        )
         conn.commit()
     conn.close()
 
+
 def get_all_albums():
     conn = sqlite3.connect(ALBUM_DATABASE_PATH)
     cursor = conn.cursor()
 
     cursor.execute("SELECT album_name, image_ids FROM albums")
     results = cursor.fetchall()
-    albums = [{"album_name": name, "image_paths": [get_path_from_id(id) for id in json.loads(ids)]} for name, ids in results]
+    albums = [
+        {
+            "album_name": name,
+            "image_paths": [get_path_from_id(id) for id in json.loads(ids)],
+        }
+        for name, ids in results
+    ]
 
     conn.close()
     return albums
 
+
 @album_exists
 def edit_album_description(album_name, new_description):
     conn = sqlite3.connect(ALBUM_DATABASE_PATH)
     cursor = conn.cursor()
-    
-    cursor.execute("UPDATE albums SET description = ? WHERE album_name = ?",
-                   (new_description, album_name))
+
+    cursor.execute(
+        "UPDATE albums SET description = ? WHERE album_name = ?",
+        (new_description, album_name),
+    )
+    conn.commit()
+    conn.close()
+
+
+def remove_image_from_all_albums(image_id):
+    conn = sqlite3.connect(ALBUM_DATABASE_PATH)
+    cursor = conn.cursor()
+
+    cursor.execute("SELECT album_name, image_ids FROM albums")
+    albums = cursor.fetchall()
+
+    for album_name, image_ids_json in albums:
+        image_ids = json.loads(image_ids_json)
+        if image_id in image_ids:
+            image_ids.remove(image_id)
+            cursor.execute(
+                "UPDATE albums SET image_ids = ? WHERE album_name = ?",
+                (json.dumps(image_ids), album_name),
+            )
+
     conn.commit()
     conn.close()

backend/app/database/faces.py

@@ -1,80 +1,131 @@
 import sqlite3
 import json
 import numpy as np
-from app.config.settings import FACES_DATABASE_PATH
-from app.database.images import get_id_from_path, get_path_from_id
+from app.config.settings import FACES_DATABASE_PATH, IMAGES_DATABASE_PATH
+
 
 def create_faces_table():
     conn = sqlite3.connect(FACES_DATABASE_PATH)
     cursor = conn.cursor()
-    cursor.execute("""
+    cursor.execute(
+        """
         CREATE TABLE IF NOT EXISTS faces (
             id INTEGER PRIMARY KEY AUTOINCREMENT,
             image_id INTEGER,
             embeddings TEXT,
-            FOREIGN KEY (image_id) REFERENCES image_id_mapping(id)
+            FOREIGN KEY (image_id) REFERENCES image_id_mapping(id) ON DELETE CASCADE
         )
-    """)
+    """
+    )
     conn.commit()
     conn.close()
 
+
 def insert_face_embeddings(image_path, embeddings):
+    from app.database.images import get_id_from_path
+
     conn = sqlite3.connect(FACES_DATABASE_PATH)
     cursor = conn.cursor()
-    
+
     image_id = get_id_from_path(image_path)
     if image_id is None:
         conn.close()
         raise ValueError(f"Image '{image_path}' not found in the database")
-    
+
     embeddings_json = json.dumps([emb.tolist() for emb in embeddings])
-    
-    cursor.execute("""
+
+    cursor.execute(
+        """
         INSERT OR REPLACE INTO faces (image_id, embeddings)
         VALUES (?, ?)
-    """, (image_id, embeddings_json))
-    
+    """,
+        (image_id, embeddings_json),
+    )
+
     conn.commit()
     conn.close()
 
+
 def get_face_embeddings(image_path):
+    from app.database.images import get_id_from_path
+
     conn = sqlite3.connect(FACES_DATABASE_PATH)
     cursor = conn.cursor()
-    
+
     image_id = get_id_from_path(image_path)
     if image_id is None:
         conn.close()
         return None
-    
-    cursor.execute("""
+
+    cursor.execute(
+        """
         SELECT embeddings FROM faces
         WHERE image_id = ?
-    """, (image_id,))
-    
+    """,
+        (image_id,),
+    )
+
     result = cursor.fetchone()
     conn.close()
-    
+
     if result:
         embeddings_json = result[0]
         embeddings = np.array(json.loads(embeddings_json))
         return embeddings
     else:
         return None
 
+
 def get_all_face_embeddings():
+    from app.database.images import get_path_from_id
+
     conn = sqlite3.connect(FACES_DATABASE_PATH)
     cursor = conn.cursor()
-    
-    cursor.execute("""
+
+    cursor.execute(
+        """
         SELECT image_id, embeddings FROM faces
-    """)
-    
+    """
+    )
+
     results = cursor.fetchall()
     all_embeddings = []
     for image_id, embeddings_json in results:
         image_path = get_path_from_id(image_id)
         embeddings = np.array(json.loads(embeddings_json))
         all_embeddings.append({"image_path": image_path, "embeddings": embeddings})
-    
+    print("returning")
     conn.close()
     return all_embeddings
+
+
+def delete_face_embeddings(image_id):
+    conn = sqlite3.connect(FACES_DATABASE_PATH)
+    cursor = conn.cursor()
+
+    cursor.execute("DELETE FROM faces WHERE image_id = ?", (image_id,))
+
+    conn.commit()
+    conn.close()
+
+
+def cleanup_face_embeddings():
+    conn_images = sqlite3.connect(IMAGES_DATABASE_PATH)
+    conn_faces = sqlite3.connect(FACES_DATABASE_PATH)
+    cursor_images = conn_images.cursor()
+    cursor_faces = conn_faces.cursor()
+
+    cursor_faces.execute("SELECT DISTINCT image_id FROM faces")
+    face_image_ids = set(row[0] for row in cursor_faces.fetchall())
+
+    cursor_images.execute("SELECT id FROM image_id_mapping")
+    valid_image_ids = set(row[0] for row in cursor_images.fetchall())
+
+    orphaned_ids = face_image_ids - valid_image_ids
+
+    for orphaned_id in orphaned_ids:
+        cursor_faces.execute("DELETE FROM faces WHERE image_id = ?", (orphaned_id,))
+
+    conn_faces.commit()
+    conn_images.close()
+    conn_faces.close()