Update dataset.py

AttributeError: module 'torchvision.transforms' has no attribute 'InterpolationMode'
pull/9/head
Dominic-ZZ 3 years ago committed by GitHub
parent 69a0b15bc0
commit e8e48a2070
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -30,10 +30,10 @@ class InpaintingData(Dataset):
transforms.ColorJitter(0.05, 0.05, 0.05, 0.05), transforms.ColorJitter(0.05, 0.05, 0.05, 0.05),
transforms.ToTensor()]) transforms.ToTensor()])
self.mask_trans = transforms.Compose([ self.mask_trans = transforms.Compose([
transforms.Resize(args.image_size, interpolation=transforms.InterpolationMode.NEAREST), transforms.Resize(args.image_size),
transforms.RandomHorizontalFlip(), transforms.RandomHorizontalFlip(),
transforms.RandomRotation( transforms.RandomRotation(
(0, 45), interpolation=transforms.InterpolationMode.NEAREST), (0, 45)),
]) ])
@ -77,4 +77,4 @@ if __name__ == '__main__':
data = InpaintingData(args) data = InpaintingData(args)
print(len(data), len(data.mask_path)) print(len(data), len(data.mask_path))
img, mask, filename = data[0] img, mask, filename = data[0]
print(img.size(), mask.size(), filename) print(img.size(), mask.size(), filename)

Loading…
Cancel
Save