senyh / ucmt Goto Github PK

unfolds  = torch.nn.Unfold(kernel_size=(h, w), stride=s).to(device)
folds = torch.nn.Fold(output_size=(args.image_size, args.image_size), kernel_size=(h, w), stride=s).to(device)
x11 = unfolds(uncertainty_map11)  # B x C*kernel_size[0]*kernel_size[1] x L  8 256 256 
x11 = x11.view(B, 1, h, w, -1)  # B x C x h x w x L
x11_mean = torch.mean(x11, dim=(1, 2, 3))  # B x L
_, x11_max_index = torch.sort(x11_mean, dim=1, descending=True)  # B x L B x L
# for student 2
x22 = unfolds(uncertainty_map22)  # B x C*kernel_size[0]*kernel_size[1] x L
x22 = x22.view(B, 1, h, w, -1)  # B x C x h x w x L
x22_mean = torch.mean(x22, dim=(1, 2, 3))  # B x L
_, x22_max_index = torch.sort(x22_mean, dim=1, descending=True)  # B x L B x L
img_unfold = unfolds(imageA1).view(B, C, h, w, -1)  # B x C x h x w x L
lab_unfold = unfolds(label.float()).view(B, 1, h, w, -1)  # B x C x h x w x L
for i in range(B):## 对8张图片进行操作
    img_unfold[i, :, :, :, x11_max_index[i, :topk]] = img_unfold[i, :, :, :, x22_max_index[i, -topk:]]
    img_unfold[i, :, :, :, x22_max_index[i, :topk]] = img_unfold[i, :, :, :, x11_max_index[i, -topk:]]
    lab_unfold[i, :, :, :, x11_max_index[i, :topk]] = lab_unfold[i, :, :, :, x22_max_index[i, -topk:]]
    lab_unfold[i, :, :, :, x22_max_index[i, :topk]] = lab_unfold[i, :, :, :, x11_max_index[i, -topk:]]
image2 = folds(img_unfold.view(B, C*h*w, -1))
label2 = folds(lab_unfold.view(B, 1*h*w, -1))