Modo vettorializzato per applicare una maschera tridimensionale a RGB in pytorch
Ho un tensore HxWx3 che rappresenta un'immagine RGB e un tensore maschera HxWx3 (booleano) come input. Si presume che per ogni (i,j) nel tensore della maschera ci sia esattamente un valore vero (che è esattamente uno di R\G\B è attivo). Voglio applicare la maschera all'immagine per ottenere un tensore HxW (o HxWx1) V dove V[i,j]='il valore R\G\B corrispondente in base alla maschera'.
Usando il problema dell'applicazione della maschera binaria a un'immagine RGB con numpy sono stato in grado di ottenere quanto segue:
>>> X*mask
tensor([[[ 9., 10.],
[ 0., 0.]],
[[ 0., 0.],
[ 0., 20.]],
[[ 0., 0.],
[30., 0.]]])
Ma come detto, voglio un singolo dim HxW e non HxWx3 come risultato.
Illustrazione:

Risposte
Supponendo che per ogni i,j venga mantenuto un solo valore R/G/B, puoi semplicemente fare:
(X*mask).sum(axis=2)
Questo dovrebbe darti l'output desiderato (HxW).