I’m not sure the way that I am using data augmentation for my semantic segmentation task is working properly. I define my own pyTorch Dataset PositiveOnly
of images and their masks. However, when I try to plot the same instance of my dataset, the image and mask do not align. I have 3 channel images, and 1 channel masks.
My function to apply simple rotations to a semantic segmentation dataset is here, as well as the custom ToTensor function, and the overall wrapper for my dataset:
class Randomize(object):
def __call__(self, sample):
imgdata = sample['img']
fptdata = sample['fpt']
# mirror horizontally
mirror = np.random.randint(0, 2)
if mirror:
imgdata = np.flip(imgdata, 2)
fptdata = np.flip(fptdata, 1)
# flip vertically
flip = np.random.randint(0, 2)
if flip:
imgdata = np.flip(imgdata, 1)
fptdata = np.flip(fptdata, 0)
# rotate by [0,1,2,3]*90 deg
rot = np.random.randint(0, 4)
imgdata = np.rot90(imgdata, rot, axes=(1,2))
fptdata = np.rot90(fptdata, rot, axes=(0,1))
return {'img': imgdata.copy(),
'fpt': fptdata.copy()}
class ToTensor(object):
def __call__(self, sample):
out = {'img': torch.from_numpy(sample['img'].copy()),
'fpt': torch.from_numpy(sample['fpt'].copy())}
return out
def create_dataset(*args, apply_transforms=True, **kwargs):
if apply_transforms:
data_transforms = transforms.Compose([
Randomize(),
ToTensor(),
])
else:
data_transforms = None
data = PositiveOnly(*args, **kwargs, transform=data_transforms)
return data
I’ve broken down the Randomize
function line by line and it does appear to work, but once i call it in my dataset prior to the dataloader I see they don’t align.
My dataset is defined as:
class PositiveOnly(Dataset):
def __init__(self, p_pkl, transform=None):
self.p_pkl = p_pkl
self.transform = transform
with open(p_pkl, "rb") as fp:
self.p_pkl = pkl.load(fp)
def __len__(self):
return len(self.p_pkl)
def __getitem__(self, idx):
sample = self.p_pkl[idx]
# print(sample)
if self.transform:
sample = self.transform(sample)
return sample
Here is the code I run to see that they don’t align:
p = create_dataset('data.pkl',
apply_transforms=True)
pic = np.moveaxis(np.array(p[0]['img']), 0,2) #is a pytorch tensor with shape (3, 288,288)
mak = p[0]['fpt']
plt.imshow(pic)
plt.show()
plt.imshow(mak)
plt.show()
Is there something wrong with my Randomize function, or the way that I am calling to plot them? I plan to run this in a deep learning model but need them to align so the results aren’t messed up.