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, return_pose=True, upsampleH=240, upsampleW=427)
| 107 | self.fc_pose = nn.Linear(512, feat_dim) |
| 108 | |
| 109 | def forward(self, x, return_feature=False, isSingleStream=False, return_pose=True, upsampleH=240, upsampleW=427): |
| 110 | ''' |
| 111 | inference DFNet. It can regress camera pose as well as extract intermediate layer features. |
| 112 | :param x: image blob (2B x C x H x W) two stream or (B x C x H x W) single stream |
| 113 | :param return_feature: whether to return features as output |
| 114 | :param isSingleStream: whether it's an single stream inference or siamese network inference |
| 115 | :param upsampleH: feature upsample size H |
| 116 | :param upsampleW: feature upsample size W |
| 117 | :return feature_maps: (2, [B, C, H, W]) or (1, [B, C, H, W]) or None |
| 118 | :return predict: [2B, 12] or [B, 12] |
| 119 | ''' |
| 120 | # normalize input data |
| 121 | mean, std = x.new_tensor(self.mean), x.new_tensor(self.std) |
| 122 | x = (x - mean[:, None, None]) / std[:, None, None] |
| 123 | |
| 124 | ### encoder ### |
| 125 | feature_maps = [] |
| 126 | for i in range(len(self.encoder)): |
| 127 | x = self.encoder[i](x) |
| 128 | |
| 129 | if i in self.hypercolumn_indices: |
| 130 | feature = x.clone() |
| 131 | feature_maps.append(feature) |
| 132 | |
| 133 | if i==self.hypercolumn_indices[-1]: |
| 134 | if return_pose==False: |
| 135 | predict = None |
| 136 | break |
| 137 | |
| 138 | ### extract and process intermediate features ### |
| 139 | if return_feature: |
| 140 | feature_maps = self.adaptation_layers(feature_maps) # (3, [B, C, H', W']), H', W' are different in each layer |
| 141 | |
| 142 | if isSingleStream: # not siamese network style inference |
| 143 | feature_stacks = [] |
| 144 | for f in feature_maps: |
| 145 | feature_stacks.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(f)) |
| 146 | feature_maps = [torch.stack(feature_stacks)] # (1, [3, B, C, H, W]) |
| 147 | else: # siamese network style inference |
| 148 | feature_stacks_t = [] |
| 149 | feature_stacks_r = [] |
| 150 | for f in feature_maps: |
| 151 | # split real and nerf batches |
| 152 | batch = f.shape[0] # should be target batch_size + rgb batch_size |
| 153 | feature_t = f[:batch//2] |
| 154 | feature_r = f[batch//2:] |
| 155 | |
| 156 | feature_stacks_t.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(feature_t)) # GT img |
| 157 | feature_stacks_r.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(feature_r)) # render img |
| 158 | feature_stacks_t = torch.stack(feature_stacks_t) # [3, B, C, H, W] |
| 159 | feature_stacks_r = torch.stack(feature_stacks_r) # [3, B, C, H, W] |
| 160 | feature_maps = [feature_stacks_t, feature_stacks_r] # (2, [3, B, C, H, W]) |
| 161 | else: |
| 162 | feature_maps = None |
| 163 | |
| 164 | if return_pose==False: |
| 165 | return feature_maps, predict |
| 166 |
nothing calls this directly
no outgoing calls
no test coverage detected