-
Notifications
You must be signed in to change notification settings - Fork 98
/
test_model.py
270 lines (228 loc) · 11.5 KB
/
test_model.py
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
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
"""
Video Face Manipulation Detection Through Ensemble of CNNs
Image and Sound Processing Lab - Politecnico di Milano
Nicolò Bonettini
Edoardo Daniele Cannas
Sara Mandelli
Luca Bondi
Paolo Bestagini
"""
import argparse
import gc
from collections import OrderedDict
from pathlib import Path
import albumentations as A
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm
from architectures import fornet
from architectures.fornet import FeatureExtractor
from isplutils import utils, split
from isplutils.data import FrameFaceDatasetTest
def main():
# Args
parser = argparse.ArgumentParser()
parser.add_argument('--testsets', type=str, help='Testing datasets', nargs='+', choices=split.available_datasets,
required=True)
parser.add_argument('--testsplits', type=str, help='Test split', nargs='+', default=['val', 'test'],
choices=['train', 'val', 'test'])
parser.add_argument('--dfdc_faces_df_path', type=str, action='store',
help='Path to the Pandas Dataframe obtained from extract_faces.py on the DFDC dataset. '
'Required for training/validating on the DFDC dataset.')
parser.add_argument('--dfdc_faces_dir', type=str, action='store',
help='Path to the directory containing the faces extracted from the DFDC dataset. '
'Required for training/validating on the DFDC dataset.')
parser.add_argument('--ffpp_faces_df_path', type=str, action='store',
help='Path to the Pandas Dataframe obtained from extract_faces.py on the FF++ dataset. '
'Required for training/validating on the FF++ dataset.')
parser.add_argument('--ffpp_faces_dir', type=str, action='store',
help='Path to the directory containing the faces extracted from the FF++ dataset. '
'Required for training/validating on the FF++ dataset.')
# Specify trained model path
parser.add_argument('--model_path', type=Path, help='Full path of the trained model', required=True)
# Common params
parser.add_argument('--batch', type=int, help='Batch size to fit in GPU memory', default=128)
parser.add_argument('--workers', type=int, help='Num workers for data loaders', default=6)
parser.add_argument('--device', type=int, help='GPU id', default=0)
parser.add_argument('--debug', action='store_true', help='Debug flag', )
parser.add_argument('--num_video', type=int, help='Number of real-fake videos to test')
parser.add_argument('--results_dir', type=Path, help='Output folder',
default='results/')
parser.add_argument('--override', action='store_true', help='Override existing results', )
args = parser.parse_args()
device = torch.device('cuda:{}'.format(args.device)) if torch.cuda.is_available() else torch.device('cpu')
num_workers: int = args.workers
batch_size: int = args.batch
max_num_videos_per_label: int = args.num_video # number of real-fake videos to test
model_path: Path = args.model_path
results_dir: Path = args.results_dir
debug: bool = args.debug
override: bool = args.override
test_sets = args.testsets
test_splits = args.testsplits
dfdc_df_path = args.dfdc_faces_df_path
ffpp_df_path = args.ffpp_faces_df_path
dfdc_faces_dir = args.dfdc_faces_dir
ffpp_faces_dir = args.ffpp_faces_dir
# get arguments from the model path
face_policy = str(model_path).split('face-')[1].split('_')[0]
patch_size = int(str(model_path).split('size-')[1].split('_')[0])
net_name = str(model_path).split('net-')[1].split('_')[0]
model_name = '_'.join(model_path.with_suffix('').parts[-2:])
# Load net
net_class = getattr(fornet, net_name)
# load model
print('Loading model...')
state_tmp = torch.load(model_path, map_location='cpu')
if 'net' not in state_tmp.keys():
state = OrderedDict({'net': OrderedDict()})
[state['net'].update({'model.{}'.format(k): v}) for k, v in state_tmp.items()]
else:
state = state_tmp
net: FeatureExtractor = net_class().eval().to(device)
incomp_keys = net.load_state_dict(state['net'], strict=True)
print(incomp_keys)
print('Model loaded!')
# val loss per-frame
criterion = nn.BCEWithLogitsLoss(reduction='none')
# Define data transformers
test_transformer = utils.get_transformer(face_policy, patch_size, net.get_normalizer(), train=False)
# datasets and dataloaders (from train_binclass.py)
print('Loading data...')
# Check if paths for DFDC and FF++ extracted faces and DataFrames are provided
for dataset in test_sets:
if dataset.split('-')[0] == 'dfdc' and (dfdc_df_path is None or dfdc_faces_dir is None):
raise RuntimeError('Specify DataFrame and directory for DFDC faces for testing!')
elif dataset.split('-')[0] == 'ff' and (ffpp_df_path is None or ffpp_faces_dir is None):
raise RuntimeError('Specify DataFrame and directory for FF++ faces for testing!')
splits = split.make_splits(dfdc_df=dfdc_df_path, ffpp_df=ffpp_df_path, dfdc_dir=dfdc_faces_dir,
ffpp_dir=ffpp_faces_dir, dbs={'train': test_sets, 'val': test_sets, 'test': test_sets})
train_dfs = [splits['train'][db][0] for db in splits['train']]
train_roots = [splits['train'][db][1] for db in splits['train']]
val_roots = [splits['val'][db][1] for db in splits['val']]
val_dfs = [splits['val'][db][0] for db in splits['val']]
test_dfs = [splits['test'][db][0] for db in splits['test']]
test_roots = [splits['test'][db][1] for db in splits['test']]
# Output paths
out_folder = results_dir.joinpath(model_name)
out_folder.mkdir(mode=0o775, parents=True, exist_ok=True)
# Samples selection
if max_num_videos_per_label and max_num_videos_per_label > 0:
dfs_out_train = [select_videos(df, max_num_videos_per_label) for df in train_dfs]
dfs_out_val = [select_videos(df, max_num_videos_per_label) for df in val_dfs]
dfs_out_test = [select_videos(df, max_num_videos_per_label) for df in test_dfs]
else:
dfs_out_train = train_dfs
dfs_out_val = val_dfs
dfs_out_test = test_dfs
# Extractions list
extr_list = []
# Append train and validation set first
if 'train' in test_splits:
for idx, dataset in enumerate(test_sets):
extr_list.append(
(dfs_out_train[idx], out_folder.joinpath(dataset + '_train.pkl'), train_roots[idx], dataset + ' TRAIN')
)
if 'val' in test_splits:
for idx, dataset in enumerate(test_sets):
extr_list.append(
(dfs_out_val[idx], out_folder.joinpath(dataset + '_val.pkl'), val_roots[idx], dataset + ' VAL')
)
if 'test' in test_splits:
for idx, dataset in enumerate(test_sets):
extr_list.append(
(dfs_out_test[idx], out_folder.joinpath(dataset + '_test.pkl'), test_roots[idx], dataset + ' TEST')
)
for df, df_path, df_root, tag in extr_list:
if override or not df_path.exists():
print('\n##### PREDICT VIDEOS FROM {} #####'.format(tag))
print('Real frames: {}'.format(sum(df['label'] == False)))
print('Fake frames: {}'.format(sum(df['label'] == True)))
print('Real videos: {}'.format(df[df['label'] == False]['video'].nunique()))
print('Fake videos: {}'.format(df[df['label'] == True]['video'].nunique()))
dataset_out = process_dataset(root=df_root, df=df, net=net, criterion=criterion,
patch_size=patch_size,
face_policy=face_policy, transformer=test_transformer,
batch_size=batch_size,
num_workers=num_workers, device=device, )
df['score'] = dataset_out['score'].astype(np.float32)
df['loss'] = dataset_out['loss'].astype(np.float32)
print('Saving results to: {}'.format(df_path))
df.to_pickle(str(df_path))
if debug:
plt.figure()
plt.title(tag)
plt.hist(df[df.label == True].score, bins=100, alpha=0.6, label='FAKE frames')
plt.hist(df[df.label == False].score, bins=100, alpha=0.6, label='REAL frames')
plt.legend()
del (dataset_out)
del (df)
gc.collect()
if debug:
plt.show()
print('Completed!')
def process_dataset(df: pd.DataFrame,
root: str,
net: FeatureExtractor,
criterion,
patch_size: int,
face_policy: str,
transformer: A.BasicTransform,
batch_size: int,
num_workers: int,
device: torch.device,
) -> dict:
if isinstance(device, (int, str)):
device = torch.device(device)
dataset = FrameFaceDatasetTest(
root=root,
df=df,
size=patch_size,
scale=face_policy,
transformer=transformer,
)
# Preallocate
score = np.zeros(len(df))
loss = np.zeros(len(df))
loader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False, drop_last=False)
with torch.no_grad():
idx0 = 0
for batch_data in tqdm(loader):
batch_images = batch_data[0].to(device)
batch_labels = batch_data[1].to(device)
batch_samples = len(batch_images)
batch_out = net(batch_images)
batch_loss = criterion(batch_out, batch_labels)
score[idx0:idx0 + batch_samples] = batch_out.cpu().numpy()[:, 0]
loss[idx0:idx0 + batch_samples] = batch_loss.cpu().numpy()[:, 0]
idx0 += batch_samples
out_dict = {'score': score, 'loss': loss}
return out_dict
def select_videos(df: pd.DataFrame, max_videos_per_label: int) -> pd.DataFrame:
"""
Select up to a maximum number of videos
:param df: DataFrame of frames. Required columns: 'video','label'
:param max_videos_per_label: maximum number of real and fake videos
:return: DataFrame of selected frames
"""
# Save random state
st0 = np.random.get_state()
# Set seed for this selection only
np.random.seed(42)
df_fake = df[df.label == True]
fake_videos = df_fake['video'].unique()
selected_fake_videos = np.random.choice(fake_videos, min(max_videos_per_label, len(fake_videos)), replace=False)
df_selected_fake_frames = df_fake[df_fake['video'].isin(selected_fake_videos)]
df_real = df[df.label == False]
real_videos = df_real['video'].unique()
selected_real_videos = np.random.choice(real_videos, min(max_videos_per_label, len(real_videos)), replace=False)
df_selected_real_frames = df_real[df_real['video'].isin(selected_real_videos)]
# Restore random state
np.random.set_state(st0)
return pd.concat((df_selected_fake_frames, df_selected_real_frames), axis=0, verify_integrity=True).copy()
if __name__ == '__main__':
main()