| 43 | """ Localization Network of RARE, which predicts C' (K x 2) from I (I_width x I_height) """ |
| 44 | |
| 45 | def __init__(self, F, I_channel_num): |
| 46 | super(LocalizationNetwork, self).__init__() |
| 47 | self.F = F |
| 48 | self.I_channel_num = I_channel_num |
| 49 | self.conv = nn.Sequential( |
| 50 | nn.Conv2d(in_channels=self.I_channel_num, out_channels=64, kernel_size=3, stride=1, padding=1, |
| 51 | bias=False), nn.BatchNorm2d(64), nn.ReLU(True), |
| 52 | nn.MaxPool2d(2, 2), # batch_size x 64 x I_height/2 x I_width/2 |
| 53 | nn.Conv2d(64, 128, 3, 1, 1, bias=False), nn.BatchNorm2d(128), nn.ReLU(True), |
| 54 | nn.MaxPool2d(2, 2), # batch_size x 128 x I_height/4 x I_width/4 |
| 55 | nn.Conv2d(128, 256, 3, 1, 1, bias=False), nn.BatchNorm2d(256), nn.ReLU(True), |
| 56 | nn.MaxPool2d(2, 2), # batch_size x 256 x I_height/8 x I_width/8 |
| 57 | nn.Conv2d(256, 512, 3, 1, 1, bias=False), nn.BatchNorm2d(512), nn.ReLU(True), |
| 58 | nn.AdaptiveAvgPool2d(1) # batch_size x 512 |
| 59 | ) |
| 60 | |
| 61 | self.localization_fc1 = nn.Sequential(nn.Linear(512, 256), nn.ReLU(True)) |
| 62 | self.localization_fc2 = nn.Linear(256, self.F * 2) |
| 63 | |
| 64 | # Init fc2 in LocalizationNetwork |
| 65 | self.localization_fc2.weight.data.fill_(0) |
| 66 | """ see RARE paper Fig. 6 (a) """ |
| 67 | ctrl_pts_x = np.linspace(-1.0, 1.0, int(F / 2)) |
| 68 | ctrl_pts_y_top = np.linspace(0.0, -1.0, num=int(F / 2)) |
| 69 | ctrl_pts_y_bottom = np.linspace(1.0, 0.0, num=int(F / 2)) |
| 70 | ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1) |
| 71 | ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1) |
| 72 | initial_bias = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0) |
| 73 | self.localization_fc2.bias.data = torch.from_numpy(initial_bias).float().view(-1) |
| 74 | |
| 75 | def forward(self, batch_I): |
| 76 | """ |