clean/video/mintime/preprocessing/split_dataset.py

# ForgeryNet provided a training set and a validation set but not a complete test set. With this code, the validation set is moved into a folder so that it can be used as a test set,
# while a new validation set is derived from the training set. 
# The latter is constructed so that it has a distribution of deepfake generation methods equal to that of the training set and is composed of a number of samples equal to 10% 
# of those in the training set.
# A plot is also generated to show the distribution of the three datasets.

import os
import argparse
import pandas as pd
import math
import matplotlib.pyplot as plt
import collections
import random
import shutil
import glob
import csv

seed = 42

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--train_list_file', default="../../datasets/ForgeryNet/Training/video_list_complete.txt", type=str,
                        help='Videos List txt file path for training set (to be splitted in train and validation)')
    parser.add_argument('--validation_list_file', default="../../datasets/ForgeryNet/Validation/video_list.txt", type=str,
                        help='Videos List txt file path for validation set (our test set)')
    parser.add_argument('--plots_output_path', default="../outputs", type=str,
                        help='Plots output path')
    parser.add_argument('--faces_path', default="../../datasets/ForgeryNet/faces", type=str,
                        help='Images path')
    parser.add_argument('--validation_set_output_path', default="../../datasets/ForgeryNet/faces/val", type=str,
                        help='Test set output path')
    parser.add_argument('--train_faces_path', default="../../datasets/ForgeryNet/faces/train", type=str,
                        help='Train images path')
    parser.add_argument('--test_faces_path', default="../../datasets/ForgeryNet/faces/test", type=str,
                        help='Test images path')

    opt = parser.parse_args()
    print(opt)
    datasets = {"train": {}, "val": {}, "test": {}}
    
    # Reading of the training set and extraction of its distribution excluding videos in which no faces were found.
    paths = glob.glob(f'{opt.train_faces_path}/*/**/*.mp4', recursive=True)
    with open(opt.train_list_file, 'r') as temp_f:
        col_count = [ len(l.split(" ")) for l in temp_f.readlines() ]

        column_names = [i for i in range(0, max(col_count))]
        df = pd.read_csv(opt.train_list_file, sep=' ', names=column_names)
      
        training_counter = {}        
        column_names.reverse()
        skipped = 0
        for index, row in df.iterrows():
            video_name = os.path.join(opt.train_faces_path, row[1].split("train_video_release/")[-1])
            if video_name not in paths:
                skipped += 1
                continue

            for column_name in column_names:
                if not math.isnan(row[column_name]):
                    deepfake_class = row[column_name]
                    break
            
            if deepfake_class in training_counter:
                training_counter[deepfake_class] += 1
            else:
                training_counter[deepfake_class] = 1

            if deepfake_class in datasets["train"]:
                datasets["train"][deepfake_class].append(video_name.replace("train_video_release", "train").replace(opt.train_faces_path, "train"))
            else:
                datasets["train"][deepfake_class] = [video_name.replace("train_video_release", "train").replace(opt.train_faces_path, "train")]

        print(skipped, "videos in training set without detected faces skipped.")
        training_counter = collections.OrderedDict(sorted(training_counter.items()))

        # Construction of the validation set from the training set distribution
        total_training_samples = len(df)
        validation_size = total_training_samples/10
        total = 0
        validation_counter = {}
        for key in training_counter:
            percentage = training_counter[key]/total_training_samples
            elements = validation_size*percentage
            validation_counter[key] = int(elements)
            training_counter[key] -= elements

        validation_counter = collections.OrderedDict(sorted(validation_counter.items()))

        # Plotting training set distribution
        names = list(training_counter.keys())
        values = list(training_counter.values())
        x = [i-0.3 for i in range(len(training_counter))]
        plt.bar(x, values, 0.3, tick_label=names, label = "Training Set")

        # Plotting validation set distribution        
        names = list(validation_counter.keys())
        values = list(validation_counter.values())
        x = [i for i in range(len(training_counter))]
        plt.bar(x, values, 0.3, tick_label=names, label = "Validation Set")
    
    # Reading of the validation set (which will be used as a test set) and extraction of its distribution excluding videos in which no faces were found.
    skipped = 0
    with open(opt.validation_list_file, 'r') as temp_f:
        col_count = [ len(l.split(" ")) for l in temp_f.readlines() ]

        column_names = [i for i in range(0, max(col_count))]
        df = pd.read_csv(opt.validation_list_file, sep=' ', names=column_names)
      
        test_counter = {}        
        column_names.reverse()

        paths = glob.glob(f'{opt.test_faces_path}/*/**/*.mp4', recursive=True)
        for index, row in df.iterrows():
            video_name = os.path.join(opt.test_faces_path, row[1].split("val_video_release/")[-1])
            if video_name not in paths:
                skipped += 1
                continue

            for column_name in column_names:
                if not math.isnan(row[column_name]):
                    deepfake_class = row[column_name]
                    break
            
            if deepfake_class in test_counter:
                test_counter[deepfake_class] += 1
            else:
                test_counter[deepfake_class] = 1

            if deepfake_class in datasets["test"]:
                datasets["test"][deepfake_class].append(video_name.replace("val_video_release", "test").replace(opt.test_faces_path, "test"))
            else:
                datasets["test"][deepfake_class] = [video_name.replace("val_video_release", "test").replace(opt.test_faces_path, "test")]
    
        print(skipped, "videos in test set without detected faces skipped.")
    test_counter = collections.OrderedDict(sorted(test_counter.items()))

    # Plotting test set distribution
    names = list(test_counter.keys())
    values = list(test_counter.values())
    
    x = [i+0.3 for i in range(len(test_counter))]
    plt.bar(x, values, 0.3, tick_label=names, label = "Test Set")

    plt.legend()
    plt.savefig(os.path.join(opt.plots_output_path, "distribution"))


# Move selected training files for the validation set construction into validation folder
for deepfake_class in datasets["train"]:
    number_of_elements = validation_counter[deepfake_class]
    extracted_elements = random.Random(seed).sample(datasets["train"][deepfake_class],number_of_elements)
    for index, video_name in enumerate(extracted_elements):
        out_path = os.path.join(opt.validation_set_output_path, video_name.split("Training/video")[-1]).replace("val/train", "val")
        src_path = os.path.join(opt.faces_path, video_name).replace("train_video_release", "train")
        datasets["train"][deepfake_class].remove(video_name)
        if deepfake_class in datasets["val"]:
            datasets["val"][deepfake_class].append(video_name.replace("train", "val"))
        else:
            datasets["val"][deepfake_class] = [video_name.replace("train", "val")]
        if index % 500 == 0:
            print("Moved", index, "videos into validation set.")
        shutil.move(src_path, out_path)
    
# Generate labels csv files for the three sets
for key in datasets:
    f = open(os.path.join(opt.faces_path, key+".csv"), 'w+')
    dataset = datasets[key]
    for deepfake_class in dataset:
        if deepfake_class == 0:
            binary_class = "0"
        else:
            binary_class = "1"
        for video in dataset[deepfake_class]:
            row = video + " " + binary_class + " " + str(int(deepfake_class)) + "\n"
            f.write(row)
            
    f.close()