Доброго вечера. Есть вопрос по наложению маски в торче.
Допустим, у меня есть трехканальное RGB изображение-тензор [1, 3, 512, 512] (B, C, H, W) и обученная модель для сегментации, которая дает на выходе тензор [1, 20, 512, 51] (B, C, H, W). Каждый из 20 каналов C - сегментация того или иного класса. Как можно оставаясь в пределах торча применить к исходному изображению маску сегментации, оставив на изображении только регионы одного или нескольких желаемых классов (а все остальное закрасить черным, (0, 0, 0))?
Используя numpy и opencv решить эту задачу просто, но тогда теряются градиенты торча. А мне нужно изображение с наложенной маской использовать для расчетов в алгоритме оптимизации.