Hi i tried implementing code to get hflip and vflip of both my images and masks … but i keep getting this error … any advice is much appreciated
My masks are numpy arrays where each pixel is represented by the class index
My images are jpg files
#Dataset class for BiopsyData (used by the data loader)
class BiopsyDataset(Dataset):
def __init__(self, root_data_dir, ids):
files = []
files2 = []
#Get all file names
for file in os.listdir(root_data_dir+'/CroppedImages'):
if file.endswith('.jpg'):
files.append( file )
files = np.sort(files).tolist()
#Get all file names masks
for file in os.listdir(root_data_dir+'/Masks'):
if file.endswith('.npy'):
files2.append( file )
files2 = np.sort(files2).tolist()
#Create dataset from specified ids
self.files = [files[i] for i in ids]
self.files2 = [files2[i] for i in ids]
#Transforms
self.to_tensor = torchvision.transforms.ToTensor()
self.root_data_dir = root_data_dir
def __len__(self):
return len(self.files)
def transform(self, image, mask):
image = TF.to_pil_image(image)
# Random horizontal flipping
if random.random() > 0.5:
image = TF.hflip(image)
mask = np.flip(mask, 1)
# Random vertical flipping
if random.random() > 0.5:
image = TF.vflip(image)
mask = np.flip(mask, 0)
# Transform to tensor
image = TF.to_tensor(image)
image = TF.normalize(image, mean, std)
return image, mask
#Returns a single image and label pair
def __getitem__(self, index):
#Read image and labels
image = Image.open(self.root_data_dir + '/CroppedImages/' + self.files[index])
Mask = np.load(self.root_data_dir + '/Masks/' + self.files2[index])
image = self.to_tensor(image)
image, Mask = self.transform(image, Mask)
#Cross entropy loss needs labels as LongTensor type
Mask = self.to_tensor(Mask).type(torch.LongTensor)
return image, Mask