File size: 2,226 Bytes
a8bbf69 f89004e dd85347 a8bbf69 2980986 a8bbf69 f89004e 7a6fee1 4f2ae83 f89004e 4f2ae83 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 |
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_drawing = mp.solutions.drawing_utils
# Adjust these values to make the drawing thinner and less obtrusive
landmark_drawing_spec = mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=1, circle_radius=1)
if landmarks:
for landmark in landmarks.landmark:
x = int(landmark.x * image.shape[1])
y = int(landmark.y * image.shape[0])
cv2.circle(image, (x, y), landmark_drawing_spec.circle_radius, landmark_drawing_spec.color, landmark_drawing_spec.thickness)
return image |