Is it possible to combine aug_transforms with some transforms of Albumentations for Segmentation?

I have tried the next code without exit:

aug_transforms(mult=1.0, do_flip=True, 
                                flip_vert=True, max_rotate=10., 
                                max_zoom=1.1,max_warp=0.2, 
                                p_affine=0.75, p_lighting=0, 
                                xtra_tfms=[ElasticTransform,
                                           GridDistortion, 
                                           OpticalDistortion])

ElasticTransform,GridDistortion, OpticalDistortion are the ones from Albumentations that I would like to use.

Albumentations is used as follows: aug(image=image, mask=mask)

You would need to wrap them as ItemTransform to make them applied on the tuple as a whole and not on each part of the tuple. Also I’m not sure albumentation provides augmentation on a batch of images on the GPU, so you should probably pass those as item_tfms and not inside aug_transforms.

Okey, I will let you know when I try it.

I have tried the following:

    from albumentations import (
        ElasticTransform,
        GridDistortion, 
        OpticalDistortion 
    )
    from albumentations.pytorch.transforms import ToTensorV2


    def tensor2np(x):
        np_image = x.cpu().numpy()
        np_image = (np_image * 255).astype(np.uint8)
        
        return np_image

    class FastAIElasticTransform(ItemTransform):
        def __init__(self,probability):
            self.probability=probability
            self.transform=ElasticTransform(p=self.probability)
            
        def encodes(self,x):
            
            images = x[0]
            masks = x[1]
            
            tensor_images = []
            tensor_masks = []
            for i in range(0,images.shape[0]):
                np_image = images[i]
                np_mask = masks[i]
                np_image = tensor2np(np_image)
                np_mask = tensor2np(np_mask)


                # apply transform
                transformed = self.transform(image=np_image,mask=np_mask)
                transformed = ToTensorV2()(image=transformed["image"],mask=transformed["mask"])

                # back to tensor
                tensor_image = transformed["image"]
                tensor_image = tensor_image.reshape(3,1002,1002)
                tensor_images.append(tensor_image)
                
                tensor_mask = transformed["mask"]
                tensor_masks.append(tensor_mask)
                
            
            return (torch.stack(tensor_images),torch.stack(tensor_masks))
        
    class FastAIGridDistortion(ItemTransform):
        def __init__(self,probability):
            self.probability=probability
            self.transform=GridDistortion(p=self.probability)
            
        def encodes(self,x):
            images = x[0]
            masks = x[1]
            
            tensor_images = []
            tensor_masks = []
            for i in range(0,images.shape[0]):
                np_image = images[i]
                np_mask = masks[i]
                np_image = tensor2np(np_image)
                np_mask = tensor2np(np_mask)


                # apply transform
                transformed = self.transform(image=np_image,mask=np_mask)
                transformed = ToTensorV2()(image=transformed["image"],mask=transformed["mask"])

                # back to tensor
                tensor_image = transformed["image"]
                tensor_image = tensor_image.reshape(3,1002,1002)
                tensor_images.append(tensor_image)
                
                tensor_mask = transformed["mask"]
                tensor_masks.append(tensor_mask)
                
            
            return (torch.stack(tensor_images),torch.stack(tensor_masks))
        
    class FastAIOpticalDistortion(ItemTransform):
        def __init__(self,probability):
            self.probability=probability
            self.transform=OpticalDistortion(p=self.probability)
            
        def encodes(self,x):
            images = x[0]
            masks = x[1]
            
            tensor_images = []
            tensor_masks = []
            for i in range(0,images.shape[0]):
                np_image = images[i]
                np_mask = masks[i]
                np_image = tensor2np(np_image)
                np_mask = tensor2np(np_mask)


                # apply transform
                transformed = self.transform(image=np_image,mask=np_mask)
                transformed = ToTensorV2()(image=transformed["image"],mask=transformed["mask"])

                # back to tensor
                tensor_image = transformed["image"]
                tensor_image = tensor_image.reshape(3,1002,1002)
                tensor_images.append(tensor_image)
                
                tensor_mask = transformed["mask"]
                tensor_masks.append(tensor_mask)
                
            
            return (torch.stack(tensor_images),torch.stack(tensor_masks))

    manual = DataBlock(blocks=(ImageBlock, MaskBlock(codes)),
                       get_items=partial(get_image_files,folders=[manual_name]),
                       get_y=get_y_fn,
                       splitter=RandomSplitter(valid_pct=0.1,seed=2020),
                       item_tfms=[Resize((size,size))],
                       batch_tfms=[*aug_transforms(mult=1.0, do_flip=True, 
                                    flip_vert=True, max_rotate=10., 
                                    max_zoom=1.1,max_warp=0.2, 
                                    p_affine=0.75, p_lighting=0, 
                                    xtra_tfms=[FastAIElasticTransform(0.25),
                                              FastAIGridDistortion(0.25),
                                              FastAIOpticalDistortion(0.25)])
                                    ,Normalize.from_stats(*imagenet_stats)]
                      )
    manual.summary(path_images)
    dls = manual.dataloaders(path_images,bs=bs)
    dls.show_batch(vmin=0,vmax=1,figsize=(12, 9))

It raised the next error:

    Setting-up type transforms pipelines
    Collecting items from ../datasets/Images
    Found 621 items
    2 datasets of sizes 559,62
    Setting up Pipeline: PILBase.create
    Setting up Pipeline: <lambda> -> PILBase.create

    Building one sample
      Pipeline: PILBase.create
        starting from
          ../datasets/Images/manual/165.png
        applying PILBase.create gives
          PILImage mode=RGB size=1002x1004
      Pipeline: <lambda> -> PILBase.create
        starting from
          ../datasets/Images/manual/165.png
        applying <lambda> gives
          ../datasets/Labels/manual/165.png
        applying PILBase.create gives
          PILMask mode=L size=1002x1004

    Final sample: (PILImage mode=RGB size=1002x1004, PILMask mode=L size=1002x1004)


    Setting up after_item: Pipeline: AddMaskCodes -> Resize -> ToTensor
    Setting up before_batch: Pipeline: 
    Setting up after_batch: Pipeline: FastAIElasticTransform -> FastAIGridDistortion -> FastAIOpticalDistortion -> IntToFloatTensor -> AffineCoordTfm -> LightingTfm -> Normalize
    Could not do one pass in your dataloader, there is something wrong in it

    Building one batch
    Applying item_tfms to the first sample:
      Pipeline: AddMaskCodes -> Resize -> ToTensor
        starting from
          (PILImage mode=RGB size=1002x1004, PILMask mode=L size=1002x1004)
        applying AddMaskCodes gives
          (PILImage mode=RGB size=1002x1004, PILMask mode=L size=1002x1004)
        applying Resize gives
          (PILImage mode=RGB size=1002x1002, PILMask mode=L size=1002x1002)
        applying ToTensor gives
          (TensorImage of size 3x1002x1002, TensorMask of size 1002x1002)

    Adding the next 3 samples

    No before_batch transform to apply

    Collating items in a batch

    Applying batch_tfms to the batch built
      Pipeline: FastAIElasticTransform -> FastAIGridDistortion -> FastAIOpticalDistortion -> IntToFloatTensor -> AffineCoordTfm -> LightingTfm -> Normalize
        starting from
          (TensorImage of size 4x3x1002x1002, TensorMask of size 4x1002x1002)
        applying FastAIElasticTransform gives
          (Tensor of size 4x3x1002x1002, Tensor of size 4x1002x1002)
        applying FastAIGridDistortion gives
          (Tensor of size 4x3x1002x1002, Tensor of size 4x1002x1002)
        applying FastAIOpticalDistortion gives
          (Tensor of size 4x3x1002x1002, Tensor of size 4x1002x1002)
        applying IntToFloatTensor gives
          (Tensor of size 4x3x1002x1002, Tensor of size 4x1002x1002)
        applying AffineCoordTfm failed.

    ---------------------------------------------------------------------------
    RuntimeError                              Traceback (most recent call last)
    <ipython-input-141-0b19c2cb3a2d> in <module>
         13                                 ,Normalize.from_stats(*imagenet_stats)]
         14                   )
    ---> 15 manual.summary(path_images)
         16 dls = manual.dataloaders(path_images,bs=bs)
         17 dls.show_batch(vmin=0,vmax=1,figsize=(12, 9))

    ~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastai2/data/block.py in summary(self, source, bs, **kwargs)
        186         print("\nApplying batch_tfms to the batch built")
        187         b = to_device(b, dls.device)
    --> 188         b = _apply_pipeline(dls.train.after_batch, b)
        189     else: print("\nNo batch_tfms to apply")

    ~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastai2/data/block.py in _apply_pipeline(p, x)
        131         except Exception as e:
        132             print(f"    applying {name} failed.")
    --> 133             raise e
        134     return x
        135 

    ~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastai2/data/block.py in _apply_pipeline(p, x)
        127         name = f.name
        128         try:
    --> 129             x = f(x)
        130             if name != "noop": print(f"    applying {name} gives\n      {_short_repr(x)}")
        131         except Exception as e:

    ~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastai2/vision/augment.py in __call__(self, b, split_idx, **kwargs)
         30 
         31     def __call__(self, b, split_idx=None, **kwargs):
    ---> 32         self.before_call(b, split_idx=split_idx)
         33         return super().__call__(b, split_idx=split_idx, **kwargs) if self.do else b
         34 

    ~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastai2/vision/augment.py in before_call(self, b, split_idx)
        351         while isinstance(b, tuple): b = b[0]
        352         self.split_idx = split_idx
    --> 353         self.do,self.mat = True,self._get_affine_mat(b)
        354         for t in self.coord_fs: t.before_call(b)
        355 

    ~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastai2/vision/augment.py in _get_affine_mat(self, x)
        362         aff_m = _init_mat(x)
        363         if self.split_idx: return _prepare_mat(x, aff_m)
    --> 364         ms = [f(x) for f in self.aff_fs]
        365         ms = [m for m in ms if m is not None]
        366         for m in ms: aff_m = aff_m @ m

    ~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastai2/vision/augment.py in <listcomp>(.0)
        362         aff_m = _init_mat(x)
        363         if self.split_idx: return _prepare_mat(x, aff_m)
    --> 364         ms = [f(x) for f in self.aff_fs]
        365         ms = [m for m in ms if m is not None]
        366         for m in ms: aff_m = aff_m @ m

    ~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastai2/vision/augment.py in rotate_mat(x, max_deg, p, draw, batch)
        515     def _def_draw(x):   return x.new(x.size(0)).uniform_(-max_deg, max_deg)
        516     def _def_draw_b(x): return x.new_zeros(x.size(0)) + random.uniform(-max_deg, max_deg)
    --> 517     thetas = _draw_mask(x, _def_draw_b if batch else _def_draw, draw=draw, p=p, batch=batch) * math.pi/180
        518     return affine_mat(thetas.cos(), thetas.sin(), t0(thetas),
        519                      -thetas.sin(), thetas.cos(), t0(thetas))

    ~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastai2/vision/augment.py in _draw_mask(x, def_draw, draw, p, neutral, batch)
        427 def _draw_mask(x, def_draw, draw=None, p=0.5, neutral=0., batch=False):
        428     if draw is None: draw=def_draw
    --> 429     if callable(draw): res=draw(x)
        430     elif is_listy(draw):
        431         test_eq(len(draw), x.size(0))

    ~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastai2/vision/augment.py in _def_draw(x)
        513 def rotate_mat(x, max_deg=10, p=0.5, draw=None, batch=False):
        514     "Return a random rotation matrix with `max_deg` and `p`"
    --> 515     def _def_draw(x):   return x.new(x.size(0)).uniform_(-max_deg, max_deg)
        516     def _def_draw_b(x): return x.new_zeros(x.size(0)) + random.uniform(-max_deg, max_deg)
        517     thetas = _draw_mask(x, _def_draw_b if batch else _def_draw, draw=draw, p=p, batch=batch) * math.pi/180

    RuntimeError: _th_uniform_ not supported on CPUType for Byte

I have tried this also:

class FastAIElasticTransform(ItemTransform):
    def __init__(self,probability):
        self.probability=probability
        self.transform=ElasticTransform(p=self.probability)
        
    def encodes(self,x):
        
        np_image = x[0]
        np_mask = x[1]

        # apply transform
        transformed = self.transform(image=np_image,mask=np_mask)
        tensor_image = transformed["image"]
            
        tensor_mask = transformed["mask"]
            
        
        return (tensor_image,tensor_mask)
    
class FastAIGridDistortion(ItemTransform):
    def __init__(self,probability):
        self.probability=probability
        self.transform=GridDistortion(p=self.probability)
        
    def encodes(self,x):
        np_image = x[0]
        np_mask = x[1]

        # apply transform
        transformed = self.transform(image=np_image,mask=np_mask)
        tensor_image = transformed["image"]
            
        tensor_mask = transformed["mask"]
            
        
        return (tensor_image,tensor_mask)
    
class FastAIOpticalDistortion(ItemTransform):
    def __init__(self,probability):
        self.probability=probability
        self.transform=OpticalDistortion(p=self.probability)
        
    def encodes(self,x):
        np_image = x[0]
        np_mask = x[1]

        # apply transform
        transformed = self.transform(image=np_image,mask=np_mask)
        tensor_image = transformed["image"]
            
        tensor_mask = transformed["mask"]
            
        
        return (tensor_image,tensor_mask)


manual = DataBlock(blocks=(ImageBlock, MaskBlock(codes)),
                   get_items=partial(get_image_files,folders=[manual_name]),
                   get_y=get_y_fn,
                   splitter=RandomSplitter(valid_pct=0.1,seed=2020),
                   item_tfms=[Resize((size,size)),FastAIElasticTransform(0.25),
                                          FastAIGridDistortion(0.25),
                                          FastAIOpticalDistortion(0.25)],
                   batch_tfms=[*aug_transforms(mult=1.0, do_flip=True, 
                                flip_vert=True, max_rotate=10., 
                                max_zoom=1.1,max_warp=0.2, 
                                p_affine=0.75, p_lighting=0)
                                ,Normalize.from_stats(*imagenet_stats)]
                  )
manual.summary(path_images)
dls = manual.dataloaders(path_images,bs=bs)
dls.show_batch(vmin=0,vmax=1,figsize=(12, 9))

With this error

Setting-up type transforms pipelines
Collecting items from ../datasets/Images
Found 621 items
2 datasets of sizes 559,62
Setting up Pipeline: PILBase.create
Setting up Pipeline: <lambda> -> PILBase.create

Building one sample
  Pipeline: PILBase.create
    starting from
      ../datasets/Images/manual/165.png
    applying PILBase.create gives
      PILImage mode=RGB size=1002x1004
  Pipeline: <lambda> -> PILBase.create
    starting from
      ../datasets/Images/manual/165.png
    applying <lambda> gives
      ../datasets/Labels/manual/165.png
    applying PILBase.create gives
      PILMask mode=L size=1002x1004

Final sample: (PILImage mode=RGB size=1002x1004, PILMask mode=L size=1002x1004)


Setting up after_item: Pipeline: AddMaskCodes -> FastAIElasticTransform -> FastAIGridDistortion -> FastAIOpticalDistortion -> Resize -> ToTensor
Setting up before_batch: Pipeline: 
Setting up after_batch: Pipeline: IntToFloatTensor -> AffineCoordTfm -> LightingTfm -> Normalize

Building one batch
Applying item_tfms to the first sample:
  Pipeline: AddMaskCodes -> FastAIElasticTransform -> FastAIGridDistortion -> FastAIOpticalDistortion -> Resize -> ToTensor
    starting from
      (PILImage mode=RGB size=1002x1004, PILMask mode=L size=1002x1004)
    applying AddMaskCodes gives
      (PILImage mode=RGB size=1002x1004, PILMask mode=L size=1002x1004)
    applying FastAIElasticTransform gives
      (PILImage mode=RGB size=1002x1004, PILMask mode=L size=1002x1004)
    applying FastAIGridDistortion gives
      (PILImage mode=RGB size=1002x1004, PILMask mode=L size=1002x1004)
    applying FastAIOpticalDistortion gives
      (PILImage mode=RGB size=1002x1004, PILMask mode=L size=1002x1004)
    applying Resize gives
      (PILImage mode=RGB size=1002x1002, PILMask mode=L size=1002x1002)
    applying ToTensor gives
      (TensorImage of size 3x1002x1002, TensorMask of size 1002x1002)

Adding the next 3 samples

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-150-4a5443614798> in <module>
     12                                 ,Normalize.from_stats(*imagenet_stats)]
     13                   )
---> 14 manual.summary(path_images)
     15 dls = manual.dataloaders(path_images,bs=bs)
     16 dls.show_batch(vmin=0,vmax=1,figsize=(12, 9))

~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastai2/data/block.py in summary(self, source, bs, **kwargs)
    163         s = [_apply_pipeline(dls.train.after_item, dsets.train[0])]
    164         print(f"\nAdding the next {bs-1} samples")
--> 165         s += [dls.train.after_item(dsets.train[i]) for i in range(1, bs)]
    166     else:
    167         print("No item_tfms to apply")

~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastai2/data/block.py in <listcomp>(.0)
    163         s = [_apply_pipeline(dls.train.after_item, dsets.train[0])]
    164         print(f"\nAdding the next {bs-1} samples")
--> 165         s += [dls.train.after_item(dsets.train[i]) for i in range(1, bs)]
    166     else:
    167         print("No item_tfms to apply")

~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastcore/transform.py in __call__(self, o)
    183         self.fs.append(t)
    184 
--> 185     def __call__(self, o): return compose_tfms(o, tfms=self.fs, split_idx=self.split_idx)
    186     def __repr__(self): return f"Pipeline: {' -> '.join([f.name for f in self.fs if f.name != 'noop'])}"
    187     def __getitem__(self,i): return self.fs[i]

~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastcore/transform.py in compose_tfms(x, tfms, is_enc, reverse, **kwargs)
    136     for f in tfms:
    137         if not is_enc: f = f.decode
--> 138         x = f(x, **kwargs)
    139     return x
    140 

~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastcore/transform.py in __call__(self, x, **kwargs)
    102     def __call__(self, x, **kwargs):
    103         if not _is_tuple(x): return super().__call__(x, **kwargs)
--> 104         return retain_type(super().__call__(list(x), **kwargs), x)
    105 
    106     def decode(self, x, **kwargs):

~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastcore/transform.py in __call__(self, x, **kwargs)
     70     @property
     71     def name(self): return getattr(self, '_name', _get_name(self))
---> 72     def __call__(self, x, **kwargs): return self._call('encodes', x, **kwargs)
     73     def decode  (self, x, **kwargs): return self._call('decodes', x, **kwargs)
     74     def __repr__(self): return f'{self.name}: {self.encodes} {self.decodes}'

~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastcore/transform.py in _call(self, fn, x, split_idx, **kwargs)
     80     def _call(self, fn, x, split_idx=None, **kwargs):
     81         if split_idx!=self.split_idx and self.split_idx is not None: return x
---> 82         return self._do_call(getattr(self, fn), x, **kwargs)
     83 
     84     def _do_call(self, f, x, **kwargs):

~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastcore/transform.py in _do_call(self, f, x, **kwargs)
     84     def _do_call(self, f, x, **kwargs):
     85         if not _is_tuple(x):
---> 86             return x if f is None else retain_type(f(x, **kwargs), x, f.returns_none(x))
     87         res = tuple(self._do_call(f, x_, **kwargs) for x_ in x)
     88         return retain_type(res, x)

~/anaconda3/envs/segmentation/lib/python3.7/site-packages/fastcore/dispatch.py in __call__(self, *args, **kwargs)
     96         if not f: return args[0]
     97         if self.inst is not None: f = MethodType(f, self.inst)
---> 98         return f(*args, **kwargs)
     99 
    100     def __get__(self, inst, owner):

<ipython-input-149-dc21e5018226> in encodes(self, x)
     10 
     11         # apply transform
---> 12         transformed = self.transform(image=np_image,mask=np_mask)
     13         tensor_image = transformed["image"]
     14 

~/anaconda3/envs/segmentation/lib/python3.7/site-packages/albumentations/core/transforms_interface.py in __call__(self, force_apply, **kwargs)
     85                     )
     86                 kwargs[self.save_key][id(self)] = deepcopy(params)
---> 87             return self.apply_with_params(params, **kwargs)
     88 
     89         return kwargs

~/anaconda3/envs/segmentation/lib/python3.7/site-packages/albumentations/core/transforms_interface.py in apply_with_params(self, params, force_apply, **kwargs)
     98                 target_function = self._get_target_function(key)
     99                 target_dependencies = {k: kwargs[k] for k in self.target_dependence.get(key, [])}
--> 100                 res[key] = target_function(arg, **dict(params, **target_dependencies))
    101             else:
    102                 res[key] = None

~/anaconda3/envs/segmentation/lib/python3.7/site-packages/albumentations/augmentations/transforms.py in apply(self, img, random_state, interpolation, **params)
   1301             self.value,
   1302             np.random.RandomState(random_state),
-> 1303             self.approximate,
   1304         )
   1305 

~/anaconda3/envs/segmentation/lib/python3.7/site-packages/albumentations/augmentations/functional.py in wrapped_function(img, *args, **kwargs)
     52     def wrapped_function(img, *args, **kwargs):
     53         shape = img.shape
---> 54         result = func(img, *args, **kwargs)
     55         result = result.reshape(shape)
     56         return result

~/anaconda3/envs/segmentation/lib/python3.7/site-packages/albumentations/augmentations/functional.py in elastic_transform(img, alpha, sigma, alpha_affine, interpolation, border_mode, value, random_state, approximate)
   1140         cv2.warpAffine, M=matrix, dsize=(width, height), flags=interpolation, borderMode=border_mode, borderValue=value
   1141     )
-> 1142     img = warp_fn(img)
   1143 
   1144     if approximate:

~/anaconda3/envs/segmentation/lib/python3.7/site-packages/albumentations/augmentations/functional.py in __process_fn(img)
    178             img = np.dstack(chunks)
    179         else:
--> 180             img = process_fn(img, **kwargs)
    181         return img
    182 

TypeError: Expected Ptr<cv::UMat> for argument 'src'

As indicated by the line

Pipeline: AddMaskCodes -> FastAIElasticTransform -> FastAIGridDistortion -> FastAIOpticalDistortion -> Resize -> ToTensor

you are applying those transforms before ToTensor is applied, so you should use a version that will expect PIL images, not tensors. You should also try your transform individually on a tuple (image, mask) to make sure they actually work before putting them in a Pipeline.

1 Like

Yes, if you look at it again, you can see that these transforms are working with PILImages

Maybe the names of variables are not descriptive

@sgugger I have managed to make them work. However, they are very slow?

Has fastai2.vision any transform like ElasticTransform, GridDistortion and OpticalDistortion from Albumentations?

Probably because they’re being done on the CPU. If you can get it working as a batch transform you can utilize the GPU instead. Also Sylvain recently ported over initial work for Albumentions (but they’re not done on the GPU)

See here:

Yes, this operations are done in CPU and I need to convert them to PILIMAGES, then to tensor and back to cuda

I will look into it!!!

No albumentations does not use the GPU for the data augmentation, it won’t work on a batch of images.

1 Like

I have managed to make them work as batch transform @sgugger :

def tensor2np(x):
    np_image = x.cpu().numpy()
    np_image = (np_image * 255).astype(np.uint8)
    
    return np_image

class FastAIElasticTransform(ItemTransform):
    def __init__(self,p):
        self.probability=p
        self.transform=ElasticTransform(p=self.probability)
        
    def encodes(self,x):
        images = x[0]
        masks = x[1]
        
        tensor_images = []
        tensor_masks = []
        for i in range(0,images.shape[0]):
            np_image = tensor2np(images[i])
            np_mask = tensor2np(masks[i])
            np_image = np_image
            np_mask = np_mask

            # apply transform
            transformed = self.transform(image=np_image,mask=np_mask)
            transformed = ToTensor()(image=transformed["image"],mask=transformed["mask"])

            # back to tensor
            tensor_image = transformed["image"]
            tensor_image = tensor_image.reshape(3,1002,1002).cuda()
            tensor_images.append(tensor_image)
            
            tensor_mask = transformed["mask"].long()
            tensor_mask = tensor_mask.reshape(1002,1002).cuda()
            tensor_masks.append(tensor_mask)
            
        
        return (torch.stack(tensor_images),torch.stack(tensor_masks))

However it runs very slow and show_batch not working, owing to the fact that returns normal tensor and not the ones from FastAI

Sorry I meant to seperate the two, batch transforms are done on the GPU but not those transforms :slight_smile:

Yes it runs in a for loop in python. Those are really slow.
Putting them as item_tfms would give you the speed of multiprocessing when assembling the batches.

1 Like

I am looking at your tutorial of using Albumentations an your code looks better than mine :sweat_smile:

However, I would like them to use the same probability as the probability cascade of aug_transforms