Mask is not being transformed (only image)

[Geethen]

I am trying to apply the flipnslide augmentation, however, the image only get transformed. I have tried passing the inputs as a dict and as separate args but the output mask remains the same shape as the input. Instead I am expecting the output tensor to have a shape of (58, 1, 64,64). See reproducible example below.

import torch
from flipnslide.tiling import FlipnSlide
from kornia.augmentation.container import AugmentationSequential
from kornia.constants import DataKey


class FlipnSlideAug(torch.nn.Module):
    def __init__(self, tilesize: int, viz: bool = False):
        super().__init__()
        self.tilesize = tilesize
        self.viz = viz
    
    def forward(self, tensor):
        # Process the input tensor directly
        return self.process_tensor(tensor)
    
    def process_tensor(self, tensor):
        # Ensure input is on CPU and convert to numpy
        np_array = tensor.squeeze(0).cpu().numpy()
        
        sample_tiled = FlipnSlide(
            tile_size=self.tilesize, 
            data_type='tensor',
            save=False, 
            image=np_array,
            viz=self.viz
        )
        
        # Return the tiles directly
        return sample_tiled.tiles

# Usage with AugmentationSequential
flipnslide = FlipnSlideAug(tilesize=64, viz=False)
tfms = AugmentationSequential(
    flipnslide,
    data_keys=[DataKey.INPUT, DataKey.MASK]# Specify both keys
)

# Example usage
train_batch = {
    'image': torch.rand(1, 3, 256, 256),
    'mask': torch.rand(1, 1, 256, 256)
}

transformed = tfms(train_batch['image'], train_batch['mask'])

print("Transformed image shape:", transformed[0].shape)
print("Transformed mask shape:", transformed[1].shape)

OUTPUT
Transformed image shape: torch.Size([58, 3, 64, 64])
Transformed mask shape: torch.Size([1, 1, 256, 256])

[adamjstewart]

This might be a better question to ask in the Kornia GitHub, but I'm guessing you need to subclass kornia.augmentation.GeometricAugmentationBase2D instead of nn.Module. See the docs for example custom augmentations. Not sure if there is a more detailed tutorial somewhere @johnnv1

[johnnv1]

Yes, it should subclass GeometricAugmentationBase2D -- and indeed we need a better/complete tutorial around this one!

[isaaccorley]

@Geethen You can see some examples of our transforms that subclass this here https://github.com/microsoft/torchgeo/blob/main/torchgeo/transforms/transforms.py

[edgarriba]

@lappemic maybe this is a new tutorial to add in the backlog list

[Geethen]

Thanks to you all @johnnv1 @adamjstewart @isaaccorley I have got it to work. If anyone is interested in implementing this see code below for implementation.

import torch
from flipnslide.tiling import FlipnSlide
from typing import Any
from torch import Tensor
import kornia as K

class _FlipnSlide(K.augmentation.GeometricAugmentationBase2D):
    """Flip and slide a tensor."""

    def __init__(self, tilesize: int, viz: bool = False) -> None:
        """Initialize a new _FlipnSlide instance.

        Args:
            tilesize: desired tile size
            viz: visualization flag
        """
        super().__init__(same_on_batch = True )
        self.flags = {'tilesize': tilesize, 'viz': viz}

    def compute_transformation(
        self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any]
    ) -> Tensor:
        """Compute the transformation.

        Args:
            input: the input tensor
            params: generated parameters
            flags: static parameters

        Returns:
            the transformation
        """
        
        out: Tensor = self.identity_matrix(input)
        return out

    def apply_transform(
        self,
        input: Tensor,
        params: dict[str, Tensor],
        flags: dict[str, Any],
        transform: Tensor = None,
    ) -> Tensor:
        """Apply the transform.

        Args:
            input: the input tensor
            params: generated parameters
            flags: static parameters
            transform: the geometric transformation tensor

        Returns:
            the augmented input
        """
        np_array = input.squeeze(0).cpu().numpy()
        sample_tiled = FlipnSlide(
            tile_size=flags['tilesize'], 
            data_type='tensor',
            save=False, 
            image=np_array,
            viz=flags['viz']
        )
        return sample_tiled.tiles

# Usage with AugmentationSequential
flipnslide = _FlipnSlide(tilesize=64, viz=False)
tfms = AugmentationSequential(
    flipnslide,
    data_keys=None
)
# Example usage
train_batch = {
    'image': torch.rand(1, 3, 256, 256),
    'mask': torch.rand(1, 1, 256, 256)
}

transformed = tfms(train_batch)

print("Transformed image shape:", transformed['image'].shape)
print("Transformed mask shape:", transformed['mask'].shape)
print(transformed)