test_model.py

"""
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()