Hello, Pytorch community,
I want to apply some argumentations transform for one image and two ground truths say mask and mask1. Transformations are applied on image and mask but mask1 is not undergoing through any transformation.
Thank You
My code looks like this:
train_transforms = A.Compose(
[
A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH),
ToTensorV2(),
],
)
class Dataset_(Dataset):
def init(self, image_dir, mask_dir,mask_dir1 transform=None):
self.image_dir = image_dir
self.mask_dir = mask_dir
self.mask_dir1 = mask_dir1
self.transform = transform
self.images = os.listdir(image_dir)
self.masks = os.listdir(mask_dir)
self.masks1 = os.listdir(mask_dir1)
def __len__(self):
return len(self.images)
def __getitem__(self, index):
img_path = os.path.join(self.image_dir, self.images[index])
mask_path = os.path.join(self.mask_dir, self.images[index][:-4]+'_gt.npy')
mask_path1 = os.path.join(self.mask_dir1, self.images[index][:-4]+'_gt.npy')
image = np.load(img_path,allow_pickle=True, fix_imports=True)
mask = np.load(mask_path,allow_pickle=True, fix_imports=True)
mask1 = np.load(mask_path1,allow_pickle=True, fix_imports=True)
#mask=np.expand_dims(mask, axis=0)
if self.transform is not None:
augmentations = self.transform(image=image,mask=mask,mask1=mask1)
image = augmentations["image"]
mask = augmentations["mask"]
mask1 = augmentations["mask1"]
return image, mask, mask1
def get_loaders(
train_dir=TRAIN_IMG_DIR,
train_maskdir=TRAIN_MASK_DIR,
train_maskdir1=TRAIN_MASK_DIR1,
val_dir=VAL_IMG_DIR,
val_maskdir=VAL_MASK_DIR,
batch_size=BATCH_SIZE,
train_transform=train_transforms,
val_transform=val_transforms,
num_workers=NUM_WORKERS,
pin_memory=PIN_MEMORY,
):
train_ds = Dataset_(
image_dir=train_dir,
mask_dir=train_maskdir,train_maskdir1,
transform=train_transform,
)