Input: -img (N, C, H, W) -shiftx, shifty (N, c, H, W)
(self, img, shiftx, shifty, weight)
| 72 | |
| 73 | |
| 74 | def sample_one(self, img, shiftx, shifty, weight): |
| 75 | """ |
| 76 | Input: |
| 77 | -img (N, C, H, W) |
| 78 | -shiftx, shifty (N, c, H, W) |
| 79 | """ |
| 80 | |
| 81 | N, C, H, W = img.size() |
| 82 | |
| 83 | # flatten all (all restored as Tensors) |
| 84 | flat_shiftx = shiftx.view(-1) |
| 85 | flat_shifty = shifty.view(-1) |
| 86 | flat_basex = torch.arange(0, H).view(-1, 1)[None, None].cuda().long().repeat(N, C, 1, W).view(-1) |
| 87 | flat_basey = torch.arange(0, W).view(1, -1)[None, None].cuda().long().repeat(N, C, H, 1).view(-1) |
| 88 | flat_weight = weight.view(-1) |
| 89 | flat_img = img.contiguous().view(-1) |
| 90 | |
| 91 | # The corresponding positions in I1 |
| 92 | idxn = torch.arange(0, N).view(N, 1, 1, 1).long().cuda().repeat(1, C, H, W).view(-1) |
| 93 | idxc = torch.arange(0, C).view(1, C, 1, 1).long().cuda().repeat(N, 1, H, W).view(-1) |
| 94 | # ttype = flat_basex.type() |
| 95 | idxx = flat_shiftx.long() + flat_basex |
| 96 | idxy = flat_shifty.long() + flat_basey |
| 97 | |
| 98 | |
| 99 | # recording the inside part the shifted |
| 100 | mask = idxx.ge(0) & idxx.lt(H) & idxy.ge(0) & idxy.lt(W) |
| 101 | |
| 102 | # Mask off points out of boundaries |
| 103 | ids = (idxn*C*H*W + idxc*H*W + idxx*W + idxy) |
| 104 | ids_mask = torch.masked_select(ids, mask).clone().cuda() |
| 105 | |
| 106 | #(zero part - gt) -> difference |
| 107 | # difference back propagate -> No influence! Whether we do need mask? mask? |
| 108 | # put (add) them together |
| 109 | # Note here! accmulate fla must be true for proper bp |
| 110 | img_warp = torch.zeros([N*C*H*W, ]).cuda() |
| 111 | img_warp.put_(ids_mask, torch.masked_select(flat_img*flat_weight, mask), accumulate=True) |
| 112 | |
| 113 | one_warp = torch.zeros([N*C*H*W, ]).cuda() |
| 114 | one_warp.put_(ids_mask, torch.masked_select(flat_weight, mask), accumulate=True) |
| 115 | |
| 116 | |
| 117 | |
| 118 | return img_warp.view(N, C, H, W), one_warp.view(N, C, H, W) |
| 119 | |
| 120 |