Pytorch에서 RGB에 3 차원 마스크를 적용하는 벡터화 된 방법

Question

RGB 이미지를 나타내는 HxWx3 텐서와 입력으로 HxWx3 마스크 (부울) 텐서가 있습니다. 마스크 텐서의 각 (i, j)에 대해 정확히 하나의 참 값이 있다고 가정합니다 (즉, R \ G \ B 중 정확히 하나가 켜져 있음). 이미지에 마스크를 적용하여 HxW (또는 HxWx1) 텐서 V를 생성합니다. 여기서 V [i, j] = '마스크에 따라 일치하는 R \ G \ B 값'이 생성됩니다.

numpy로 RGB 이미지에 바이너리 마스크를 적용하는 문제를 사용 하여 다음을 달성 할 수있었습니다.

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

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

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

그러나 언급했듯이 결과적으로 HxWx3이 아닌 단일 희미한 HxW를 원합니다.

삽화:

Answer 1

각 i, j에 대해 단일 R / G / B 값만 유지된다고 가정하면 다음을 수행 할 수 있습니다.

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

원하는 (HxW) 출력을 제공해야합니다.