Masked Sliding Window tensor

pvalois: <p>Hello all,</p>
<p>My problem should be simple, but most solutions I found use a for loop, which I think is making it too slow for my case.</p>
<p>Given an image batch, I want to create a masked batch which masks parts in a sliding window fashion. I tried it with unfold and fold, but I still cannot understand how to do it in a simple way.</p>
<p>For example, given an image batch of shape (16, 3, 224, 224) and a sliding window with kernel size 32 and stride 16. I want a new tensor with shape (16, 169, 3, 224, 224). Each of the 169 images having a different mask.</p>
<p>Could anybody help? Thanks!</p>

J_Johnson: <p>Can you write the complete code that gives the result you want as an iterable? You can substitute <code>img = torch.rand(16, 3, 224, 224)</code> as the input.</p>

pvalois: <aside class="quote no-group" data-username="J_Johnson" data-post="2" data-topic="180517">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/j_johnson/48/55494_2.png" class="avatar"> J_Johnson:</div>
<blockquote>
<p>img = torch.rand(16, 3, 224, 224)</p>
</blockquote>
</aside>
<p>Sure! I was able to rework the code a bit and there is only one for loop now, which I think there might be a way to remove.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch

class SlidingWindowMask(torch.nn.Module):

    def __init__(self, kernel_size: float, stride: float) -&gt; None:
        super().__init__()
        self.kernel_size = kernel_size
        self.stride = stride

    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        top, left, bottom, right = self.generate_coordinates(x)
        num_masks = len(top)

        x_masked = torch.stack([x] * num_masks, dim=0)
        for i in range(num_masks):
            x_masked[i, ..., left[i] : right[i], top[i] : bottom[i]] = 0

        return x_masked

    def get_dimensions(self, x: torch.Tensor) -&gt; tuple[int, int, int, int, int, int]:
        *_, height, width = x.size()

        mask_height = int(self.kernel_size * height)
        mask_width = int(self.kernel_size * width)

        stride_height = int(self.stride * height)
        stride_width = int(self.stride * width)

        return height, width, mask_height, mask_width, stride_height, stride_width
    
    def generate_coordinates(
        self,
        x: torch.Tensor,
    ) -&gt; torch.Tensor:
        h, w, m_h, m_w, s_h, s_w = self.get_dimensions(x)

        top, left = self.compute_top_left(h, w, m_h, m_w, s_h, s_w, device=x.device)
        
        bottom = top + m_h
        right = left + m_w

        return torch.stack([top, left, bottom, right])

    def compute_top_left(self, h, w, m_h, m_w, s_h, s_w, *args, **kwargs) -&gt; tuple[torch.Tensor, torch.Tensor]:
        top = torch.arange(0, w - m_w + 1, s_h, *args, **kwargs)
        left = torch.arange(0, h - m_h + 1, s_w, *args, **kwargs)

        grid_x, grid_y = torch.meshgrid(left, top, indexing="ij")

        return grid_y.flatten(), grid_x.flatten()
</code></pre>
<p>Below you have how I am testing it:</p>
<pre data-code-wrap="python"><code class="lang-python">from PIL import Image

from .masking.sliding_window import SlidingWindowMask

import torchvision.transforms.functional as TF

import torchvision

import seaborn_image as isns

# img = torch.rand(16, 3, 224, 224)

img = Image.open(ROOT / "files/images/turtle.png")

img = TF.to_tensor(img)

print(f"Image shape: {img.shape}")

masking = SlidingWindowMask(kernel_size=0.25, stride=.15)

x = masking(img)

print(f"Masked image shape: {x.shape}")

x = torchvision.transforms.functional.vflip(x)

img_grid = torchvision.utils.make_grid(x, nrow=4, normalize=True, pad_value=1)

img_grid = img_grid.permute(1, 2, 0)

isns.imshow(img_grid, showticks=False, despine=True, cbar=False)
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">Image shape: torch.Size([4, 319, 319])
Masked image shape: torch.Size([16, 4, 319, 319])
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/1/7/1732b85548fbe9e161b7ca129d6e761c24d55276.png" alt="sliding_window_mask" data-base62-sha1="3jdEeRP7klK7bqueRt5fTD5z1si" width="389" height="389"></p>

J_Johnson: <p>Here is one way you could do this, but just take note that the kernel_size and stride should be equal, or the F.fold operation winds up overwriting the mask, in most cases.</p>
<pre><code class="lang-auto">#define values
img = torch.rand(16, 4, 319, 319)
kernel_size=32
stride=32

assert kernel_size==stride, "stride and kernel_size must be equal"

#create mask
mask = torch.ones_like(img)
mask = F.unfold(mask, kernel_size=kernel_size, stride=stride) # batch_size, channels*kernel*kernel, patches
N, chw, p = mask.shape
mask = mask.unsqueeze(3).expand(N, chw, p, p).clone()

diag = torch.diag(torch.ones(p)).bool()
mask[:,:,diag]=0

N, c, h, w = img.shape
mask = mask.reshape(N, c, kernel_size, kernel_size, p, p)
mask = mask.permute(0,4,1,2,3,5).reshape(N*p*c, kernel_size*kernel_size, p)
mask = F.fold(mask, (h,w), kernel_size=kernel_size, stride=stride)
mask = mask.reshape(N,p, c, h, w)

#expand img
masked_img=img.unsqueeze(1).expand(N,p,c,h,w).clone()

#apply mask
masked_img[~mask.bool()]=0.
</code></pre>
<p>The above steps can also be simplified to substitute the <code>img</code> in, instead of using a mask. However, if you make the mask once(since it’s deterministic), you can store it in memory and apply it repeatedly on new batches of images without needing to recreate it, which will be faster if this is going to be repeated.</p>
<p>Also, on a side note, I noticed you’re loading a png file, which means one of the channels is an alpha channel(used for transparency in png files). Not sure if you intended to keep that, but if you would like that channel removed, see here:</p>
<aside class="quote" data-post="2" data-topic="28509">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/what-is-color-axis/28509/2">What is color axis?</a> <a class="badge-category__wrapper " href="/c/vision/5"><span data-category-id="5" data-drop-close="true" class="badge-category " title="Topics related to either pytorch/vision or vision research related topics"><span class="badge-category__name">vision</span></span></a>
  </div>
  <blockquote>
    The axis swap is most likely right. 
I guess you have 4 channels, because your image might have an additional alpha channel. 
Usually it’s the last channel, so you could just remove it with: 
image2 = image2[:3]

Try to slice the fourth channel and have a look at the values. Probably they are all ones.
  </blockquote>
</aside>