`torchvision` v2 transforms not applying to both image and mask

ayoubft: <p><code>torchvision.transforms.v2.Compose</code> with paired transforms like <code>RandomHorizontalFlip</code> and <code>RandomVerticalFlip</code> applies only to the image, not to the mask.</p>
<p>Expected behavior: when using paired transforms and passing a tuple <code>(image, mask)</code>, both should be flipped identically.</p>
<p>Observed: only the image is flipped; the mask is unchanged.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms.v2 import RandomHorizontalFlip, RandomVerticalFlip, Compose
from torchvision.transforms.v2.functional import pil_to_tensor
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw


def create_triangle_image(size=64):
    img = Image.new("RGB", (size, size), "black")
    draw = ImageDraw.Draw(img)
    draw.polygon([(63, 0), (63, 31), (32, 0)], fill="white")
    return img


class TriangleDataset(Dataset):
    def __init__(self, transform=None):
        pil_img = create_triangle_image()
        self.img = pil_to_tensor(pil_img)
        self.mask = self.img[0].unsqueeze(0)  # 1 channel mask
        self.transform = transform

    def __getitem__(self, idx):
        if self.transform:
            img_tr, mask_tr = self.transform(self.img, self.mask)
        return img_tr, mask_tr

    def __len__(self):
        return 1


transforms = Compose([
    RandomHorizontalFlip(p=.50),
    RandomVerticalFlip(p=.50),
])

dataset = TriangleDataset(transform=transforms)
img, mask = next(iter(DataLoader(dataset)))

# Plot
fig, axs = plt.subplots(1, 2, figsize=(6, 3))
axs[0].imshow(img[0].permute(1, 2, 0), cmap='gray')
axs[0].set_title("Image")
axs[1].imshow(mask[0][0], cmap='gray')
axs[1].set_title("Mask")
for ax in axs:
    ax.axis("off")
plt.tight_layout()
plt.show()
</code></pre>
<p>Results (multiple runs):</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/a/aa438f69720af1e7b444b27e12a7938bf08dba87.png" data-download-href="https://discuss.pytorch.org/uploads/default/aa438f69720af1e7b444b27e12a7938bf08dba87" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/a/aa438f69720af1e7b444b27e12a7938bf08dba87.png" alt="image" data-base62-sha1="oidUR72ZaAb8JCt7zNa8heYIrOv" width="571" height="299"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">571×299 1.75 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/8/682195e4f1aa88b4f203defa5e0ce975ae4241df.png" data-download-href="https://discuss.pytorch.org/uploads/default/682195e4f1aa88b4f203defa5e0ce975ae4241df" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/8/682195e4f1aa88b4f203defa5e0ce975ae4241df.png" alt="image" data-base62-sha1="eRbC2BDfhRDJ911rq80n7hMbH1J" width="571" height="299"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">571×299 1.75 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/6/463d5b51e8830aeb885073216f6a70d724d62c8c.png" data-download-href="https://discuss.pytorch.org/uploads/default/463d5b51e8830aeb885073216f6a70d724d62c8c" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/4/6/463d5b51e8830aeb885073216f6a70d724d62c8c.png" alt="image" data-base62-sha1="a1mSveSH2zILWYMp7VExkzxAIQQ" width="571" height="299"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">571×299 1.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>How to fix this?</p>

ayoubft: <p><a href="https://discuss.pytorch.org/t/torchvision-transfors-how-to-perform-identical-transform-on-both-image-and-target/10606/7">The functional API</a> works.</p>

ptrblck: <p>This is expected behavior in the <code>transforms.v2</code> API according to the <a href="https://docs.pytorch.org/vision/main/auto_examples/transforms/plot_transforms_getting_started.html#what-do-i-pass-as-input">docs</a>:</p>
<blockquote>
<p>If there is no <a href="https://docs.pytorch.org/vision/main/generated/torchvision.tv_tensors.Image.html#torchvision.tv_tensors.Image" title="torchvision.tv_tensors.Image"><code>Image</code></a> or <a href="https://docs.pytorch.org/vision/main/generated/torchvision.tv_tensors.Video.html#torchvision.tv_tensors.Video" title="torchvision.tv_tensors.Video"><code>Video</code></a> instance, only the first pure <a href="https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor" title="(in PyTorch v2.7)"><code>torch.Tensor</code></a> will be transformed as image or video, while all others will be passed-through. Here “first” means “first in a depth-wise traversal”.</p>
</blockquote>
<p>You could use the <code>tv_tensors</code> classes instead:</p>
<pre data-code-wrap="python"><code class="lang-python">img_tv = torchvision.tv_tensors.Image(img)
mask_tv = torchvision.tv_tensors.Mask(mask)

out1, out2 = transforms(img_tv, mask_tv)
</code></pre>