Overlap two tensors of different size based on an offset in PyTorch (or how to do PIL.paste but with tensors)

nintendoguy: <p>I have the following structure:</p>
<pre><code>torch.Size([channels, width, height])
</code></pre>
<p>Let’s say I have a tensor <code>a</code></p>
<pre><code>torch.Size([4, 512, 512])
</code></pre>
<p>And tensor <code>b</code></p>
<pre><code>torch.Size([4, 100, 100])
</code></pre>
<p>What I would like to do is to create a tensor <code>c</code> that is the result of “placing” tensor <code>b</code> on an arbitrary (width, height) coordinate offset of tensor <code>a</code>. For example, let’s say I would like to place tensor <code>b</code> on (300,100) of tensor <code>a</code></p>
<p>So for tensor <code>a</code>'s width between the 300-400 position, the values on tensor <code>a</code> should be replaced by the 100 values of tensor <code>b</code> width.</p>
<p>For tensor <code>a</code>'s height between the 100-200 position, the values of tensor <code>a</code> should be replaced by the 100 values of tensor <code>b</code> height.</p>
<p>I would also like to choose for which channels I want to do this substitution and for which channels I would keep tensor <code>a</code>'s value</p>
<p>(PS: The image is just an easy to illustrate example, but I would like to do it in a more generalisable way, so I’m no interested in converting to PIL, using <code>PIL.paste</code> and back to tensor but I would like to do all operations directly with Tensors)</p>

Matias_Vasquez: <p>You could access those indices directly and change the value.</p>
<pre data-code-wrap="python"><code class="lang-python"># Size of new_values have to match the patch selected from 'a'
a[channels, y1:y2, x1:x2] = new_values 
</code></pre>
<p>For something more general, you can do something like this ↓</p>
<pre data-code-wrap="python"><code class="lang-python">def paste(a, b, x1, y1, channels=None):
    assert len(a.shape) &gt;= 3, f"Expected at least 3 dimensions for tensor 'a', got {len(a.shape)}"
    assert len(b.shape) == 3 or len(b.shape) == 2, f"Expected [2 or 3] dimensions for tensor 'b', got {len(a.shape)}"

    channels = channels if not channels is None else list(range(a.shape[-3]))
    if len(b.shape) == 3:
        assert a.shape[-3]==b.shape[-3] or len(channels)==b.shape[-3] or b.shape[-3]==1, "tensors a and b must have the same number of channels or 'b' 1."


    # Patch size
    _h, _w = b.shape[-2], b.shape[-1]

    h = _h if y1 + _h &lt; a.shape[-2] else a.shape[-2] - y1
    w = _w if x1 + _w &lt; a.shape[-1] else a.shape[-1] - x1

    if len(b.shape) == 3  and a.shape[-3] == b.shape[-3]:
        a[..., channels, y1:y1+h, x1:x1+w] = b[channels, :h, :w]  
    else:
        a[..., channels, y1:y1+h, x1:x1+w] = b[..., :h, :w]

    return a
</code></pre>
<p>Here you can give an image (or a batch of images) and the patch <code>b</code> to be pasted into the image(s).<br>
It does not matter if the patch is too big and this should take car of it.</p>
<p>Also, you can select which channels you want to apply this to, else it will be applied to every channel.</p>
<p><code>b</code> can have the same number of channels as <code>a</code> or as the <code>channels mask</code> or just be <code>1</code> and it will be cast to every channel.</p>
<h2><a name="testing-1" class="anchor" href="#testing-1"></a>Testing</h2>
<pre data-code-wrap="python"><code class="lang-python">from PIL import Image
import requests
from io import BytesIO
import torch

url = "https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcQiFzoiIq7R99280lBVEcLAHRpb3LTz0ybB0A&amp;usqp=CAU"
response = requests.get(url)
img = Image.open(BytesIO(response.content))

display(img)

img_t = torchvision.transforms.ToTensor()(img)

x1, y1 = 150, 65
b = torch.ones(3, 30, 80)
mask = [1]

img_p = paste(img_t, b, x1, y1, channels=mask)
display(torchvision.transforms.ToPILImage()(img_p))
</code></pre>
<ul>
<li>
<p>Image <strong>before</strong> patching<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/0/b/0b5f22f4ef25900890613eccbdfffd2e306891bb.png" alt="image" data-base62-sha1="1CB52upodG4KUsOoj8f8Az8KN5h" width="311" height="162"></p>
</li>
<li>
<p>Image <strong>after</strong> patching<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/8/a/8aa90229ac0a3c7525fd70cc6eea4daf7df0ebd0.png" alt="image" data-base62-sha1="jMDZmZAJo1aTzTKHUhsMrj8ucfe" width="311" height="162"></p>
</li>
</ul>
<p>If this does not solve your use-case please let me know.</p>