File size: 5,301 Bytes
9e14838 | 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | # Since several subjects can be found within a video, it is necessary to cluster them into groups based on similarity.
# This operation is carried out in the following code with additional attention to maintaining the temporal coherence of faces.
# The extracted faces are reorganised into consecutive sequences of similar faces so as to be more suitable for network processing.
import argparse
import os
import glob
import torch
import numpy as np
import pandas as pd
import shutil
from functools import partial
from multiprocessing.pool import Pool
from numpy.linalg import norm
from PIL import Image
from torchvision import transforms
from facenet_pytorch import InceptionResnetV1, fixed_image_standardization
from collections import OrderedDict
from sklearn.cluster import KMeans
from torch.utils.data.dataloader import DataLoader
from progress.bar import ChargingBar
from utils import preprocess_images, _generate_connected_components
seed = 42
def move_files(face_paths):
src_path, dst_path = face_paths
os.makedirs(os.path.dirname(dst_path), exist_ok=True)
shutil.move(src_path, dst_path)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--faces_path', default="../../datasets/ForgeryNet/faces", type=str,
help='Path of folder containing train/val/test with extracted cropped faces to be clustered.')
parser.add_argument('--gpu_id', default=0, type=int,
help='ID of GPU to be used.')
parser.add_argument('--similarity_threshold', default=0.45, type=float,
help='Threshold to discard faces with high distance.')
parser.add_argument('--valid_cluster_size_ratio', default=0.20, type=int,
help='Valid cluster size ratio.')
parser.add_argument('--workers', default=40, type=int,
help='Number of data loader workers.')
opt = parser.parse_args()
print(opt)
# Get all the paths of the videos to be clustered
for dataset in os.listdir(opt.faces_path):
dataset_path = os.path.join(opt.faces_path, dataset)
if not os.path.isdir(dataset_path):
continue
print()
print("Clustering videos in ", dataset_path)
set_paths = glob.glob(f'{dataset_path}/*/**/*.mp4', recursive=True)
excluded_videos = []
for path in set_paths:
if os.path.exists(os.path.join(path, "0")):
excluded_videos.append(path)
set_paths = [video_path for video_path in set_paths if video_path not in excluded_videos]
print("Excluded already clustered videos: ", len(excluded_videos))
# For each video in each set, perform faces clustering
bar = ChargingBar('Clustered videos', max=(len(set_paths)))
for path in set_paths:
# Read all faces, load them into a dictionary
faces_files = [face_file for face_file in os.listdir(path) if not os.path.isdir(os.path.join(path, face_file))]
faces_files = sorted(faces_files, key=lambda x:(int(x.split("_")[0]), int(os.path.splitext(x)[0].split("_")[1])))
mapping = {}
faces = []
for index, face_file in enumerate(faces_files):
face_path = os.path.join(path, face_file)
frame_number = int(os.path.splitext(face_file)[0].split("_")[0])
face = Image.open(face_path)
faces.append(face)
mapping[index] = face_path
# Extract the embeddings
embeddings_extractor = InceptionResnetV1(pretrained='vggface2').eval().to(opt.gpu_id)
faces = [preprocess_images(face) for face in faces]
faces = np.stack([np.uint8(face) for face in faces])
faces = torch.as_tensor(faces)
faces = faces.permute(0, 3, 1, 2).float()
faces = fixed_image_standardization(faces)
face_recognition_input = faces.cuda()
embeddings = []
embeddings = embeddings_extractor(face_recognition_input).detach().cpu().numpy()
# Clustering
valid_cluster_size = int(len(mapping) * opt.valid_cluster_size_ratio)
similarities = np.dot(np.array(embeddings), np.array(embeddings).T)
components = _generate_connected_components(
similarities, similarity_threshold=opt.similarity_threshold
)
components = [sorted(component) for component in components]
mapped_components = []
for identity_index, component in enumerate(components):
for index in component:
src_path = mapping[index]
folder_path = os.path.dirname(src_path)
file_name = os.path.basename(src_path)
dst_path = os.path.join(folder_path, str(identity_index), file_name)
mapped_components.append((src_path, dst_path))
# Organize the clusters inside the folder
with Pool(processes=opt.workers) as p:
for v in p.imap_unordered(move_files, mapped_components):
continue
bar.next()
print()
|