DFNet implementation
| 72 | return features |
| 73 | |
| 74 | class DFNet(nn.Module): |
| 75 | ''' DFNet implementation ''' |
| 76 | default_conf = { |
| 77 | 'hypercolumn_layers': ["conv1_2", "conv3_3", "conv5_3"], |
| 78 | 'output_dim': 128, |
| 79 | } |
| 80 | mean = [0.485, 0.456, 0.406] |
| 81 | std = [0.229, 0.224, 0.225] |
| 82 | |
| 83 | def __init__(self, feat_dim=12, places365_model_path=''): |
| 84 | super(DFNet, self).__init__() |
| 85 | |
| 86 | self.layer_to_index = {k: v for v, k in enumerate(vgg16_layers.keys())} |
| 87 | self.hypercolumn_indices = [self.layer_to_index[n] for n in self.default_conf['hypercolumn_layers']] # [2, 14, 28] |
| 88 | |
| 89 | # Initialize architecture |
| 90 | vgg16 = models.vgg16(pretrained=True) |
| 91 | |
| 92 | self.encoder = nn.Sequential(*list(vgg16.features.children())) |
| 93 | |
| 94 | self.scales = [] |
| 95 | current_scale = 0 |
| 96 | for i, layer in enumerate(self.encoder): |
| 97 | if isinstance(layer, torch.nn.MaxPool2d): |
| 98 | current_scale += 1 |
| 99 | if i in self.hypercolumn_indices: |
| 100 | self.scales.append(2**current_scale) |
| 101 | |
| 102 | ## adaptation layers, see off branches from fig.3 in S2DNet paper |
| 103 | self.adaptation_layers = AdaptLayers(self.default_conf['hypercolumn_layers'], self.default_conf['output_dim']) |
| 104 | |
| 105 | # pose regression layers |
| 106 | self.avgpool = nn.AdaptiveAvgPool2d(1) |
| 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) |