Modo vettorializzato per applicare una maschera tridimensionale a RGB in pytorch

Aug 18 2020

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

4 Mercury Aug 18 2020 at 18:31

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).