test.py

import sys
import time
import os
import csv
import torch
from util import Logger, printSet
from validate import validate
from thop import profile
from torchsummary import summary
from networks.LFMSN import LFMSN
from options.test_options import TestOptions
import numpy as np
import random
import random


def seed_torch(seed=1029):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.enabled = False


seed_torch(100)
DetectionTests = {
    'Diffusion1kStep': {'dataroot': './datasets/Generalization_Test/Diffusion1kStep/',
                        'no_resize': False,
                        # Due to the different shapes of images in the dataset, resizing is required during batch detection.
                        'no_crop': True,
                        },

    'ForenSynths': {'dataroot': './datasets/Generalization_Test/ForenSynths_test/',
                    'no_resize': False,
                    # Due to the different shapes of images in the dataset, resizing is required during batch detection.
                    'no_crop': True,
                    },

    'GANGen-Detection': {'dataroot': './datasets/Generalization_Test/GANGen-Detection/',
                         'no_resize': True,
                         'no_crop': True,
                         },

    'DiffusionForensics': {'dataroot': './datasets/Generalization_Test/DiffusionForensics/',
                           'no_resize': False,
                           # Due to the different shapes of images in the dataset, resizing is required during batch detection.
                           'no_crop': True,
                           },

    'UniversalFakeDetect': {'dataroot': './datasets/Generalization_Test/UniversalFakeDetect/',
                            'no_resize': False,
                            # Due to the different shapes of images in the dataset, resizing is required during batch detection.
                            'no_crop': True,
                            },

}

opt = TestOptions().parse(print_options=False)
print(f'Model_path {opt.model_path}')

# get model
model = LFMSN()
model.load_state_dict(torch.load(opt.model_path, map_location='cpu'), strict=True)
model.cuda()
model.eval()
#summary(model.cuda(), (3, 256, 256))
#input = torch.randn(1, 3, 256, 256).cuda()
#flops, params = profile(model, inputs=(input,))
#print(f'FLOPs: {flops / 1e9} GFLOPs, Params: {params / 1e6} M')
tatol_accs = []
tatol_aps = []
for testSet in DetectionTests.keys():
    dataroot = DetectionTests[testSet]['dataroot']
    printSet(testSet)

    accs = []
    aps = []
    print(time.strftime("%Y_%m_%d_%H_%M_%S", time.localtime()))
    for v_id, val in enumerate(os.listdir(dataroot)):
        opt.dataroot = '{}/{}'.format(dataroot, val)
        opt.classes = ''  # os.listdir(opt.dataroot) if multiclass[v_id] else ['']
        opt.no_resize = DetectionTests[testSet]['no_resize']
        opt.no_crop = DetectionTests[testSet]['no_crop']
        acc, ap, _, _, _, _ = validate(model, opt)
        accs.append(acc)
        aps.append(ap)
        tatol_accs.append(acc)
        tatol_aps.append(ap)
        print("({} {:12}) acc: {:.1f}; ap: {:.1f}".format(v_id, val, acc * 100, ap * 100))
    print("({} {:10}) acc: {:.1f}; ap: {:.1f}".format(v_id + 1, 'Mean', np.array(accs).mean() * 100,
                                                      np.array(aps).mean() * 100))
    print('*' * 25)

print("(tatol acc: {:.1f}; ap: {:.1f}".format(np.array(tatol_accs).mean() * 100, np.array(tatol_aps).mean() * 100))
print('*' * 25)