Spaces:
Running
on
T4
Running
on
T4
| import torch | |
| def flow_prop(feat, flow, mode='forward'): | |
| """ | |
| Args: | |
| feat: features to be aligned | |
| flow: the filled current flow | |
| mode: `forward` or `backward`, indicates the propagation direction | |
| Returns: feature after warping | |
| """ | |
| assert mode in ['forward', 'backward'], 'Invalid mode: {}'.format(mode) | |
| feat = warp(feat, flow, mode) | |
| return feat | |
| def warp(feat, flow, mode): | |
| device = feat.device | |
| c = feat.shape[1] | |
| y = flow[:, 0:1, :, :] | |
| x = flow[:, 1:2, :, :] | |
| x = x.repeat(1, c, 1, 1) # [b, c, h, w] | |
| y = y.repeat(1, c, 1, 1) | |
| x1 = torch.floor(x) | |
| x2 = x1 + 1 | |
| y1 = torch.floor(y) | |
| y2 = y1 + 1 | |
| w11, w12, w21, w22 = get_gaussian_weights(x, y, x1, y1, x2, y2) | |
| feat11, o11 = sample_one(feat, x1, y1, w11, mode) | |
| feat12, o12 = sample_one(feat, x1, y2, w12, mode) | |
| feat21, o21 = sample_one(feat, x2, y1, w21, mode) | |
| feat22, o22 = sample_one(feat, x2, y2, w22, mode) | |
| feat_o = feat11 + feat12 + feat21 + feat22 | |
| o = o11 + o12 + o21 + o22 | |
| feat_o[o > 0] = feat_o[o > 0] / o[o > 0] | |
| return feat_o | |
| def sample_one(feat, shiftx, shifty, weight, mode): | |
| device = feat.device | |
| b, c, h, w = feat.shape | |
| flat_shiftx = shiftx.view(-1) # [b * c * h * w] | |
| flat_shifty = shifty.view(-1) | |
| flat_basex = torch.arange(0, h, requires_grad=False).view(-1, 1).long().repeat(b, c, 1, w).view(-1) | |
| flat_basey = torch.arange(0, w, requires_grad=False).view(-1, 1).long().repeat(b, c, h, 1).view(-1) | |
| flat_basex = flat_basex.to(device) | |
| flat_basey = flat_basey.to(device) | |
| flat_weight = weight.reshape(-1) | |
| flat_feat = feat.reshape(-1) | |
| idxn = torch.arange(0, b, requires_grad=False).view(b, 1, 1, 1).long().repeat(1, c, h, w).view(-1) | |
| idxc = torch.arange(0, c, requires_grad=False).view(1, c, 1, 1).long().repeat(b, 1, h, w).view(-1) | |
| idxn = idxn.to(device) | |
| idxc = idxc.to(device) | |
| if mode == 'forward': | |
| idxx = flat_shiftx.long() + flat_basex # size [-1] | |
| idxy = flat_shifty.long() + flat_basey # size [-1] | |
| else: # backward propagation | |
| idxx = -flat_shiftx.long() + flat_basex # size [-1] | |
| idxy = -flat_shifty.long() + flat_basey # size [-1] | |
| # record the shifted pixels inside the image boundaries | |
| mask = idxx.ge(0) & idxx.lt(h) & idxy.ge(0) & idxy.lt(w) | |
| # mask off points out of boundaries | |
| ids = idxn * c * h * w + idxc * h * w + idxx * w + idxy | |
| ids_mask = torch.masked_select(ids, mask).clone() | |
| # put the value into corresponding regions | |
| feat_warp = torch.zeros([b * c * h * w]) | |
| feat_warp = feat_warp.to(device) | |
| feat_warp.put_(ids_mask, torch.masked_select(flat_feat * flat_weight, mask), accumulate=True) | |
| one_warp = torch.zeros([b * c * h * w]) | |
| one_warp = one_warp.to(device) | |
| one_warp.put_(ids_mask, torch.masked_select(flat_weight, mask), accumulate=True) | |
| return feat_warp.view(b, c, h, w), one_warp.view(b, c, h, w) | |
| def get_gaussian_weights(x, y, x1, y1, x2, y2): | |
| sigma = 1 | |
| w11 = torch.exp(-((x - x1) ** 2 + (y - y1) ** 2) / (sigma ** 2)) | |
| w12 = torch.exp(-((x - x1) ** 2 + (y - y2) ** 2) / (sigma ** 2)) | |
| w21 = torch.exp(-((x - x2) ** 2 + (y - y1) ** 2) / (sigma ** 2)) | |
| w22 = torch.exp(-((x - x2) ** 2 + (y - y2) ** 2) / (sigma ** 2)) | |
| return w11, w12, w21, w22 | |