inference DFNet. It can regress camera pose as well as extract intermediate layer features. :param x: image blob (2B x C x H x W) two stream or (B x C x H x W) single stream :param return_feature: whether to return features as output :param isSingleStream
(self, x, return_feature=False, isSingleStream=False, upsampleH=120, upsampleW=213)
| 86 | self.fc_pose = nn.Linear(1536, feat_dim) |
| 87 | |
| 88 | def forward(self, x, return_feature=False, isSingleStream=False, upsampleH=120, upsampleW=213): |
| 89 | ''' |
| 90 | inference DFNet. It can regress camera pose as well as extract intermediate layer features. |
| 91 | :param x: image blob (2B x C x H x W) two stream or (B x C x H x W) single stream |
| 92 | :param return_feature: whether to return features as output |
| 93 | :param isSingleStream: whether it's an single stream inference or siamese network inference |
| 94 | :param upsampleH: feature upsample size H |
| 95 | :param upsampleW: feature upsample size W |
| 96 | :return feature_maps: (2, [B, C, H, W]) or (1, [B, C, H, W]) or None |
| 97 | :return predict: [2B, 12] or [B, 12] |
| 98 | ''' |
| 99 | # normalize input data |
| 100 | mean, std = x.new_tensor(self.mean), x.new_tensor(self.std) |
| 101 | x = (x - mean[:, None, None]) / std[:, None, None] |
| 102 | |
| 103 | ### encoder ### |
| 104 | feature_maps = [] |
| 105 | list_x = self.feature_extractor(x) |
| 106 | |
| 107 | x = list_x['reduction_6'] # features to save |
| 108 | for i in self.feature_block_index: |
| 109 | fe = list_x['reduction_'+str(i)].clone() |
| 110 | feature_maps.append(fe) |
| 111 | |
| 112 | ### extract and process intermediate features ### |
| 113 | if return_feature: |
| 114 | feature_maps = self.adaptation_layers(feature_maps) # (3, [B, C, H', W']), H', W' are different in each layer |
| 115 | |
| 116 | pdb.set_trace() |
| 117 | if isSingleStream: # not siamese network style inference |
| 118 | feature_stacks = [] |
| 119 | |
| 120 | for f in feature_maps: |
| 121 | feature_stacks.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(f)) |
| 122 | feature_maps = [torch.stack(feature_stacks)] # (1, [3, B, C, H, W]) |
| 123 | else: # siamese network style inference |
| 124 | feature_stacks_t = [] |
| 125 | feature_stacks_r = [] |
| 126 | |
| 127 | for f in feature_maps: |
| 128 | # split real and nerf batches |
| 129 | batch = f.shape[0] # should be target batch_size + rgb batch_size |
| 130 | feature_t = f[:batch//2] |
| 131 | feature_r = f[batch//2:] |
| 132 | |
| 133 | feature_stacks_t.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(feature_t)) # GT img |
| 134 | feature_stacks_r.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(feature_r)) # render img |
| 135 | feature_stacks_t = torch.stack(feature_stacks_t) # [3, B, C, H, W] |
| 136 | feature_stacks_r = torch.stack(feature_stacks_r) # [3, B, C, H, W] |
| 137 | feature_maps = [feature_stacks_t, feature_stacks_r] # (2, [3, B, C, H, W]) |
| 138 | |
| 139 | else: |
| 140 | feature_maps = None |
| 141 | |
| 142 | ### pose regression head ### |
| 143 | x = self.avgpool(x) |
| 144 | x = x.reshape(x.size(0), -1) |
| 145 | predict = self.fc_pose(x) |
nothing calls this directly
no outgoing calls
no test coverage detected