从图像张量复制片段

Question

我有三个张量： A - (1, 3, 256, 256) B - (1, 3, 256, 256) - 这是一个白色图像张量 C - (256, 256) - 这是段张量

例如 C 看起来像：

tensor([[ 337,  337,  337,  ...,  340,  340,  340],
        [ 337,  337,  337,  ...,  340,  340,  340],
        [ 337,  337,  337,  ...,  340,  340,  340],
        ...,
        [1022, 1022, 1022,  ..., 1010, 1010, 1010],
        [1022, 1022, 1022,  ..., 1010, 1010, 1010],
        [1022, 1022, 1022,  ..., 1010, 1010, 1010]], device='cuda:0')

其中 37 可以表示建筑物等

张量 C 给出了线段形状的位置。我想要的是根据从张量 A 到张量 B 的位置复制相同的片段。这将是将该片段 Photoshop 到白色图像张量上。

这类似于掩蔽，我研究了 mask_select (https://pytorch.org/docs/stable/generated/torch.masked_select.html) 但只有 returns 一维张量返回。

Answer 1

您不需要 select C 中的像素，只需屏蔽它们：

select = 337  # which segment to select
select_mask = (C == select)[None, None, ...]  # create binary mask and add singleton dimensions
# this is the part where you select the right part of A
B = B * (1 - select_mask) + A * select_mask

从图像张量复制片段

copy segment from image tensor

python

image

pytorch

tensor