如何在 PyTorch 中对图像张量应用遮罩？

Question

使用 NumPy 或 OpenCV 应用掩码是一个相对简单的过程。但是，如果我需要在优化算法的损失计算中使用蒙版图像，我需要专门使用 PyTorch，否则会干扰梯度计算。 假设我有一个图像张量[1, 512, 512, 3]（批次、高度、宽度、通道）和一个掩码张量[1, 20, 512, 512]（批次、通道、高度、宽度），其中每个通道对应于 20 个分割类之一，我想得到用黑色 (0, 0, 0) 填充每个像素的蒙版图像张量，但属于一个或多个指定分割类别的除外。

这里是如何使用 numpy 完成的：

import numpy as np
import torch

# Create dummy image and mask
image_tensor = torch.randn([1, 512, 512, 3])
mask_tensor = torch.randn([1, 20, 512, 512])

# Apply argmax to mask
mask_tensor =  torch.max(mask_tensor, 1)[1] # -> 1, 512, 512

# Define mask function
def selective_mask(image_src, mask, dims=[]):
    h, w = mask.shape
    background = np.zeros([h, w, 3], dtype=np.uint8)
    for j_, j in enumerate(mask[:, :]):
        for k_, k in enumerate(j):
            if k in dims:
                background[j_, k_] = image_src[j_, k_]
    output = background
    return output

# Convert tensors to numpy:
image = image_tensor.squeeze(0).cpu().numpy()
mask = mask_tensor.squeeze(0).cpu().nmpy()

# Apply mask function for several classes
image_masked = selective_mask(image, mask, dims=[5, 6, 8])

应如何更改我的代码以使其符合 PyTorch 要求？

Answer 1

首先，函数selective_mask 的定义远非您所说的“直截了​​当”。使用numpy（和torch，它被设计为最兼容）的关键是利用操作的向量化并避免使用不可并行化的循环。

如果你以这种方式重写上述函数：

def selective_mask(image_src, mask, channels=[]):
    mask = mask[np.array(channels).astype(int)]
    return np.sign(np.sum(mask, axis=0), dtype=image_src.dtype) * image_src

事实证明，你实际上可以对 pytorch 张量做同样的事情（这里不需要挤压批次（第一）维度）：

def selective_mask_t(image_src, mask, channels=[]):
    mask = mask[:, torch.tensor(channels).long()]
    mask = torch.sgn(torch.sum(mask, dim=1)).to(dtype=image_src.dtype).unsqueeze(-1)
    return mask * image_src

另外，您可能希望以这种方式制作面具本身：
（顺便说一句，在这里使用 max 和 sgn 的组合实际上应该比设置由 argmax 索引的元素更快）

# Create dummy image and mask
image_tensor = torch.randn([1, 512, 512, 3])
mask_tensor = torch.randn([1, 20, 512, 512])

# Discreticize the mask (set to one in the channel with the highest value) -> 1, 20, 512, 512
mask_tensor = torch.sgn(mask_tensor - torch.max(mask_tensor, 1)[0].unsqueeze(1)) + 1.  

那么它应该可以正常工作：

print(selective_mask_t(image_tensor, mask_tensor, [5, 6, 8]))