def image_to_patches(img):
img = F.resize(img, size=292)
img = RandomCrop((255, 255))(img)
split_per_side = 3 # split of patches per image side
patch_jitter = 21 # jitter of each patch from each grid
h, w = img.size
h_grid = h // split_per_side
w_grid = w // split_per_side
h_patch = h_grid - patch_jitter
w_patch = w_grid - patch_jitter
assert h_patch > 0 and w_patch > 0
patches = []
for i in range(split_per_side):
for j in range(split_per_side):
p = F.crop(img, i * h_grid, j * w_grid, h_grid, w_grid)
p = RandomCrop((h_patch, w_patch))(p)
patches.append(p)
patches = [patch_opt(p) for p in patches]
perms = []
[
perms.append(torch.cat((patches[i], patches[4]), dim=0))
for i in range(9) if i != 4
]
#patch_labels = torch.LongTensor([0, 1, 2, 3, 4, 5, 6, 7])
patches = torch.stack(perms)
return patches
def __init__(self):
self.train_transform = Compose(
[RandomCrop(size=28, padding=2),
ToTensor()])
self.test_transform = Compose(
[RandomCrop(size=28, padding=2),
ToTensor()])
def __init__(self, data_path):
super(InputPair, self).__init__()
self.config = self.setup_config()
self.data_path = data_path
imdb = self._load_pickle()
self.videos_num = imdb['n_videos']
self.videos = imdb['videos']
self.indices = np.random.permutation(self.videos_num)
# https://github.com/bilylee/SiamFC-TensorFlow/issues/69
safe_size = self.config['instance_sz'] - \
int(self.config['instance_sz'] * self.config['max_stretch_scale']) # 243
perturbation_size = safe_size - 8 # size for center perturbation, 235
self.transform_z = Compose([
RandomStretch(max_stretch=self.config['max_stretch_scale']),
# after RandomStretch, the size may smaller than 255, make sure CenterCrop area in it.
CenterCrop(safe_size),
RandomCrop(perturbation_size),
CenterCrop(self.config['exemplar_sz']),
ToTensor()
])
# input x do not have to be 255*255*3
self.transform_x = Compose([
RandomStretch(max_stretch=self.config['max_stretch_scale']),
CenterCrop(safe_size),
RandomCrop(perturbation_size),
ToTensor()
])
def get_transforms(auto_augment, input_sizes, m, mean, n, std):
if auto_augment:
# AutoAugment + Cutout
train_transforms = Compose([
RandomCrop(size=input_sizes, padding=4, fill=128),
RandomHorizontalFlip(p=0.5),
CIFAR10Policy(),
ToTensor(),
Normalize(mean=mean, std=std),
Cutout(n_holes=1, length=16)
])
else:
# RandAugment + Cutout
train_transforms = Compose([
RandomCrop(size=input_sizes, padding=4, fill=128),
RandomHorizontalFlip(p=0.5),
RandomRandAugment(n=n, m_max=m), # This version includes cutout
ToTensor(),
Normalize(mean=mean, std=std)
])
test_transforms = Compose([
ToTensor(),
Normalize(mean=mean, std=std)
])
return test_transforms, train_transforms
def __init__(self, seq_dataset, **kargs):
super(Pairwise, self).__init__()
self.cfg = self.parse_args(**kargs)
self.pairs_per_seq = 10
self.max_dist = 100
self.exemplar_sz = 127
self.instance_sz = 255
self.context = 0.5
self.seq_dataset = seq_dataset
self.indices = np.random.permutation(len(seq_dataset))
# augmentation for exemplar and instance images
self.transform_z = Compose([
RandomStretch(max_stretch=0.05),
CenterCrop(self.instance_sz - 8),
RandomCrop(self.instance_sz - 2 * 8),
CenterCrop(self.exemplar_sz),
ToTensor()
])
self.transform_x = Compose([
RandomStretch(max_stretch=0.05),
CenterCrop(self.instance_sz - 8),
RandomCrop(self.instance_sz - 2 * 8),
ToTensor()
])
def prepare_data(self):
"Prepare supervised and unsupervised datasets from cifar"
dataset_path = self.hparams.dataset_path
n_labeled = self.hparams.n_labeled
n_overlap = self.hparams.n_overlap
seed = self.hparams.seed
if self.hparams.dataset == "cifar":
n = self.hparams.randaug_n
m = self.hparams.randaug_m
uda_tfm = Compose([RandAugment(n, m), ToTensor(), Normalize(*CIFAR_STATS)])
sup_tfm = Compose([RandomCrop(32, 4, padding_mode="reflect"), RandomHorizontalFlip(), ToTensor(), Normalize(*CIFAR_STATS)])
val_tfm = Compose([ToTensor(), Normalize(*CIFAR_STATS)])
sup_ds, unsup_ds = Cifar.uda_ds(dataset_path, n_labeled, n_overlap, sup_tfm, uda_tfm, seed=seed)
val_ds = Cifar.val_ds(dataset_path, val_tfm)
if self.hparams.dataset == "quickdraw":
uda_tfm = Compose([ExpandChannels, SketchDeformation, RandomHorizontalFlip(), RandomRotation(30), RandomCrop(128, 18), ToTensor()])
sup_tfm = Compose([ExpandChannels, RandomCrop(128, 9), RandomHorizontalFlip(), ToTensor()])
val_tfm = Compose([ExpandChannels, ToTensor()])
sup_ds, unsup_ds = QuickDraw.uda_ds(dataset_path, n_labeled, n_overlap, sup_tfm, uda_tfm, seed=seed)
val_ds = QuickDraw.val_ds(dataset_path, val_tfm)
self.train_ds_sup = sup_ds
self.train_ds_unsup = unsup_ds
self.valid_ds = val_ds
def __init__(self,
num_views,
random_seed,
dataset,
additional_face=True,
jittering=False):
if dataset == 1:
self.ids = np.load('../Datasets/voxceleb1_ori/train.npy')
if dataset == 2:
self.ids = np.load('../Datasets/voxceleb1_ori/val.npy')
if dataset == 3:
self.ids = np.load('../Datasets/voxceleb1_ori/test.npy')
self.rng = np.random.RandomState(random_seed)
self.num_views = num_views
#self.base_file = os.environ['VOX_CELEB_LOCATION'] + '/%s/'
self.base_file = VOX_CELEB_LOCATION + '/%s/'
crop = 200
if jittering == True:
precrop = crop + 24
crop = self.rng.randint(crop, precrop)
self.pose_transform = Compose([
Scale((256, 256)),
Pad((20, 80, 20, 30)),
CenterCrop(precrop),
RandomCrop(crop),
Scale((256, 256)),
ToTensor()
])
self.transform = Compose([
Scale((256, 256)),
Pad((24, 24, 24, 24)),
CenterCrop(precrop),
RandomCrop(crop),
Scale((256, 256)),
ToTensor()
])
else:
precrop = crop + 24
self.pose_transform = Compose([
Scale((256, 256)),
Pad((20, 80, 20, 30)),
CenterCrop(crop),
Scale((256, 256)),
ToTensor()
])
self.transform = Compose([
Scale((256, 256)),
Pad((24, 24, 24, 24)),
CenterCrop(precrop),
Scale((256, 256)),
ToTensor()
])
def transform_target(crop_size):
"""Ground truth image
"""
return Compose([
RandomCrop(crop_size),
RandomHorizontalFlip(),
])
def build_dataset(source_domain_name,
target_domain_name):
""" Build torch DataSet
Args:
source_domain_name (string): name of source domain dataset.
target_domain_name (string): name of target domain dataset.
Returns:
datasets (dict): dictionary mapping domain_name (string) to torch Dataset.
"""
# Define transforms for training and evaluation
transform_train = Compose([Resize([256, 256]),
RandomCrop([224, 224]),
RandomHorizontalFlip(),
RandomRotation(degrees=30, fill=128),
ToTensor(),
Normalize(IMAGENET_MEAN, IMAGENET_STD)])
transform_eval = Compose([Resize([256, 256]),
CenterCrop([224, 224]),
ToTensor(),
Normalize(IMAGENET_MEAN, IMAGENET_STD)])
datasets = {}
datasets['train_source'] = ImageFolder(root=root_dir[source_domain_name],
transform=transform_train)
datasets['train_target'] = ImageFolder(root=root_dir[target_domain_name],
transform=transform_train)
datasets['test'] = ImageFolder(root=root_dir[target_domain_name],
transform=transform_eval)
return datasets
def transform_train(self, data, targets=None, batch_size=None):
""" Transform the training data, perform random cropping data augmentation and basic random flip augmentation.
Args:
data: Numpy array. The data to be transformed.
batch_size: int batch_size.
targets: the target of training set.
Returns:
A DataLoader class instance.
"""
short_edge_length = min(data.shape[1], data.shape[2])
common_list = [Normalize(torch.Tensor(self.mean), torch.Tensor(self.std))]
if self.augment:
compose_list = [ToPILImage(),
RandomCrop(data.shape[1:3], padding=4),
RandomHorizontalFlip(),
ToTensor()
] + common_list + [Cutout(n_holes=Constant.CUTOUT_HOLES,
length=int(short_edge_length * Constant.CUTOUT_RATIO))]
else:
compose_list = common_list
if len(data.shape) != 4:
compose_list = []
dataset = self._transform(compose_list, data, targets)
if batch_size is None:
batch_size = Constant.MAX_BATCH_SIZE
batch_size = min(len(data), batch_size)
return DataLoader(dataset, batch_size=batch_size, shuffle=True)
def __getitem__(self, index):
input = self.images[index % self.len]
if self.augmentations is not None:
input = self.augmentations(input)
if self.crop_size != -1:
# For training, take random crop.
if input.size[0] < self.crop_size or input.size[1] < self.crop_size:
input = input.resize((self.crop_size, self.crop_size), Image.BICUBIC)
else:
input = RandomCrop(self.crop_size)(input)
else:
# For testing, we want to take the whole image.
width, height = input.size[:2]
width = width - (width % self.upscale_factor)
height = height - (height % self.upscale_factor)
input = CenterCrop((height, width))(input)
# Make a high-resolution copy.
target = input.copy()
# Downsample to create image at low-res.
# We already make sure that crop_size divides upscale_factor.
input = input.resize(
(input.size[0]//self.upscale_factor, input.size[1]//self.upscale_factor),
Image.BICUBIC)
# Upsample using bicubic interpolation.
input = input.resize(target.size, Image.BICUBIC)
input = ToTensor()(input)
target = ToTensor()(target)
return input, target
def __init__(self,
*args,
num_classes=100,
keep_index=True,
img_size=224,
train=True,
load_features_path=None,
save_features_path=None,
extract_at_layer=None,
feature_extractor=None,
**kwargs):
super(ExtendedDataset, self).__init__(*args, train=train, **kwargs)
self.train = train
self.num_classes = num_classes
self.data_index = None
self.keep_index = keep_index
if keep_index:
self.data_index = np.arange(len(self))
self.mean_image = None
self.resize = Resize(img_size)
self.augment = Compose(
[RandomCrop(img_size, padding=8),
RandomHorizontalFlip()])
self._set_mean_image()
self._set_data()
self.layer = extract_at_layer
self.use_feature_data = False
self.feature_data_set = False
if extract_at_layer:
self.use_feature_data = True
self._set_feature_data(feature_extractor,
extract_at_layer,
load_features_path=load_features_path,
save_features_path=save_features_path)
def __init__(self, opt, val=False):
super(CustomCIFAR100, self).__init__()
dir_dataset = opt.dir_dataset
if val:
self.dataset = CIFAR100(root=dir_dataset,
train=False,
download=True)
self.transform = Compose([
ToTensor(),
Normalize(mean=[0.507, 0.487, 0.441],
std=[0.267, 0.256, 0.276])
])
else:
self.dataset = CIFAR100(root=dir_dataset,
train=True,
download=True)
self.transform = Compose([
RandomCrop((32, 32),
padding=4,
fill=0,
padding_mode='constant'),
RandomHorizontalFlip(),
ToTensor(),
Normalize(mean=[0.507, 0.487, 0.441],
std=[0.267, 0.256, 0.276])
])
def get_train_test_loaders(dataset_name, path, batch_size, num_workers):
assert dataset_name in datasets.__dict__, "Unknown dataset name {}".format(dataset_name)
fn = datasets.__dict__[dataset_name]
train_transform = Compose([
Pad(2),
RandomCrop(32),
RandomHorizontalFlip(),
ToTensor(),
Normalize((0.5, 0.5, 0.5), (0.25, 0.25, 0.25)),
])
test_transform = Compose([
ToTensor(),
Normalize((0.5, 0.5, 0.5), (0.25, 0.25, 0.25)),
])
train_ds = fn(root=path, train=True, transform=train_transform, download=True)
test_ds = fn(root=path, train=False, transform=test_transform, download=False)
train_loader = DataLoader(train_ds, batch_size=batch_size, num_workers=num_workers, pin_memory=True)
test_loader = DataLoader(test_ds, batch_size=batch_size * 2, num_workers=num_workers, pin_memory=True)
return train_loader, test_loader
def transform_train(self, data, targets=None, batch_size=None):
short_edge_length = min(data.shape[1], data.shape[2])
common_list = [
Normalize(torch.Tensor(self.mean), torch.Tensor(self.std))
]
if self.augment:
compose_list = [
ToPILImage(),
RandomCrop(data.shape[1:3], padding=4),
RandomHorizontalFlip(),
ToTensor()
] + common_list + [
Cutout(n_holes=Constant.CUTOUT_HOLES,
length=int(short_edge_length * Constant.CUTOUT_RATIO))
]
else:
compose_list = common_list
dataset = self._transform(compose_list, data, targets)
if batch_size is None:
batch_size = Constant.MAX_BATCH_SIZE
batch_size = min(len(data), batch_size)
return DataLoader(dataset, batch_size=batch_size, shuffle=True)
def _video_transform(self, mode: str):
"""
This function contains example transforms using both PyTorchVideo and TorchVision
in the same Callable. For 'train' mode, we use augmentations (prepended with
'Random'), for 'val' mode we use the respective determinstic function.
"""
args = self.args
return ApplyTransformToKey(
key="video",
transform=Compose(
[
UniformTemporalSubsample(args.video_num_subsampled),
Normalize(args.video_means, args.video_stds),
]
+ (
[
RandomShortSideScale(
min_size=args.video_min_short_side_scale,
max_size=args.video_max_short_side_scale,
),
RandomCrop(args.video_crop_size),
RandomHorizontalFlip(p=args.video_horizontal_flip_p),
]
if mode == "train"
else [
ShortSideScale(args.video_min_short_side_scale),
CenterCrop(args.video_crop_size),
]
)
),
)
def __getitem__(self, index):
image = Image.open(os.path.join(
self.root,
self.list_paths[index])) # Open image from the given path.
# Get transform list
list_transforms = list()
if self.crop_size > 0:
list_transforms.append(RandomCrop(
(self.crop_size, self.crop_size)))
if self.flip:
coin = random.random() > 0.5
if coin:
list_transforms.append(RandomHorizontalFlip())
transforms = Compose(list_transforms)
image = transforms(image) # Implement common transform
input_image = Grayscale(num_output_channels=1)(
image) # For input image, we need to make it B/W.
input_tensor, target_tensor = ToTensor()(input_image), ToTensor()(
image) # Make input, target as torch.Tensor
input_tensor = Normalize(mean=[0.5], std=[0.5])(
input_tensor) # As the input tensor has only one channel,
# Normalize parameters also have one value each.
target_tensor = Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])(
target_tensor) # As the target tensor has
# three channels Normalize parameters also have three values each.
return input_tensor, target_tensor
def __init__(
self,
wrapped_dataset: VisionDataset,
scale_factor: int = 2,
output_size: Optional[Union[int, Tuple[int, int]]] = None,
deterministic: bool = False,
scale_mode: int = Image.BICUBIC,
base_image_transform=None,
transform=None,
target_transform=None
):
super(IsrDataset, self).__init__(wrapped_dataset.root,
transform=transform,
target_transform=target_transform)
assert scale_factor > 1, 'scale factor must be >= 2'
self.scale_factor = scale_factor
self.scale_mode = scale_mode
self._dataset = wrapped_dataset
self.base_img_transform = base_image_transform
if output_size is None:
self._crop = None
elif deterministic:
self._crop = CenterCrop(size=output_size)
else:
self._crop = RandomCrop(size=output_size)
def train_hr_transform(crop_size):
return Compose([
#ToPILImage(mode = 'RGB'),
Resize(crop_size),
RandomCrop(crop_size),
ToTensor(),
])
def img_transform(self, crop_size):
normalize = Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
return Compose(
[Resize([512, 384]),
RandomCrop(crop_size),
ToTensor(), normalize])
def __call__(self, img):
for attempt in range(10):
area = img.size[0] * img.size[1]
target_area = random.uniform(0.80, 1.0) * area
aspect_ratio = random.uniform(7. / 8, 8. / 7)
w = int(round(math.sqrt(target_area * aspect_ratio)))
h = int(round(math.sqrt(target_area / aspect_ratio)))
if random.random() < 0.5:
w, h = h, w
if w <= img.size[0] and h <= img.size[1]:
x1 = random.randint(0, img.size[0] - w)
y1 = random.randint(0, img.size[1] - h)
img = img.crop((x1, y1, x1 + w, y1 + h))
assert (img.size == (w, h))
return img.resize((self.size, self.size), self.interpolation)
# Fallback
#scale = Scale(self.size, interpolation=self.interpolation)
#crop = CenterCrop(self.size)
#return crop(scale(img))
rcrop = RandomCrop(self.size)
return rcrop(img)
def default_transforms(self) -> Dict[str, Callable]:
if self.training:
post_tensor_transform = [
RandomShortSideScale(min_size=256, max_size=320),
RandomCrop(244),
RandomHorizontalFlip(p=0.5),
]
else:
post_tensor_transform = [
ShortSideScale(256),
]
return {
"post_tensor_transform": Compose([
ApplyTransformToKey(
key="video",
transform=Compose([UniformTemporalSubsample(8)] + post_tensor_transform),
),
]),
"per_batch_transform_on_device": Compose([
ApplyTransformToKey(
key="video",
transform=K.VideoSequential(
K.Normalize(torch.tensor([0.45, 0.45, 0.45]), torch.tensor([0.225, 0.225, 0.225])),
data_format="BCTHW",
same_on_frame=False
)
),
]),
}
def ImageTransform(loadSize, cropSize):
return Compose([
Resize(size=loadSize, interpolation=Image.BICUBIC),
RandomCrop(size=cropSize),
RandomHorizontalFlip(p=0.5),
ToTensor(),
])
def HR_8_transform(crop_size):
return Compose([
RandomCrop(crop_size),
RandomScale(),
#RandomRotate(),
RandomHorizontalFlip(),
])
def __init__(self, opts):
self.dataroot = opts.dataroot
# A
images_A = os.listdir(os.path.join(self.dataroot, opts.phase + 'A'))
self.A = [os.path.join(self.dataroot, opts.phase + 'A', x) for x in images_A]
# B
images_B = os.listdir(os.path.join(self.dataroot, opts.phase + 'B'))
self.B = [os.path.join(self.dataroot, opts.phase + 'B', x) for x in images_B]
self.A_size = len(self.A)
self.B_size = len(self.B)
self.dataset_size = max(self.A_size, self.B_size)
self.input_dim_A = opts.input_dim_a
self.input_dim_B = opts.input_dim_b
# setup image transformation
transforms = [Resize((opts.resize_size, opts.resize_size), Image.BICUBIC)]
if opts.phase == 'train':
transforms.append(RandomCrop(opts.crop_size))
else:
transforms.append(CenterCrop(opts.crop_size))
if not opts.no_flip:
transforms.append(RandomHorizontalFlip())
transforms.append(ToTensor())
transforms.append(Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]))
self.transforms = Compose(transforms)
print('A: %d, B: %d images'%(self.A_size, self.B_size))
return
def __init__(self, opts):
self.dataroot = opts.dataroot
self.num_domains = opts.num_domains
self.input_dim = opts.input_dim
self.nz = opts.input_nz
domains = [chr(i) for i in range(ord('A'),ord('Z')+1)]
self.images = [None]*self.num_domains
stats = ''
for i in range(self.num_domains):
img_dir = os.path.join(self.dataroot, opts.phase + domains[i])
ilist = os.listdir(img_dir)
self.images[i] = [os.path.join(img_dir, x) for x in ilist]
stats += '{}: {}'.format(domains[i], len(self.images[i]))
stats += ' images'
self.dataset_size = max([len(self.images[i]) for i in range(self.num_domains)])
# setup image transformation
transforms = [Resize((opts.resize_size, opts.resize_size), Image.BICUBIC)]
if opts.phase == 'train':
transforms.append(RandomCrop(opts.img_size))
else:
transforms.append(CenterCrop(opts.img_size))
if not opts.no_flip:
transforms.append(RandomHorizontalFlip())
transforms.append(ToTensor())
transforms.append(Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]))
self.transforms = Compose(transforms)
return
def get_cifar10_loaders(data_root: str, batch_size: int, num_workers: int, augment=True) -> (DataLoader, DataLoader):
""" Function for retrieving CIFAR10 data as data loaders.
Training set is augmented. Note that ToTensor() automatically divides by 255.
Args:
data_root: Path to CIFAR10 data.
batch_size: Batch size of data.
num_workers: Number of workers to pre-process data.
augment: Whether to use data augmentation on the training data.
Returns:
Training and evaluation data loaders for the CIFAR10 dataset.
"""
# Normalize(mean=[0.491, 0.482, 0.447], std=[0.247, 0.243, 0.262])
if augment:
train_transform = Compose([RandomHorizontalFlip(), RandomCrop(size=32, padding=4), ToTensor()])
else:
train_transform = ToTensor()
eval_transform = ToTensor()
train_set = CIFAR10(root=data_root, train=True, transform=train_transform, download=True)
eval_set = CIFAR10(root=data_root, train=False, transform=eval_transform, download=True)
train_loader = DataLoader(train_set, batch_size, shuffle=True, num_workers=num_workers, pin_memory=True)
eval_loader = DataLoader(eval_set, batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)
return train_loader, eval_loader
def __init__(self, seq_name, vis_threshold, P, K, max_per_person, crop_H, crop_W,
transform, normalize_mean=None, normalize_std=None):
self.data_dir = osp.join(cfg.DATA_DIR, 'Market-1501-v15.09.15')
self.seq_name = seq_name
self.P = P
self.K = K
self.max_per_person = max_per_person
self.crop_H = crop_H
self.crop_W = crop_W
if transform == "random":
self.transform = Compose([RandomCrop((crop_H, crop_W)), RandomHorizontalFlip(), ToTensor(), Normalize(normalize_mean, normalize_std)])
elif transform == "center":
self.transform = Compose([CenterCrop((crop_H, crop_W)), ToTensor(), Normalize(normalize_mean, normalize_std)])
else:
raise NotImplementedError("Tranformation not understood: {}".format(transform))
if seq_name:
assert seq_name in ['bounding_box_test', 'bounding_box_train', 'gt_bbox'], \
'Image set does not exist: {}'.format(seq_name)
self.data = self.load_images()
else:
self.data = []
self.build_samples()
def train_data(self) -> Tuple[DataLoader, NoisyDataset, Sampler]:
"""Configure the training set using the current configuration.
Returns:
Tuple[Dataset, DataLoader, Sampler]: Returns a NoisyDataset object
wrapping either a folder or HDF5 dataset, a DataLoader for that
dataset that uses a FixedLengthSampler (also returned).
"""
cfg = self.cfg
transform = RandomCrop(
cfg[ConfigValue.TRAIN_PATCH_SIZE],
pad_if_needed=True,
padding_mode="reflect",
)
# Load dataset
if cfg[ConfigValue.TRAIN_DATASET_TYPE] == DatasetType.FOLDER:
dataset = UnlabelledImageFolderDataset(
cfg[ConfigValue.TRAIN_DATA_PATH],
channels=cfg[ConfigValue.IMAGE_CHANNELS],
transform=transform,
recursive=True,
)
elif cfg[ConfigValue.TRAIN_DATASET_TYPE] == DatasetType.HDF5:
# It is assumed that the created dataset does not violate the minimum patch size
dataset = HDF5Dataset(
cfg[ConfigValue.TRAIN_DATA_PATH],
transform=transform,
channels=cfg[ConfigValue.IMAGE_CHANNELS],
)
else:
raise NotImplementedError("Dataset type not implemented")
# Wrap dataset for creating noisy samples
dataset = NoisyDataset(dataset,
cfg[ConfigValue.NOISE_STYLE],
cfg[ConfigValue.ALGORITHM],
pad_uniform=False,
pad_multiple=NoiseNetwork.input_wh_mul(),
square=cfg[ConfigValue.BLINDSPOT],
training_mode=True)
# Ensure dataset initialised by loading first bit of data
_ = dataset[0]
# Create a dataloader that will sample from this dataset for a fixed number of samples
sampler = FixedLengthSampler(
dataset,
num_samples=cfg[ConfigValue.TRAIN_ITERATIONS],
shuffled=True,
)
# Resume train sampler
if self._train_iter is not None:
sampler.for_next_iter(self._train_iter)
self._train_iter = None
dataloader = DataLoader(
dataset,
sampler=sampler,
batch_size=cfg[ConfigValue.TRAIN_MINIBATCH_SIZE],
num_workers=cfg[ConfigValue.DATALOADER_WORKERS],
pin_memory=cfg[ConfigValue.PIN_DATA_MEMORY],
)
return dataloader, dataset, sampler
def train_hr_transform(crop_size):
return Compose([
RandomCrop(crop_size),
# Resize((128,128), interpolation=Image.BICUBIC),
ToTensor()
# Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
