Manière vectorisée d'appliquer un masque en 3 dimensions à RVB dans pytorch

Aug 18 2020

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

4 Mercury Aug 18 2020 at 18:31

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.