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.pull/626/head
Vignesh Skanda
GitHub
parent
a11ccf9c49
commit
883839c206
|
|
@ -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(
|
||||
name='buffalo_l',
|
||||
providers=modules.globals.execution_providers
|
||||
)
|
||||
FACE_ANALYSER = insightface.app.FaceAnalysis(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]:
|
||||
faces = get_face_analyser().get(frame)
|
||||
return min(faces, key=lambda x: x.bbox[0], default=None)
|
||||
# Function to get one face from a frame
|
||||
def get_one_face(frame: Frame) -> Any:
|
||||
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)
|
||||
return faces if faces else None
|
||||
|
||||
def get_many_faces(frame: Frame) -> Any:
|
||||
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
|
||||
|
|
|
|||
Loading…
Reference in New Issue