Spaces:
Sleeping
Sleeping
File size: 6,388 Bytes
a479aa5 | 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 | import os
import numpy as np
from torchvision import datasets, transforms
from utils.toolkit import split_images_labels
import json
import bisect
class iData(object):
train_trsf = []
test_trsf = []
common_trsf = []
class_order = None
class CDDB_benchmark(object):
use_path = True
train_trsf = [
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=63 / 255),
]
test_trsf = [
transforms.Resize(256, interpolation=transforms.InterpolationMode.BICUBIC),
transforms.CenterCrop(224),
]
common_trsf = [
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
def __init__(self, args):
self.args = args
class_order = args["class_order"]
self.class_order = class_order
def download_data(self):
train_dataset = []
test_dataset = []
for id, name in enumerate(self.args["task_name"]):
root_ = os.path.join(self.args["data_path"], name, "train")
sub_classes = os.listdir(root_) if self.args["multiclass"][id] else [""]
for cls in sub_classes:
for imgname in os.listdir(os.path.join(root_, cls, "0_real")):
train_dataset.append(
(os.path.join(root_, cls, "0_real", imgname), 0 + 2 * id)
)
for imgname in os.listdir(os.path.join(root_, cls, "1_fake")):
train_dataset.append(
(os.path.join(root_, cls, "1_fake", imgname), 1 + 2 * id)
)
for id, name in enumerate(self.args["task_name"]):
root_ = os.path.join(self.args["data_path"], name, "val")
sub_classes = os.listdir(root_) if self.args["multiclass"][id] else [""]
for cls in sub_classes:
for imgname in os.listdir(os.path.join(root_, cls, "0_real")):
test_dataset.append(
(os.path.join(root_, cls, "0_real", imgname), 0 + 2 * id)
)
for imgname in os.listdir(os.path.join(root_, cls, "1_fake")):
test_dataset.append(
(os.path.join(root_, cls, "1_fake", imgname), 1 + 2 * id)
)
self.train_data, self.train_targets = split_images_labels(train_dataset)
self.test_data, self.test_targets = split_images_labels(test_dataset)
def parse_dataset(data_keys):
gen_keys = {
'gan1':['StyleGAN'],
'gan2':['StyleGAN2'],
'gan3':['StyleGAN3'],
'sd15':['StableDiffusion1.5'],
'sd2':['StableDiffusion2'],
'sd3':['StableDiffusion3'],
'sdXL':['StableDiffusionXL'],
'flux':['FLUX.1'],
'realFFHQ':['FFHQ'],
'realFORLAB':['FORLAB']
}
gen_keys['all'] = [gen_keys[key][0] for key in gen_keys.keys()]
# gen_keys['gan'] = [gen_keys[key][0] for key in gen_keys.keys() if 'gan' in key]
# gen_keys['sd'] = [gen_keys[key][0] for key in gen_keys.keys() if 'sd' in key]
gen_keys['real'] = [gen_keys[key][0] for key in gen_keys.keys() if 'real' in key]
mod_keys = {
'pre': ['PreSocial'],
'fb': ['Facebook'],
'tl': ['Telegram'],
'tw': ['X'],
}
mod_keys['all'] = [mod_keys[key][0] for key in mod_keys.keys()]
mod_keys['shr'] = [mod_keys[key][0] for key in mod_keys.keys() if key in ['fb', 'tl', 'tw']]
dataset_list = []
for data in data_keys.split('&'):
gen, mod = data.split(':')
dataset_list.append({'gen':gen_keys[gen], 'mod':mod_keys[mod]})
return dataset_list
class TrueFake_benchmark(object):
use_path = True
train_trsf = [
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=63 / 255),
]
test_trsf = [
transforms.Resize(256, interpolation=transforms.InterpolationMode.BICUBIC),
transforms.CenterCrop(224),
]
common_trsf = [
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
def __init__(self, args):
self.args = args
class_order = args["class_order"]
self.class_order = class_order
def _in_list(self, split, elem):
i = bisect.bisect_left(split, elem)
return i != len(split) and split[i] == elem
def download_data(self):
with open(self.args["split_file"], "r") as f:
splits = json.load(f)
train_split = sorted(splits["train"])
val_split = sorted(splits["val"])
train_dataset = []
test_dataset = []
for id, name in enumerate(self.args["task_name"]):
dataset_list = parse_dataset(name)
for dict in dataset_list:
generators = dict['gen']
modifiers = dict['mod']
for mod in modifiers:
for dataset_root, dataset_dirs, dataset_files in os.walk(os.path.join(self.args["data_path"], mod), topdown=True, followlinks=True):
if len(dataset_dirs):
continue
(label, gen, sub) = f'{dataset_root}/'.replace(os.path.join(self.args["data_path"], mod) + os.sep, '').split(os.sep)[:3]
if gen in generators:
for filename in sorted(dataset_files):
if os.path.splitext(filename)[1].lower() in ['.png', '.jpg', '.jpeg']:
if self._in_list(train_split, os.path.join(gen, sub, os.path.splitext(filename)[0])):
train_dataset.append((os.path.join(dataset_root, filename), (1 if label == 'Fake' else 0) + 2 * id))
if self._in_list(val_split, os.path.join(gen, sub, os.path.splitext(filename)[0])):
test_dataset.append((os.path.join(dataset_root, filename), (1 if label == 'Fake' else 0) + 2 * id))
self.train_data, self.train_targets = split_images_labels(train_dataset)
self.test_data, self.test_targets = split_images_labels(test_dataset)
|