Векторизованный способ применения 3-х мерной маски к RGB в pytorch

Question

У меня есть тензор HxWx3, представляющий изображение RGB, и тензор маски HxWx3 (логический) в качестве входных данных. Предполагается, что для каждого (i, j) в тензоре маски есть ровно одно истинное значение (то есть ровно одно из R \ G \ B включено). Я хочу применить маску к изображению, чтобы получить тензор V HxW (или HxWx1), где V [i, j] = 'соответствующее значение R \ G \ B в соответствии с маской'.

Используя проблему с применением двоичной маски к изображению RGB с помощью numpy, мне удалось добиться следующего:

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

        [[ 0.,  0.],
         [ 0., 20.]],

        [[ 0.,  0.],
         [30.,  0.]]])

Но, как уже говорилось, я хочу получить один тусклый HxW, а не HxWx3.

Иллюстрация:

Answer 1

Предполагая, что для каждого i, j сохраняется только одно значение R / G / B, вы можете просто сделать:

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

Это должно дать вам желаемый результат (HxW).