Spaces:
Sleeping
Sleeping
| """ | |
| Testing the performance deep learning models on images/features to predict whether an image is fake/synthetic or real/natural. | |
| """ | |
| # Importing Libraries | |
| import numpy as np | |
| import torch | |
| import os, sys, warnings | |
| warnings.filterwarnings("ignore") | |
| sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) | |
| from yaml import safe_load | |
| import prior_methods.prior_functions.prior_networks as prior_networks | |
| import prior_methods.prior_functions.prior_module as prior_module | |
| import functions.utils as utils | |
| import defaults | |
| # Calling Main function | |
| if __name__ == '__main__': | |
| # ----------------------------------------------------------------- | |
| # Flushing Output | |
| import functools | |
| print = functools.partial(print, flush=True) | |
| # Saving stdout | |
| sys.stdout = open('prior_results/{}.log'.format(os.path.basename(__file__)[:-3]), 'w') | |
| # ----------------------------------------------------------------- | |
| # Parsing Argumens | |
| args = utils.parser_args() | |
| # Iterate | |
| train_test_dataset_types_list = [ | |
| ("DRCT", "UnivFD"), ("DRCT", 'GenImage'), | |
| ("GenImage", "DRCT"), ("GenImage", "UnivFD"), | |
| ("GenImage", "GenImage"), ("DRCT", "DRCT"), | |
| ] | |
| # For each train and test datasets | |
| for train_dataset_type, test_dataset_type in train_test_dataset_types_list: | |
| # Save folder prefix | |
| if train_dataset_type == test_dataset_type: | |
| prefix = "" | |
| else: | |
| prefix = 'Cross_' | |
| # Pre-Process Settings based on dataset | |
| if train_dataset_type == test_dataset_type: | |
| if test_dataset_type == "UnivFD": | |
| preprocess_settings_list = [ | |
| # Default | |
| ({"probability": -1, "gaussian_blur_range": None, "jpeg_compression_qfs": None, "input_image_dimensions": (224,224), "resize": None}, "default"), | |
| ] | |
| else: | |
| preprocess_settings_list = [ | |
| # Default | |
| ({"probability": -1, "gaussian_blur_range": None, "jpeg_compression_qfs": None, "input_image_dimensions": (224,224), "resize": None}, "default"), | |
| # Gaussian-Blur | |
| ({"probability": 1, "gaussian_blur_range": [1,1], "jpeg_compression_qfs": None, "input_image_dimensions": (224,224), "resize": None}, "sigma=1"), | |
| ({"probability": 1, "gaussian_blur_range": [2,2], "jpeg_compression_qfs": None, "input_image_dimensions": (224,224), "resize": None}, "sigma=2"), | |
| ({"probability": 1, "gaussian_blur_range": [3,3], "jpeg_compression_qfs": None, "input_image_dimensions": (224,224), "resize": None}, "sigma=3"), | |
| ({"probability": 1, "gaussian_blur_range": [4,4], "jpeg_compression_qfs": None, "input_image_dimensions": (224,224), "resize": None}, "sigma=4"), | |
| ({"probability": 1, "gaussian_blur_range": [5,5], "jpeg_compression_qfs": None, "input_image_dimensions": (224,224), "resize": None}, "sigma=5"), | |
| # JPEG-Compression | |
| ({"probability": 1, "gaussian_blur_range": None, "jpeg_compression_qfs": [90,90], "input_image_dimensions": (224,224), "resize": None}, "jpegQF=90"), | |
| ({"probability": 1, "gaussian_blur_range": None, "jpeg_compression_qfs": [80,80], "input_image_dimensions": (224,224), "resize": None}, "jpegQF=80"), | |
| ({"probability": 1, "gaussian_blur_range": None, "jpeg_compression_qfs": [70,70], "input_image_dimensions": (224,224), "resize": None}, "jpegQF=70"), | |
| ({"probability": 1, "gaussian_blur_range": None, "jpeg_compression_qfs": [60,60], "input_image_dimensions": (224,224), "resize": None}, "jpegQF=60"), | |
| ({"probability": 1, "gaussian_blur_range": None, "jpeg_compression_qfs": [50,50], "input_image_dimensions": (224,224), "resize": None}, "jpegQF=50"), | |
| ({"probability": 1, "gaussian_blur_range": None, "jpeg_compression_qfs": [40,40], "input_image_dimensions": (224,224), "resize": None}, "jpegQF=40"), | |
| ({"probability": 1, "gaussian_blur_range": None, "jpeg_compression_qfs": [30,30], "input_image_dimensions": (224,224), "resize": None}, "jpegQF=30"), | |
| ] | |
| else: | |
| preprocess_settings_list = [ | |
| # Default | |
| ({"probability": -1, "gaussian_blur_range": None, "jpeg_compression_qfs": None, "input_image_dimensions": (224,224), "resize": None}, "default"), | |
| ] | |
| # For each preprocess_settings | |
| for preprocess_settings, suffix in preprocess_settings_list: | |
| # Inference-Restriction-1: Config Files | |
| """ | |
| - Inference only on limited feature extractors for various kinds of image distortions | |
| """ | |
| if train_dataset_type == "DRCT": | |
| # Config Filenames | |
| config_filenames = [ | |
| "drct-convnext-b", | |
| "drct-clip-vit-l-14", | |
| ] | |
| elif train_dataset_type == "GenImage": | |
| config_filenames = [ | |
| "clip-resnet50", | |
| "clip-vit-l-14", | |
| "drct-convnext-b", | |
| "drct-clip-vit-l-14", | |
| ] | |
| else: | |
| config_filenames = [ | |
| "clip-resnet50", | |
| "clip-vit-l-14", | |
| ] | |
| # Iterating for each config_filename | |
| for config_filename in config_filenames: | |
| # Loading Config file | |
| dir_path = os.path.dirname(os.path.realpath(__file__)) | |
| args.config = os.path.join(dir_path, "prior_configs/{}.yaml".format(config_filename)) | |
| with open(args.config, 'r') as f: | |
| config:dict = safe_load(f) | |
| # Inference-Restriction-2: Variants of Training: Removed | |
| """ | |
| - Inference only on limited feature extractors for basic list of image distortions | |
| """ | |
| if config_filename == "clip-vit-l-14" or config_filename == "clip-resnet50": | |
| checkpoint_directories = ["p=0.5_standard/CrossEntropy"] | |
| else: | |
| checkpoint_directories = ["extensive/MarginContrastiveLoss_CrossEntropy"] | |
| # For each training variant | |
| for ckpt_dir in checkpoint_directories: | |
| # Changes: (resume_ckpt_path, checkpoint_dirname, checkpoint_filename, dataset_type) | |
| config["checkpoints"]["resume_dirname"] = os.path.join(train_dataset_type, ckpt_dir) | |
| config["checkpoints"]["resume_filename"] = "best_model.pth" | |
| config["checkpoints"]["checkpoint_dirname"] = ckpt_dir | |
| config["checkpoints"]["checkpoint_filename"] = "best_model.pth" | |
| config["dataset"]["dataset_type"] = test_dataset_type | |
| # Threshold for calculating metrics | |
| if test_dataset_type == 'UnivFD': | |
| best_threshold = None | |
| else: | |
| best_threshold = 0.5 | |
| # Setting model_name and preprocess_type for Pre-processing | |
| preprocess_settings["model_name"] = config["dataset"]["model_name"] | |
| preprocess_settings["selected_transforms_name"] = "test" | |
| # Dataset-Type | |
| dataset_type = config["dataset"]["dataset_type"] | |
| # Model | |
| model_name = config["dataset"]["model_name"] | |
| f_model_name = config["dataset"]["f_model_name"] | |
| # Model | |
| feature_extractor, classifier = prior_networks.get_model(model_name=config["dataset"]["model_name"], device="cuda") | |
| # Loading Resume-Checkpoints | |
| config["checkpoints"]["resume_ckpt_path"] = os.path.join(defaults.main_prior_checkpoints_dir, config["checkpoints"]["resume_dirname"], f_model_name, config["checkpoints"]["resume_filename"]) | |
| # Loading Weights of Classifier | |
| if model_name.__contains__("drct"): | |
| # Loading entire model i.e CLIP/ConvNext weights as well as additional layer weights. | |
| initial_weights = classifier.state_dict() | |
| trained_weights = torch.load(config["checkpoints"]["resume_ckpt_path"]) | |
| for (k1,v1), (k2,v2) in zip(initial_weights.items(), trained_weights.items()): | |
| assert v1.shape == v2.shape, "Weights Mismatch" | |
| initial_weights[k1] = trained_weights[k2] | |
| classifier.load_state_dict(initial_weights, strict=True) | |
| elif model_name == "clip-vit-l-14" and train_dataset_type == "UnivFD": | |
| # In this case only the weights of classifier are available. | |
| initial_weights = classifier.state_dict() | |
| trained_weights = torch.load(config["checkpoints"]["resume_ckpt_path"]) | |
| for k in trained_weights.keys(): | |
| if k == "weight": | |
| assert trained_weights[k].shape == initial_weights["classifier.weight"].shape, "Weights Mismatch" | |
| initial_weights["classifier.weight"] = trained_weights[k] | |
| elif k == 'bias': | |
| assert trained_weights[k].shape == initial_weights["classifier.bias"].shape, "Weights Mismatch" | |
| initial_weights["classifier.bias"] = trained_weights[k] | |
| else: | |
| None | |
| classifier.load_state_dict(initial_weights, strict=True) | |
| else: | |
| # Cases that we trained for comparisions | |
| initial_weights = classifier.state_dict() | |
| trained_weights = torch.load(config["checkpoints"]["resume_ckpt_path"]) | |
| for k in trained_weights["model"].keys(): | |
| if k == "fc.weight": | |
| assert trained_weights["model"][k].shape == initial_weights["classifier.weight"].shape, "Weights Mismatch" | |
| initial_weights["classifier.weight"] = trained_weights["model"][k] | |
| elif k == 'fc.bias': | |
| assert trained_weights["model"][k].shape == initial_weights["classifier.bias"].shape, "Weights Mismatch" | |
| initial_weights["classifier.bias"] = trained_weights["model"][k] | |
| else: | |
| None | |
| classifier.load_state_dict(initial_weights, strict=True) | |
| # Log | |
| print ( | |
| "\n", | |
| "Classifier:", "\n", | |
| classifier, "\n", | |
| "\n" | |
| ) | |
| # Assertions | |
| for key in ["dataset_type", "model_name"]: | |
| assert key in config["dataset"], "{} not provided".format(key) | |
| # Image-Sources and Classes | |
| if config["dataset"]["dataset_type"] == "GenImage": | |
| # GenImage Dataset | |
| train_image_sources, test_image_sources = utils.get_GenImage_options() | |
| elif config["dataset"]["dataset_type"] == "UnivFD": | |
| # UnivFD Dataset | |
| train_image_sources, test_image_sources = utils.get_UnivFD_options() | |
| elif config["dataset"]["dataset_type"] == "DRCT": | |
| # DRCT Dataset | |
| train_image_sources, test_image_sources = utils.get_DRCT_options() | |
| else: | |
| assert False, "Invalid Dataset" | |
| # Log | |
| print ( | |
| "\n", | |
| "Test-Settings:", "\n", | |
| " "*2, "dataset_type:", dataset_type, "\n", | |
| " "*2, "model_name:", model_name, "\n", | |
| " "*2, "f_model_name:", f_model_name, "\n", | |
| " "*2, "train_image_sources:", train_image_sources, "\n", | |
| " "*2, "test_image_sources:", test_image_sources, "\n", | |
| " "*2, "resume_dirname", config["checkpoints"]["resume_dirname"], "\n", | |
| " "*2, "best_threshold", best_threshold, "\n", | |
| "\n" | |
| ) | |
| # Testing | |
| config["train_settings"]["train"] = False | |
| config["train_loss_fn"]["name"] = "CrossEntropy" | |
| config["val_loss_fn"]["name"] = "CrossEntropy" | |
| test_set_metrics, best_threshold = prior_module.run( | |
| feature_extractor=feature_extractor, | |
| classifier=classifier, | |
| config=config, | |
| train_image_sources=train_image_sources, | |
| test_image_sources=test_image_sources, | |
| preprocess_settings=preprocess_settings, | |
| best_threshold=best_threshold, | |
| verbose=False | |
| ) | |
| # Saving Results | |
| utils.write_results_csv( | |
| test_set_metrics=test_set_metrics, | |
| test_image_sources=test_image_sources, | |
| f_model_name=f_model_name, | |
| save_path=os.path.join("prior_results", "{}{}_{}_{}".format(prefix, ckpt_dir, train_dataset_type, test_dataset_type), "{}.csv".format(suffix)), | |
| ) | |
| print ("\n"*2) |