Update face_analyser.py

Updated the code file. Added few explanatory comments in the code so that the user can understand the code pretty easily.

modules/face_analyser.py

@@ -1,27 +1,199 @@
-from typing import Any, Optional
+import os
+import shutil
+from typing import Any
 import insightface
 
+import cv2
+import numpy as np
 import modules.globals
+from tqdm import tqdm
 from modules.typing import Frame
+from modules.cluster_analysis import find_cluster_centroids, find_closest_centroid
+from modules.utilities import get_temp_directory_path, create_temp, extract_frames, clean_temp, get_temp_frame_paths
+from pathlib import Path
 
-FACE_ANALYSER: Optional[insightface.app.FaceAnalysis] = None
+# Initialize the face analyzer
+FACE_ANALYSER = None
 
-def get_face_analyser() -> insightface.app.FaceAnalysis:
+# Function to get the face analyzer
+def get_face_analyser() -> Any:
     global FACE_ANALYSER
 
+     # If the face analyzer is not initialized, initialize it
     if FACE_ANALYSER is None:
-        FACE_ANALYSER = insightface.app.FaceAnalysis(
+        FACE_ANALYSER = insightface.app.FaceAnalysis(name='buffalo_l', providers=modules.globals.execution_providers)
-            name='buffalo_l', 
-            providers=modules.globals.execution_providers
-        )
         FACE_ANALYSER.prepare(ctx_id=0, det_size=(640, 640))
-    
     return FACE_ANALYSER
 
-def get_one_face(frame: Frame) -> Optional[Any]:
+# Function to get one face from a frame
-    faces = get_face_analyser().get(frame)
+def get_one_face(frame: Frame) -> Any:
-    return min(faces, key=lambda x: x.bbox[0], default=None)
+    face = get_face_analyser().get(frame)
+    try:
+        # If there are multiple faces, return the one with the smallest bounding box
+        return min(face, key=lambda x: x.bbox[0])
+    except ValueError:
+        return None
 
-def get_many_faces(frame: Frame) -> Optional[Any]:
+
-    faces = get_face_analyser().get(frame)
+def get_many_faces(frame: Frame) -> Any:
-    return faces if faces else None
+    try:
+        return get_face_analyser().get(frame)
+    except IndexError:
+        return None
+
+# Function to check if the source-target map has valid entries
+def has_valid_map() -> bool:
+    for map in modules.globals.souce_target_map:
+        if "source" in map and "target" in map:
+            return True
+    return False
+
+# Function to get the default source face
+def default_source_face() -> Any:
+    for map in modules.globals.souce_target_map:
+        if "source" in map:
+            return map['source']['face']
+    return None
+
+# Function to simplify the source-target map
+def simplify_maps() -> Any:
+    centroids = []
+    faces = []
+    for map in modules.globals.souce_target_map:
+        if "source" in map and "target" in map:
+            centroids.append(map['target']['face'].normed_embedding)
+            faces.append(map['source']['face'])
+
+    modules.globals.simple_map = {'source_faces': faces, 'target_embeddings': centroids}
+    return None
+
+# Function to add a blank map to the source-target map
+def add_blank_map() -> Any:
+    try:
+        max_id = -1
+        if len(modules.globals.souce_target_map) > 0:
+            max_id = max(modules.globals.souce_target_map, key=lambda x: x['id'])['id']
+
+        modules.globals.souce_target_map.append({
+                'id' : max_id + 1
+                })
+    except ValueError:
+        return None
+
+# Function to get unique faces from a target image
+def get_unique_faces_from_target_image() -> Any:
+    try:
+        modules.globals.souce_target_map = []
+        target_frame = cv2.imread(modules.globals.target_path)
+        many_faces = get_many_faces(target_frame)
+        i = 0
+
+        for face in many_faces:
+            x_min, y_min, x_max, y_max = face['bbox']
+            modules.globals.souce_target_map.append({
+                'id' : i, 
+                'target' : {
+                            'cv2' : target_frame[int(y_min):int(y_max), int(x_min):int(x_max)],
+                            'face' : face
+                            }
+                })
+            i = i + 1
+    except ValueError:
+        return None
+    
+    
+def get_unique_faces_from_target_video() -> Any:
+    try:
+        modules.globals.souce_target_map = []
+        frame_face_embeddings = []
+        face_embeddings = []
+    
+        print('Creating temp resources...')
+        clean_temp(modules.globals.target_path)
+        create_temp(modules.globals.target_path)
+        print('Extracting frames...')
+        extract_frames(modules.globals.target_path)
+
+        temp_frame_paths = get_temp_frame_paths(modules.globals.target_path)
+
+        i = 0
+        for temp_frame_path in tqdm(temp_frame_paths, desc="Extracting face embeddings from frames"):
+            temp_frame = cv2.imread(temp_frame_path)
+            many_faces = get_many_faces(temp_frame)
+
+            for face in many_faces:
+                face_embeddings.append(face.normed_embedding)
+            
+            frame_face_embeddings.append({'frame': i, 'faces': many_faces, 'location': temp_frame_path})
+            i += 1
+
+        centroids = find_cluster_centroids(face_embeddings)
+
+        for frame in frame_face_embeddings:
+            for face in frame['faces']:
+                closest_centroid_index, _ = find_closest_centroid(centroids, face.normed_embedding)
+                face['target_centroid'] = closest_centroid_index
+
+        for i in range(len(centroids)):
+            modules.globals.souce_target_map.append({
+                'id' : i
+            })
+
+            temp = []
+            for frame in tqdm(frame_face_embeddings, desc=f"Mapping frame embeddings to centroids-{i}"):
+                temp.append({'frame': frame['frame'], 'faces': [face for face in frame['faces'] if face['target_centroid'] == i], 'location': frame['location']})
+
+            modules.globals.souce_target_map[i]['target_faces_in_frame'] = temp
+
+        # dump_faces(centroids, frame_face_embeddings)
+        default_target_face()
+    except ValueError:
+        return None
+    
+
+def default_target_face():
+    for map in modules.globals.souce_target_map:
+        best_face = None
+        best_frame = None
+        for frame in map['target_faces_in_frame']:
+            if len(frame['faces']) > 0:
+                best_face = frame['faces'][0]
+                best_frame = frame
+                break
+
+        for frame in map['target_faces_in_frame']:
+            for face in frame['faces']:
+                if face['det_score'] > best_face['det_score']:
+                    best_face = face
+                    best_frame = frame
+
+        x_min, y_min, x_max, y_max = best_face['bbox']
+
+        target_frame = cv2.imread(best_frame['location'])
+        map['target'] = {
+                        'cv2' : target_frame[int(y_min):int(y_max), int(x_min):int(x_max)],
+                        'face' : best_face
+                        }
+
+# Function to dump faces to a temporary directory
+def dump_faces(centroids: Any, frame_face_embeddings: list):
+    temp_directory_path = get_temp_directory_path(modules.globals.target_path)
+
+    for i in range(len(centroids)):
+        if os.path.exists(temp_directory_path + f"/{i}") and os.path.isdir(temp_directory_path + f"/{i}"):
+            shutil.rmtree(temp_directory_path + f"/{i}")
+        # Create a new directory for the current centroid    
+        Path(temp_directory_path + f"/{i}").mkdir(parents=True, exist_ok=True)
+
+        for frame in tqdm(frame_face_embeddings, desc=f"Copying faces to temp/./{i}"):
+            temp_frame = cv2.imread(frame['location'])
+
+            # Initialize a counter for the faces in the frame
+            j = 0
+            for face in frame['faces']:
+                if face['target_centroid'] == i:
+                    x_min, y_min, x_max, y_max = face['bbox']
+
+                    if temp_frame[int(y_min):int(y_max), int(x_min):int(x_max)].size > 0:
+                        cv2.imwrite(temp_directory_path + f"/{i}/{frame['frame']}_{j}.png", temp_frame[int(y_min):int(y_max), int(x_min):int(x_max)])
+                j += 1