Proposal for extending transforms

[fmassa]

Up to now, we compose different transforms via Compose transform, which works like nn.Sequential but for transforms. Each transform is applied independently to input and target.
While simple and efficient, there has been an increasing demand on extending the available transforms to accept both input and target, see #9, #115, #221 for some examples. This would allow performing the same random transformations in both the input and the target, for tasks such as semantic segmentation.

There are a few possible approaches and I'll summarize the ones that have been mentioned in here:

1. provide a set of random transforms working on pairs (or triplets, etc) of images, as proposed for example in [add] A PairRandomCrop for both input and target. #221. There are a couple of downsides of this approach: (a) it doesn't scale when we want to combine images coming from different domains (in object detection, we have image and bounding boxes for example, and we need to reimplement a new transform for each pair, introducing some redundancy in the code), (b) hard-coded to work with 2 inputs (but could be extended to n-inputs without difficulties)
2. factor out the randomness in the transforms, so that we can use the same functionality we currently have, just needing to add a generate_seed call in __getitem__, as was proposed in Separate random generation from transforms #115 . The drawback from this approach is that in some cases, we need the information of input to be able to perform the transformation in target (imagine flipping a bounding box horizontally, it requires the image width, which is not available right away).
3. provide a generic base class that samples the random transformation parameters, and passes those parameters to the input arguments, as proposed by @chsasank in the slack channel. While this is an improvement over both [add] A PairRandomCrop for both input and target. #221 and Separate random generation from transforms #115, it still suffers from the same limitation of Separate random generation from transforms #115 as we still only pass one input argument at a time.

Also, all those options actually handle each input independently. Furthermore, for both 1. and 3., the order of the operations is fixed in the dataset class (first input transforms, then target transforms, then joint transforms, for example). There might be cases where it would be convenient to not be restricted to such orderings, and alternate between single transform and joint transform.

One possibility to address those issues would be to provide a set of functions like nn.Split, nn.Concat, nn.Select, etc, so that we can always pass all inputs to the transforms, and let the transforms be implemented by the user explicitly. This mimics the legacy nn behavior.
The downside of this approach is that it gets very complicated to write some transformations, and this doesn't buy us much.

Instead, one simpler approach (which has already been advocated in the past by @colesbury) is to let the user directly subclass the dataset and implement their complex transforms in there. This approach (4.) would look something like

class VOCDatasetSegmentation(VOCDataset):
    def __init__(self, flip=False, **kwargs):
        super(VOCDatasetSegmentation, self).__init__(**kwargs)
        self.flip = flip

    def __getitem__(self, idx):
        image, target = super(VOCDatasetSegmentation, self).__getitem__(idx)
        do_flip = np.random.random() > 0.5
        if self.flip and do_flip:
            # flip image and bbox here
        return image, target

A downside is that we can't easily re-use those transforms in a different Dataset class (COCO for example).

The question now is if there is an intermediate design that we could leverage that keeps the simplicity of 4., without having to subclass the dataset and reimplement every time the same transforms?

What about yet another possibility (5.) would be to let the user write the code of their as follows

class Dataset(object):
    def __init__(self, transforms=None):
        self.transforms = transforms
    def __getitem__(self, idx):
        # get image1, image2, bounding_box
        # the transforms takes all inputs into account
        if self.transforms:
            image1, image2, bounding_box = self.transforms(image1, image2, bounding_box)
        return image1, image2, bounding_box

from torchvision.transforms import random_horizontal_flip

class MyJointRandomFlipTransform(object):
    def __call__(self, image1, image2, bounding_box):
        # provide a functional interface for the current transforms
        # so that they can be easily reused, and have the parameters
        # of the transformation if needed
        image1, params = random_horizontal_flip(image1, return_params=True)
        # reuses the same transformations, if wanted
        image2 = random_horizontal_flip(image2, params=params)
        # no transformation in torchvision for bounding_box, have to do it
        # ourselves
        if params.flip:
            bounding_box[:, 1] = image1.size(2) - bounding_box[:, 1]
            bounding_box[:, 3] = image1.size(2) - bounding_box[:, 3]
        return image1, image2, bounding_box

In this way, we have the flexibility subclassing the dataset, while being more modular and easy to implement.
There would be some differences in the way we write our datasets currently, but we could have a fallback implementation for backward compatibility, in the lines of

class StandardTransform(object):
    def __init__(self, transform, target_transform):
        self.transform = transform
        self.target_transform = target_transform

   def __call__(self, input, target):
        if self.transform:
            input = self.transform(input)
        if self.target_transform:
            target = self.target_transform(target)
        return input, target

and we would replace in the current datasets transform and target_transform by a single transforms, while keeping the old behavior

class Dataset(object):
    def __init__(self, path, transforms=None, transform=None, target_transform=None):
        # assert that only transforms or (transform, target_transform) can be set at a time
        if transforms is None:
           transforms = StandardTransform(transform, target_transform)
        self.transforms = transforms

    # getitem only uses transforms from now on

What do you think? Do you see drawbacks on using such an approach?

[chsasank]

Wow, really nice writeup summarising possible approaches.

I think this is a cool idea! Simple to understand.

One drawback I can think of is how params are handled. Transforms usually have two kinda of params: some fixed and specified when transform is inited(like output size in RandomCrop) and some random (like top left corner of the crop). Which of them are returned by random_horizontal_flip(image1, return_params=True)? (I assume random_horizontal_flip is a instance of RandomHorizontalFlip). This can make things confusing. But I guess, we can easily work around this.

Another drawback with this is that we cannot have default transforms support bounding boxes. Although we can hack around it, it will be with 'gotchas', which is not cool.

Here's one solution: How about we separate param generation and transform, like in your proposal but the implementation is done in class as methods. We can then allow user to subclass the transform to do whatever the user wants. Let me illustrate with RandomCrop example:

class RandomCrop(object):
    def __init__(self, size):
        self.size = size

    def get_params(self, img):
        w, h = img.size
        th, tw = self.size
        x1 = random.randint(0, w - tw)
        y1 = random.randint(0, h - th)
        return x1, y1, x1 + th, y1 + tw

    @staticmethod
    def transform(img, x1, y1, x2, y2):
        # no self here, all required params will have to specified in params
        return img.crop((x1, y1, x2, y2))


    def __call__(self, img):

        params = self.get_params(img)
        img = self.transform(img, *params)
        return img

Assumptions is that there will always be an image to transform. Now Object detection subclass can be like.

class MyRandomCrop(RandomCrop):
    def __call__(self, img, bbox):

        x1, y1, x2, w2 = self.get_params(img)

        img = self.transform(img, x1, y1, x2, w2)
        # convention: bbox = (x, y, w, h)
        bbox[0] = bbox[0] - x1
        bbox[1] = bbox[1] - y1

        return img, bbox

Disadvantage to this solution is increased complexity. But then we can have a abstract class to explain about get_params().

[fmassa]

@chsasank I like you proposal, makes it simpler to write new transforms without having to mess up with the functional interface of the transforms that I proposed, and the return_params seemed quite ugly as well :)

Before we move forward, there is another point that I'd like to raise.
With the basic fallback implementation that I mentioned in the original post (StandardTransform), we can have the sequential and individual operations that are currently supported by the datasets.

For more complex cases where we want to apply joint transforms but also individual transforms (for example in segmentation, we perform colour augmentation in the input image, and random crops / flips in the input and target), I think the way to go would be to encourage the user to write their own Transform class, which leverages other Transform classes, and implement the __call__ that chains the operations together, like in pytorch, instead of trying to use the Compose transform to write complicated transformations. Something like

class MySegmentationTransform(object):

    def __init__(self, input_color_transform,
                 joint_random_crop_transform, joint_random_flip_transform):
        self.input_color_transform = input_color_transform
        self.joint_random_crop_transform = joint_random_crop_transform
        self.joint_random_flip_transform = joint_random_flip_transform

    def __call__(self, input, target):
        input, target = self.joint_random_crop_transform(input, target)
        input, target = self.joint_random_flip_transform(input, target)
        input = convert_pil_to_tensor(input)
        target = convert_pil_to_tensor(target)
        input = self.input_color_transform(input) / 255
        return input, target

What do you think?

cc @alykhantejani for feedback as well.

[chsasank]

This is cool, but I liked Compose :).
We can kind of modify Compose to work with joint transforms like this:
(Observe how my design uses dicts for returns)

class Compose():
    def __call__(image, **kwargs):
        if len(kwargs) == 0:
            for t in self.transforms:
                image = t(image)

            return image
        else:
            kwargs['image'] = image
            for t in self.transforms:
                kwargs = t(**kwargs)

            return kwargs

Don't you think we should support segmentation and bounding boxes in this repo? In this design, it's not such a big deal. We will have methods transform_bboxes and transform_segmap for each transform.

Then call will be like:

 def __call__(self, image, **kwargs):
    params = self.get_params(image)
    image = self.transform(image, *params)

    if len(kwargs) > 0:
        for k, v in kwargs:
            if k.startswith('bboxes'):
                kwargks[k] = self.transform_bboxes(image, v, *params)
            elif k.startswith('segmap'):
                kwargs[k] = self.transform_segmap(image, v, *params)

        kwargs['image'] = image
        return kwargs
    else:
        return img

User can then use transforms like this:

transform = Compose([RandomCrop, HorizontalFlip])
sample = {'image': image, 'segmap': segmap, 'bboxes': bboxes}
transformed_sample = transform(**sample)

[fmassa]

I'm not super happy with the Compose that you propose. It seems a bit obscure and error-prone to me, but I'll let @alykhantejani give his view on it.
I don't think we should be providing anything for the moment that restricts/over-complexifies the transforms, so I'd restrain from adding specific support for bounding boxes.

Also, segmentation maps are images, right? So why should we handle it differently? The only reason I can see it is if you want to use the Compose you proposed, but for some transforms you don't want to apply them to the segmentation map (like color augmentation). This shouldn't be a problem if you just let the user write their Transforms themselves, right?

[chsasank]

Sure, it does look a bit complex. Let's leave it out. (Segmentation needs nearest neighbour interpolation too).

[alykhantejani]

Really good summary + ideas on this thread @fmassa and @chsasank

I like the idea of letting the user create transform graphs in any way they want, without overwhelming them with things like JoinTransform, Split and Sequential, so in this respect, I like the idea of the MySegmentationTransform object that @fmassa proposed.

I also like the idea of separating out the params for a transform from the implementation, but was not a fan of the return_params kwarg as, as @chsasank mentioned, it’s a bit opaque as to what is actually returned. (Users will have to continuously refer to docs or source code)

I’m not a huge fan of using dicts with string args for the inputs as again I think there is potential for errors and confusion on the users part.

I would like to propose something that is a hybrid of the proposals above, which is that we create functional forms of the transforms where each function explicitly requires all of the parameters necessary for the transform. i.e. the random_crop function will look like def random_crop(img, x, y, w, h). The users will then pass a single callable to the Dataset (which we can make backwards compatible in the way proposed by @fmassa) which handles the generating of their random params and the chaining of all of their transforms. So we would have something a bit like this:

import torchvision.transforms as F

def my_segmentation_trasform(input, target):
    x, y, w, h = get_random_crop_paras(input, target)
    input = F.random_crop(input, x, y, w, h)
    target = F.random_crop(target, x, y, w, h)
    if random.random() > 0.5:
        input = F.hflip(input)
        target = F.hflip(target)
    input = F.color_transform(input)
    return F.image_to_tensor(input), F.image_to_tensor(target)

This is similar to @chsasank’s idea of using static methods for the transform (which are functions anyway), without the OOP overhead.

what do you guys think?

[chsasank]

Great. Let's do this. I like this idea. We'll create 'functional' parts of the transforms like what you mentioned and use them in class implementations. Just like how pytorch deals with nn.functional and nn.

Shall I start working on this?

[chsasank]

I've made a first cut refactor at #240 :)
Have a look and let me know what you think.

[fmassa]

@chsasank Thanks, very efficient! Having a look at it now