# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. import io, os, pdb import cv2, math, random import numpy as np import torch from PIL import Image, ImageFile from torch.utils.data import Dataset from torchvision import transforms from torchvision.transforms import functional as F from torchvision.transforms import InterpolationMode ImageFile.LOAD_TRUNCATED_IMAGES = True class RandomJPEG(): def __init__(self, quality=95, interval=1, p=0.1): if isinstance(quality, tuple): self.quality = [i for i in range(quality[0], quality[1]) if i % interval == 0] else: self.quality = quality self.p = p def __call__(self, img): if random.random() < self.p: if isinstance(self.quality, list): quality = random.choice(self.quality) else: quality = self.quality buffer = io.BytesIO() img.save(buffer, format='JPEG', quality=quality) buffer.seek(0) img = Image.open(buffer) return img class RandomGaussianBlur(): def __init__(self, kernel_size, sigma=(0.1, 2.0), p=1.0): self.blur = transforms.GaussianBlur(kernel_size=kernel_size, sigma=sigma) self.p = p def __call__(self, img): if random.random() < self.p: return self.blur(img) return img class RandomMask(object): def __init__(self, ratio=0.5, patch_size=16, p=0.5): """ Args: ratio (float or tuple of float): If float, the ratio of the image to be masked. If tuple of float, random sample ratio between the two values. patch_size (int): the size of the mask (d*d). """ if isinstance(ratio, float): self.fixed_ratio = True self.ratio = (ratio, ratio) elif isinstance(ratio, tuple) and len(ratio) == 2 and all(isinstance(r, float) for r in ratio): self.fixed_ratio = False self.ratio = ratio else: raise ValueError("Ratio must be a float or a tuple of two floats.") self.patch_size = patch_size self.p = p def __call__(self, tensor): if random.random() > self.p: return tensor _, h, w = tensor.shape mask = torch.ones((h, w), dtype=torch.float32) if self.fixed_ratio: ratio = self.ratio[0] else: ratio = random.uniform(self.ratio[0], self.ratio[1]) # Calculate the number of masks needed num_masks = int((h * w * ratio) / (self.patch_size ** 2)) # Generate non-overlapping random positions selected_positions = set() while len(selected_positions) < num_masks: top = random.randint(0, (h // self.patch_size) - 1) * self.patch_size left = random.randint(0, (w // self.patch_size) - 1) * self.patch_size selected_positions.add((top, left)) for (top, left) in selected_positions: mask[top:top+self.patch_size, left:left+self.patch_size] = 0 return tensor * mask.expand_as(tensor) def Get_Transforms(args): size = args.input_size TRANSFORM_DICT = { 'resize_BILINEAR': { 'train': [ transforms.RandomResizedCrop([size, size], interpolation=InterpolationMode.BILINEAR), ], 'eval': [ transforms.Resize([size, size], interpolation=InterpolationMode.BILINEAR), ], }, 'resize_NEAREST': { 'train': [ transforms.RandomResizedCrop([size, size], interpolation=InterpolationMode.NEAREST), ], 'eval': [ transforms.Resize([size, size], interpolation=InterpolationMode.NEAREST), ], }, 'crop': { 'train': [ transforms.RandomCrop([size, size], pad_if_needed=True), ], 'eval': [ transforms.CenterCrop([size, size]), ], }, 'source': { 'train': [ transforms.RandomCrop([size, size], pad_if_needed=True), ], 'eval': [ ], }, } # region [Augmentations] transform_train, transform_eval = TRANSFORM_DICT[args.transform_mode]['train'], TRANSFORM_DICT[args.transform_mode]['eval'] transform_train.extend([ transforms.RandomHorizontalFlip(p=0.5), # transforms.RandomHorizontalFlip(), transforms.RandomRotation(180), transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5), transforms.ToTensor(), RandomMask(ratio=(0.00, 0.75), patch_size=16, p=0.5), ]) transform_eval.append(transforms.ToTensor()) # endregion # region [Perturbatiocns in Testing] if args.jpeg_factor is not None: transform_eval.insert(0, RandomJPEG(quality=args.jpeg_factor, p=1.0)) if args.blur_sigma is not None: transform_eval.insert(0, transforms.GaussianBlur(kernel_size=5, sigma=args.blur_sigma)) if args.mask_ratio is not None and args.mask_patch_size is not None: transform_eval.append(RandomMask(ratio=args.mask_ratio, patch_size=args.mask_patch_size, p=1.0)) # endregion return transforms.Compose(transform_train), transforms.Compose(transform_eval) class TrainDataset(Dataset): def __init__(self, is_train, args): TRANSFORM = Get_Transforms(args) self.transform = TRANSFORM[0] if is_train else TRANSFORM[1] root = args.data_path if is_train else args.eval_data_path dataset_list = root.replace(' ', '').split(',') num_datasets = len(dataset_list) if num_datasets == 1: real_list, fake_list = self.get_real_and_fake_lists(dataset_list[0]) if is_train and args.num_train is not None: self.data_list = real_list[:args.num_train//2] + fake_list[:args.num_train//2] else: self.data_list = real_list + fake_list else: assert args.num_train is not None self.data_list = [] for dataset in dataset_list: real_list, fake_list = self.get_real_and_fake_lists(dataset) self.data_list.extend(real_list[:args.num_train//(2 * num_datasets)] + fake_list[:args.num_train//(2 * num_datasets)]) def get_image_paths(self, dir_path): image_extensions = ('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp') image_paths = [] for root, dirs, files in sorted(os.walk(dir_path)): for file in sorted(files): if file.lower().endswith(image_extensions): image_paths.append(os.path.join(root, file)) return image_paths def get_real_and_fake_lists(self, folder_path): real_list, fake_list = [], [] for root, dirs, files in sorted(os.walk(folder_path, followlinks=True)): for dir_name in sorted(dirs): if dir_name == "0_real": real_dir_path = os.path.join(root, dir_name) real_list.extend([{"image_path": image_path, "label" : 0} for image_path in self.get_image_paths(real_dir_path)]) elif dir_name == "1_fake": fake_dir_path = os.path.join(root, dir_name) fake_list.extend([{"image_path": image_path, "label" : 1} for image_path in self.get_image_paths(fake_dir_path)]) return real_list, fake_list def __len__(self): return len(self.data_list) def __getitem__(self, index): sample = self.data_list[index] image_path, targets = sample['image_path'], sample['label'] image = Image.open(image_path).convert('RGB') image = self.transform(image) return image, torch.tensor(int(targets))