title: How to use torchvision.transforms for data augmentation of segmentation task in Pytorch?
body: <p>I am a little bit confused about the data augmentation performed in PyTorch. </p>
  
  <p>Because we are dealing with segmentation tasks, we need data and mask for the same data augmentation, but some of them are random, such as random rotation. </p>
  
  <p>Keras provides a <code>random seed</code> guarantee that data and mask do the same operation, as shown in the following code:</p>
  
  <pre><code>    data_gen_args = dict(featurewise_center=True,
                           featurewise_std_normalization=True,
                           rotation_range=25,
                           horizontal_flip=True,
                           vertical_flip=True)
  
  
      image_datagen = ImageDataGenerator(**data_gen_args)
      mask_datagen = ImageDataGenerator(**data_gen_args)
  
      seed = 1
      image_generator = image_datagen.flow(train_data, seed=seed, batch_size=1)
      mask_generator = mask_datagen.flow(train_label, seed=seed, batch_size=1)
  
      train_generator = zip(image_generator, mask_generator)
  </code></pre>
  
  <p>I didn't find a similar description in the official Pytorch documentation, so I don't know how to ensure that data and mask can be processed synchronously.</p>
  
  <p>Pytorch does provide such a function, but I want to apply it to a custom Dataloader. </p>
  
  <p>For example：</p>
  
  <pre><code>def __getitem__(self, index):
      img = np.zeros((self.im_ht, self.im_wd, channel_size))
      mask = np.zeros((self.im_ht, self.im_wd, channel_size))
  
      temp_img = np.load(Image_path + '{:0&gt;4}'.format(self.patient_index[index]) + '.npy')
      temp_label = np.load(Label_path + '{:0&gt;4}'.format(self.patient_index[index]) + '.npy')
  
      for i in range(channel_size):
          img[:,:,i] = temp_img[self.count[index] + i]
          mask[:,:,i] = temp_label[self.count[index] + i]
  
      if self.transforms:
          img = np.uint8(img)
          mask = np.uint8(mask)
          img = self.transforms(img)
          mask = self.transforms(mask)
  
      return img, mask
  </code></pre>
  
  <p>In this case, img and mask will be transformed separately, because some operations such as random rotation are random, so the correspondence between mask and image may be changed. In other words, the image may have rotated but the mask did not do this.</p>
  
  <h1>EDIT 1</h1>
  
  <p>I used the method in <a href="https://github.com/a514514772/DISE-Domain-Invariant-Structure-Extraction/blob/master/util/loader/augmentations.py" rel="noreferrer">augmentations.py</a>, but I got an error:：</p>
  
  <pre><code>Traceback (most recent call last):
    File "test_transform.py", line 87, in &lt;module&gt;
      for batch_idx, image, mask in enumerate(train_loader):
    File "/home/dirk/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 314, in __next__
      batch = self.collate_fn([self.dataset[i] for i in indices])
    File "/home/dirk/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 314, in &lt;listcomp&gt;
      batch = self.collate_fn([self.dataset[i] for i in indices])
    File "/home/dirk/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/utils/data/dataset.py", line 103, in __getitem__
      return self.dataset[self.indices[idx]]
    File "/home/dirk/home/data/dirk/segmentation_unet_pytorch/data.py", line 164, in __getitem__
      img, mask = self.transforms(img, mask)
    File "/home/dirk/home/data/dirk/segmentation_unet_pytorch/augmentations.py", line 17, in __call__
      img, mask = a(img, mask)
  TypeError: __call__() takes 2 positional arguments but 3 were given
  </code></pre>
  
  <p>This is my code for <code>__getitem__()</code>：</p>
  
  <pre><code>data_transforms = {
      'train': Compose([
          RandomHorizontallyFlip(),
          RandomRotate(degree=25),
          transforms.ToTensor()
      ]),
  }
  
  train_set = DatasetUnetForTestTransform(fold=args.fold, random_index=args.random_index,transforms=data_transforms['train'])
  
  # __getitem__ in class DatasetUnetForTestTransform
  def __getitem__(self, index):
      img = np.zeros((self.im_ht, self.im_wd, channel_size))
      mask = np.zeros((self.im_ht, self.im_wd, channel_size))
      temp_img = np.load(Label_path + '{:0&gt;4}'.format(self.patient_index[index]) + '.npy')
      temp_label = np.load(Label_path + '{:0&gt;4}'.format(self.patient_index[index]) + '.npy')
      temp_img, temp_label = crop_data_label_from_0(temp_img, temp_label)
      for i in range(channel_size):
          img[:,:,i] = temp_img[self.count[index] + i]
          mask[:,:,i] = temp_label[self.count[index] + i]
  
      if self.transforms:
          img = T.ToPILImage()(np.uint8(img))
          mask = T.ToPILImage()(np.uint8(mask))
          img, mask = self.transforms(img, mask)
  
      img = T.ToTensor()(img).copy()
      mask = T.ToTensor()(mask).copy()
      return img, mask
  </code></pre>
  
  <h1>EDIT 2</h1>
  
  <p>I found that after ToTensor, the dice between the same labels becomes 255 instead of 1, how to fix it?</p>
  
  <pre><code># Dice computation
  def DSC_computation(label, pred):
      pred_sum = pred.sum()
      label_sum = label.sum()
      inter_sum = np.logical_and(pred, label).sum()
      return 2 * float(inter_sum) / (pred_sum + label_sum)
  </code></pre>
  
  <p>Feel free to ask if more code is needed to explain the problem.</p>