|
|
import cv2 |
|
|
import numpy as np |
|
|
import torchvision.datasets as datasets |
|
|
import torchvision.transforms as transforms |
|
|
import torchvision.transforms.functional as TF |
|
|
from torch.utils.data import Dataset |
|
|
from random import random, choice, shuffle |
|
|
from io import BytesIO |
|
|
from PIL import Image |
|
|
from PIL import ImageFile |
|
|
from scipy.ndimage.filters import gaussian_filter |
|
|
import pickle |
|
|
import os |
|
|
import math |
|
|
from skimage.io import imread |
|
|
from copy import deepcopy |
|
|
import torch |
|
|
|
|
|
ImageFile.LOAD_TRUNCATED_IMAGES = True |
|
|
|
|
|
|
|
|
MEAN = { |
|
|
"imagenet":[0.485, 0.456, 0.406], |
|
|
"clip":[0.48145466, 0.4578275, 0.40821073], |
|
|
"beitv2": [0.485, 0.456, 0.406], |
|
|
"siglip": [0.5, 0.5, 0.5], |
|
|
} |
|
|
|
|
|
STD = { |
|
|
"imagenet":[0.229, 0.224, 0.225], |
|
|
"clip":[0.26862954, 0.26130258, 0.27577711], |
|
|
"beitv2": [0.229, 0.224, 0.225], |
|
|
"siglip": [0.5, 0.5, 0.5], |
|
|
} |
|
|
|
|
|
|
|
|
def translate_duplicate(img, cropSize): |
|
|
if min(img.size) < cropSize: |
|
|
width, height = img.size |
|
|
|
|
|
new_width = width * math.ceil(cropSize/width) |
|
|
new_height = height * math.ceil(cropSize/height) |
|
|
|
|
|
new_img = Image.new('RGB', (new_width, new_height)) |
|
|
for i in range(0, new_width, width): |
|
|
for j in range(0, new_height, height): |
|
|
new_img.paste(img, (i, j)) |
|
|
return new_img |
|
|
else: |
|
|
return img |
|
|
|
|
|
|
|
|
def recursively_read(rootdir, must_contain, classes=[], exts=["png", "jpg", "JPEG", "jpeg"]): |
|
|
out = [] |
|
|
for r, d, f in os.walk(rootdir): |
|
|
for file in f: |
|
|
if (file.split('.')[1] in exts) and (must_contain in os.path.join(r, file)): |
|
|
if len(classes) == 0: |
|
|
out.append(os.path.join(r, file)) |
|
|
elif os.path.join(r, file).split('/')[-3] in classes: |
|
|
out.append(os.path.join(r, file)) |
|
|
return out |
|
|
|
|
|
|
|
|
def get_list(path, must_contain='', classes=[]): |
|
|
if ".pickle" in path: |
|
|
with open(path, 'rb') as f: |
|
|
image_list = pickle.load(f) |
|
|
image_list = [ item for item in image_list if must_contain in item ] |
|
|
else: |
|
|
image_list = recursively_read(path, must_contain, classes) |
|
|
return image_list |
|
|
|
|
|
|
|
|
class RealFakeDataset(Dataset): |
|
|
def __init__(self, opt): |
|
|
assert opt.data_label in ["train", "val"] |
|
|
|
|
|
self.data_label = opt.data_label |
|
|
if opt.data_mode == 'ours': |
|
|
pickle_name = "train.pickle" if opt.data_label=="train" else "val.pickle" |
|
|
real_list = get_list( os.path.join(opt.real_list_path, pickle_name) ) |
|
|
fake_list = get_list( os.path.join(opt.fake_list_path, pickle_name) ) |
|
|
elif opt.data_mode == 'wang2020': |
|
|
temp = 'train' if opt.data_label == 'train' else 'test' |
|
|
if opt.data_label == 'train': |
|
|
|
|
|
real_list = get_list( os.path.join(opt.wang2020_data_path,temp), must_contain='real' ) |
|
|
fake_list = get_list( os.path.join(opt.wang2020_data_path,temp), must_contain='fake' ) |
|
|
else: |
|
|
|
|
|
real_list = get_list( os.path.join(opt.wang2020_data_path,temp), must_contain='real' ) |
|
|
fake_list = get_list( os.path.join(opt.wang2020_data_path,temp), must_contain='fake' ) |
|
|
elif opt.data_mode == 'ours_wang2020': |
|
|
pickle_name = "train.pickle" if opt.data_label=="train" else "val.pickle" |
|
|
real_list = get_list( os.path.join(opt.real_list_path, pickle_name) ) |
|
|
fake_list = get_list( os.path.join(opt.fake_list_path, pickle_name) ) |
|
|
|
|
|
temp = 'train' if opt.data_label == 'train' else 'test' |
|
|
real_list += get_list( os.path.join(opt.wang2020_data_path,temp), must_contain='real' ) |
|
|
fake_list += get_list( os.path.join(opt.wang2020_data_path,temp), must_contain='fake' ) |
|
|
|
|
|
|
|
|
self.labels_dict = {} |
|
|
for i in real_list: |
|
|
self.labels_dict[i] = 0 |
|
|
for i in fake_list: |
|
|
self.labels_dict[i] = 1 |
|
|
|
|
|
self.total_list = real_list + fake_list |
|
|
shuffle(self.total_list) |
|
|
if opt.isTrain: |
|
|
crop_func = transforms.RandomCrop(opt.cropSize) |
|
|
elif opt.no_crop: |
|
|
crop_func = transforms.Lambda(lambda img: img) |
|
|
else: |
|
|
crop_func = transforms.CenterCrop(opt.cropSize) |
|
|
|
|
|
if opt.isTrain and not opt.no_flip: |
|
|
flip_func = transforms.RandomHorizontalFlip() |
|
|
else: |
|
|
flip_func = transforms.Lambda(lambda img: img) |
|
|
if not opt.isTrain and opt.no_resize: |
|
|
rz_func = transforms.Lambda(lambda img: img) |
|
|
else: |
|
|
rz_func = transforms.Lambda(lambda img: custom_resize(img, opt)) |
|
|
|
|
|
if opt.arch.lower().startswith("imagenet"): |
|
|
stat_from = "imagenet" |
|
|
elif opt.arch.lower().startswith("clip"): |
|
|
stat_from = "clip" |
|
|
elif opt.arch.lower().startswith("siglip"): |
|
|
stat_from = "siglip" |
|
|
elif opt.arch.lower().startswith("beitv2"): |
|
|
stat_from = "beitv2" |
|
|
|
|
|
print("mean and std stats are from: ", stat_from) |
|
|
if '2b' not in opt.arch: |
|
|
print ("using Official CLIP's normalization") |
|
|
self.transform = transforms.Compose([ |
|
|
|
|
|
transforms.Lambda(lambda img: translate_duplicate(img, opt.loadSize)), |
|
|
crop_func, |
|
|
flip_func, |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize( mean=MEAN[stat_from], std=STD[stat_from] ), |
|
|
]) |
|
|
else: |
|
|
print ("Using CLIP 2B transform") |
|
|
self.transform = None |
|
|
|
|
|
def __len__(self): |
|
|
return len(self.total_list) |
|
|
|
|
|
def __getitem__(self, idx): |
|
|
img_path = self.total_list[idx] |
|
|
label = self.labels_dict[img_path] |
|
|
img = Image.open(img_path).convert("RGB") |
|
|
img = self.transform(img) |
|
|
return img, label |
|
|
|
|
|
|
|
|
def data_augment(img, opt): |
|
|
img = np.array(img) |
|
|
if img.ndim == 2: |
|
|
img = np.expand_dims(img, axis=2) |
|
|
img = np.repeat(img, 3, axis=2) |
|
|
|
|
|
if random() < opt.blur_prob: |
|
|
sig = sample_continuous(opt.blur_sig) |
|
|
gaussian_blur(img, sig) |
|
|
|
|
|
if random() < opt.jpg_prob: |
|
|
method = sample_discrete(opt.jpg_method) |
|
|
qual = sample_discrete(opt.jpg_qual) |
|
|
img = jpeg_from_key(img, qual, method) |
|
|
|
|
|
return Image.fromarray(img) |
|
|
|
|
|
|
|
|
def sample_continuous(s): |
|
|
if len(s) == 1: |
|
|
return s[0] |
|
|
if len(s) == 2: |
|
|
rg = s[1] - s[0] |
|
|
return random() * rg + s[0] |
|
|
raise ValueError("Length of iterable s should be 1 or 2.") |
|
|
|
|
|
|
|
|
def sample_discrete(s): |
|
|
if len(s) == 1: |
|
|
return s[0] |
|
|
return choice(s) |
|
|
|
|
|
|
|
|
def gaussian_blur(img, sigma): |
|
|
gaussian_filter(img[:,:,0], output=img[:,:,0], sigma=sigma) |
|
|
gaussian_filter(img[:,:,1], output=img[:,:,1], sigma=sigma) |
|
|
gaussian_filter(img[:,:,2], output=img[:,:,2], sigma=sigma) |
|
|
|
|
|
|
|
|
def cv2_jpg(img, compress_val): |
|
|
img_cv2 = img[:,:,::-1] |
|
|
encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), compress_val] |
|
|
result, encimg = cv2.imencode('.jpg', img_cv2, encode_param) |
|
|
decimg = cv2.imdecode(encimg, 1) |
|
|
return decimg[:,:,::-1] |
|
|
|
|
|
|
|
|
def pil_jpg(img, compress_val): |
|
|
out = BytesIO() |
|
|
img = Image.fromarray(img) |
|
|
img.save(out, format='jpeg', quality=compress_val) |
|
|
img = Image.open(out) |
|
|
|
|
|
img = np.array(img) |
|
|
out.close() |
|
|
return img |
|
|
|
|
|
|
|
|
jpeg_dict = {'cv2': cv2_jpg, 'pil': pil_jpg} |
|
|
def jpeg_from_key(img, compress_val, key): |
|
|
method = jpeg_dict[key] |
|
|
return method(img, compress_val) |
|
|
|
|
|
|
|
|
rz_dict = {'bilinear': Image.BILINEAR, |
|
|
'bicubic': Image.BICUBIC, |
|
|
'lanczos': Image.LANCZOS, |
|
|
'nearest': Image.NEAREST} |
|
|
def custom_resize(img, opt): |
|
|
interp = sample_discrete(opt.rz_interp) |
|
|
return TF.resize(img, opt.loadSize, interpolation=rz_dict[interp]) |
|
|
|
