| 25 | |
| 26 | |
| 27 | class BiRAFT(nn.Module): |
| 28 | def __init__(self, args): |
| 29 | super(BiRAFT, self).__init__() |
| 30 | self.args = args |
| 31 | |
| 32 | if args.small: |
| 33 | self.hidden_dim = hdim = 96 |
| 34 | self.context_dim = cdim = 64 |
| 35 | args.corr_levels = 4 |
| 36 | args.corr_radius = 3 |
| 37 | |
| 38 | else: |
| 39 | self.hidden_dim = hdim = 128 |
| 40 | self.context_dim = cdim = 128 |
| 41 | args.corr_levels = 4 |
| 42 | args.corr_radius = 4 |
| 43 | |
| 44 | if 'dropout' not in self.args: |
| 45 | self.args.dropout = 0 |
| 46 | |
| 47 | if 'alternate_corr' not in self.args: |
| 48 | self.args.alternate_corr = False |
| 49 | |
| 50 | # feature network, context network, and update block |
| 51 | if args.small: |
| 52 | self.fnet = SmallEncoder(output_dim=128, norm_fn='instance', dropout=args.dropout) |
| 53 | self.cnet = SmallEncoder(output_dim=hdim + cdim, norm_fn='none', dropout=args.dropout) |
| 54 | self.update_block = SmallUpdateBlock(self.args, hidden_dim=hdim) |
| 55 | |
| 56 | else: |
| 57 | if self.args.fnet == 'CNN': |
| 58 | self.fnet = BasicEncoder(output_dim=256, norm_fn='instance', dropout=args.dropout) |
| 59 | self.cnet = BasicEncoder(output_dim=hdim+cdim, norm_fn='batch', dropout=args.dropout) |
| 60 | elif self.args.fnet == 'twins': |
| 61 | self.fnet = twins_svt_large(pretrained=True) |
| 62 | self.cnet = twins_svt_large(pretrained=True) |
| 63 | self.update_block = BasicUpdateBlock(self.args, hidden_dim=hdim) |
| 64 | |
| 65 | def freeze_bn(self): |
| 66 | for m in self.modules(): |
| 67 | if isinstance(m, nn.BatchNorm2d): |
| 68 | m.eval() |
| 69 | |
| 70 | def initialize_flow(self, img): |
| 71 | """ Flow is represented as difference between two coordinate grids flow = coords1 - coords0""" |
| 72 | N, C, H, W = img.shape |
| 73 | coords0 = coords_grid(N, H // 8, W // 8, img.device) |
| 74 | coords1 = coords_grid(N, H // 8, W // 8, img.device) |
| 75 | |
| 76 | # optical flow computed as difference: flow = coords1 - coords0 |
| 77 | return coords0, coords1 |
| 78 | |
| 79 | def upsample_flow(self, flow, mask): |
| 80 | """ Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """ |
| 81 | N, _, H, W = flow.shape |
| 82 | mask = mask.view(N, 1, 9, 8, 8, H, W) |
| 83 | mask = torch.softmax(mask, dim=2) |
| 84 | |