def get_loader_bsds(directory=img_dir,
batch_size=16,
train=True,
num_workers=0,
pin_memory=True,
gray_scale=True,
crop_size=[32, 32]):
# 32 x 32 crop
if train:
if gray_scale:
transform = preprocess.scale_random_crop_gray(
crop_size[0], crop_size[1])
else:
transform = preprocess.scale_random_crop(crop_size[0],
crop_size[1])
else:
if gray_scale:
transform = preprocess.central_crop_gray(crop_size[0],
crop_size[1])
else:
transform = preprocess.scale_crop(crop_size[0], crop_size[1])
dataset = img_dataset.PlainImageFolder(root=directory, transform=transform)
shuffle = train
loader = data.DataLoader(dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=shuffle,
pin_memory=False)
return loader
def get_loader_denoising(directory=img_dir,
batch_size=16,
train=True,
num_workers=0,
pin_memory=True,
gray_scale=True,
crop_size=[80, 80]):
# 32 x 32 crop
if gray_scale:
transform = preprocess.central_crop_gray(crop_size[0], crop_size[1])
else:
transform = preprocess.scale_crop(crop_size[0], crop_size[1])
# =============================================================================
# transform = transforms.Compose([
# #transforms.Grayscale(), # Convert to grayscale
# transforms.ToTensor(),
# ])
#
# =============================================================================
dataset = img_dataset.PlainImageFolder(root=directory, transform=transform)
shuffle = train
loader = data.DataLoader(dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=shuffle,
pin_memory=False)
return loader
def get_loader_mask(directory=img_dir,
batch_size=1,
train=True,
num_workers=0,
pin_memory=True,
gray_scale=True,
crop_size=[80, 80]):
if crop_size == None: #Do not crop
# 32 x 32 crop
if gray_scale:
transform = transforms.Compose([
transforms.Grayscale(), # Convert to grayscale
transforms.ToTensor(),
])
else:
transform = transforms.Compose([
transforms.ToTensor(),
])
else:
# crop
if gray_scale:
transform = transforms.Compose([
transforms.CenterCrop(crop_size[0]),
transforms.Grayscale(),
transforms.ToTensor()
])
else:
transform = transforms.Compose(
[transforms.CenterCrop(crop_size[0]),
transforms.ToTensor()])
dataset = img_dataset.PlainImageFolder(root=directory, transform=transform)
shuffle = train
loader = data.DataLoader(dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=shuffle,
pin_memory=False)
return loader