将3维蒙版应用于pytorch中的RGB的矢量化方法

问题描述

我有一个表示RGB图像的HxWx3张量和一个作为输入的HxWx3蒙版（布尔）张量。 假设对于掩码张量中的每个（i，j），都存在一个真实值（即R \ G \ B中的一个恰好处于打开状态）。 我想对图像应用遮罩，以生成HxW（或HxWx1）张量V，其中V [i，j] ='根据遮罩匹配的R \ G \ B值。

使用Problem applying binary mask to an RGB image with numpy，我可以实现以下目标：

>>> X*mask
tensor([[[ 9.,10.],[ 0.,0.]],[[ 0.,0.],20.]],[30.,0.]]])

但是如上所述，我想要一个昏暗的HxW，而不是HxWx3。

插图：

解决方法

假设对于每个i，j仅保留一个R / G / B值，您可以简单地执行以下操作：

(X*mask).sum(axis=2)

这应该为您提供所需的（HxW）输出。