Localization Network of RARE, which predicts C' (K x 2) from I (I_width x I_height)
| 40 | |
| 41 | |
| 42 | class LocalizationNetwork(nn.Module): |
| 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 | """ |
| 77 | input: batch_I : Batch Input Image [batch_size x I_channel_num x I_height x I_width] |
| 78 | output: batch_C_prime : Predicted coordinates of fiducial points for input batch [batch_size x F x 2] |
| 79 | """ |
| 80 | batch_size = batch_I.size(0) |
| 81 | features = self.conv(batch_I).view(batch_size, -1) |
| 82 | batch_C_prime = self.localization_fc2(self.localization_fc1(features)).view(batch_size, self.F, 2) |
| 83 | return batch_C_prime |
| 84 | |
| 85 | |
| 86 | class GridGenerator(nn.Module): |