File size: 2,623 Bytes
a8bbf69 f89004e dd85347 a8bbf69 2980986 a8bbf69 f89004e b6b1cef f89004e 7a6fee1 24110f1 7a00801 f89004e 24110f1 b6b1cef 24110f1 b6b1cef 24110f1 b6b1cef 24110f1 1705770 24110f1 f89004e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import torch
import numpy as np
from facenet_pytorch import InceptionResnetV1
from sklearn.cluster import DBSCAN
import os
import shutil
import mediapipe as mp
import cv2
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = InceptionResnetV1(pretrained='vggface2').eval().to(device)
def get_face_embedding(face_img):
face_tensor = torch.tensor(face_img).permute(2, 0, 1).unsqueeze(0).float() / 255
face_tensor = (face_tensor - 0.5) / 0.5
face_tensor = face_tensor.to(device)
with torch.no_grad():
embedding = model(face_tensor)
return embedding.cpu().numpy().flatten()
def cluster_faces(embeddings):
if len(embeddings) < 2:
print("Not enough faces for clustering. Assigning all to one cluster.")
return np.zeros(len(embeddings), dtype=int)
X = np.stack(embeddings)
dbscan = DBSCAN(eps=0.3, min_samples=5, metric='cosine')
clusters = dbscan.fit_predict(X)
if np.all(clusters == -1):
print("DBSCAN assigned all to noise. Considering as one cluster.")
return np.zeros(len(embeddings), dtype=int)
return clusters
def organize_faces_by_person(embeddings_by_frame, clusters, aligned_faces_folder, organized_faces_folder):
for (frame_num, embedding), cluster in zip(embeddings_by_frame.items(), clusters):
person_folder = os.path.join(organized_faces_folder, f"person_{cluster}")
os.makedirs(person_folder, exist_ok=True)
src = os.path.join(aligned_faces_folder, f"frame_{frame_num}_face.jpg")
dst = os.path.join(person_folder, f"frame_{frame_num}_face.jpg")
shutil.copy(src, dst)
def draw_facial_landmarks(image, landmarks):
mp_face_mesh = mp.solutions.face_mesh
mp_drawing = mp.solutions.drawing_utils
# Adjust these values to make the drawing thinner and less obtrusive
connection_drawing_spec = mp_drawing.DrawingSpec(color=(255, 0, 0), thickness=1)
if landmarks:
# Create a blank image with the same dimensions as the input image
connections_layer = np.zeros_like(image)
# Draw connections on the blank image
mp_drawing.draw_landmarks(
image=connections_layer,
landmark_list=landmarks,
connections=mp_face_mesh.FACEMESH_TESSELATION,
landmark_drawing_spec=None, # Set to None to avoid drawing points
connection_drawing_spec=connection_drawing_spec)
# Blend the connections layer with the original image
alpha = 0.2
image = cv2.addWeighted(image, 1, connections_layer, alpha, 0)
return image |