Hi,
I have segmentation maps with size = [5,64,64] where batch size is 5, the spatial size is 64x64.
Given a feature map F with the same size 64x64 and N channels.
Let select locations in the segmentation map whose labels are v. Now I want to select features in F based on the selected locations whose labels are v in the segmentation map.
Here’s how it looks like:
seg = torch.randint(0,5,size=(5,64,64))

F = torch.rand(5,25,64,64)

locations = seg==1
F[locations]??

Examples:
seg with batch_size=2, spatial_size=2x2=4
seg = [[[0,1,2,3],[1,2,1,2]
F with 3 channels, the size is (2,3,4)
F = [[2,3,4,3],
     [4,1,2,4],
     [4,0,5,1]],

     [1,2,0,4],
     [0,1,10,3],
     [1,2,0,5]],
location = seg == 1

output = F[location]= [[3,1,0],[1,0,1],[0,10,0]]