Manière vectorisée d'appliquer un masque en 3 dimensions à RVB dans pytorch
J'ai un tenseur HxWx3 représentant une image RVB et un tenseur de masque HxWx3 (booléen) en entrée. On suppose que pour chaque (i,j) dans le tenseur de masque, il y a exactement une vraie valeur (c'est-à-dire qu'exactement l'un des R\G\B est activé). Je veux appliquer le masque à l'image pour obtenir un tenseur V HxW (ou HxWx1) où V[i,j]='la valeur R\G\B correspondante selon le masque'.
En utilisant le problème d'application d'un masque binaire à une image RVB avec numpy, j'ai pu obtenir ce qui suit :
>>> X*mask
tensor([[[ 9., 10.],
[ 0., 0.]],
[[ 0., 0.],
[ 0., 20.]],
[[ 0., 0.],
[30., 0.]]])
Mais comme indiqué, je veux un seul dim HxW et non HxWx3 comme résultat.
Illustration:

Réponses
En supposant que pour chaque i,j une seule valeur R/G/B est retenue, vous pouvez simplement faire :
(X*mask).sum(axis=2)
Cela devrait vous donner la sortie (HxL) souhaitée.