Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,565 Bytes
c509e76 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
from __future__ import absolute_import
import math
import numpy as np
import sys
import torch
from torch import nn
from torch.nn import functional as F
from torch.nn import init
def conv3x3_block(in_planes, out_planes, stride=1):
"""3x3 convolution with padding"""
conv_layer = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=1, padding=1)
block = nn.Sequential(
conv_layer,
nn.BatchNorm2d(out_planes),
nn.ReLU(inplace=True),
)
return block
class STNHead(nn.Module):
def __init__(self, in_planes, num_ctrlpoints, activation='none'):
super(STNHead, self).__init__()
self.in_planes = in_planes
self.num_ctrlpoints = num_ctrlpoints
self.activation = activation
self.stn_convnet = nn.Sequential(
conv3x3_block(in_planes, 32), # 32*64
nn.MaxPool2d(kernel_size=2, stride=2),
conv3x3_block(32, 64), # 16*32
nn.MaxPool2d(kernel_size=2, stride=2),
conv3x3_block(64, 128), # 8*16
nn.MaxPool2d(kernel_size=2, stride=2),
conv3x3_block(128, 256), # 4*8
nn.MaxPool2d(kernel_size=2, stride=2),
conv3x3_block(256, 256), # 2*4,
nn.MaxPool2d(kernel_size=2, stride=2),
conv3x3_block(256, 256)) # 1*2 > 256*8*8
self.stn_fc1 = nn.Sequential(
# nn.Linear(2*256, 512),
nn.Linear(8*8*256, 512),
nn.BatchNorm1d(512),
nn.ReLU(inplace=True))
self.stn_fc2 = nn.Linear(512, num_ctrlpoints*2)
self.init_weights(self.stn_convnet)
self.init_weights(self.stn_fc1)
self.init_stn(self.stn_fc2)
def init_weights(self, module):
for m in module.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
if m.bias is not None:
m.bias.data.zero_()
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
m.weight.data.normal_(0, 0.001)
m.bias.data.zero_()
def init_stn(self, stn_fc2):
# margin = 0.01
# sampling_num_per_side = int(self.num_ctrlpoints / 2)
# ctrl_pts_x = np.linspace(margin, 1.-margin, sampling_num_per_side)
# ctrl_pts_y_top = np.ones(sampling_num_per_side) * margin
# ctrl_pts_y_bottom = np.ones(sampling_num_per_side) * (1-margin)
# ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
# ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
# ctrl_points = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0).astype(np.float32)
margin_x, margin_y = 0.35,0.35
# margin_x, margin_y = 0,0
num_ctrl_pts_per_side = (self.num_ctrlpoints-4) // 4 +2
ctrl_pts_x = np.linspace(margin_x, 1.0 - margin_x, num_ctrl_pts_per_side)
ctrl_pts_y_top = np.ones(num_ctrl_pts_per_side) * margin_y
ctrl_pts_y_bottom = np.ones(num_ctrl_pts_per_side) * (1.0 - margin_y)
ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
ctrl_pts_x_left = np.ones(num_ctrl_pts_per_side) * margin_x
ctrl_pts_x_right = np.ones(num_ctrl_pts_per_side) * (1.0-margin_x)
ctrl_pts_left = np.stack([ctrl_pts_x_left[1:-1], ctrl_pts_x[1:-1]], axis=1)
ctrl_pts_right = np.stack([ctrl_pts_x_right[1:-1], ctrl_pts_x[1:-1]], axis=1)
ctrl_points = np.concatenate([ctrl_pts_top, ctrl_pts_bottom, ctrl_pts_left, ctrl_pts_right], axis=0).astype(np.float32)
if self.activation is 'none':
pass
elif self.activation == 'sigmoid':
ctrl_points = -np.log(1. / ctrl_points - 1.)
stn_fc2.weight.data.zero_()
stn_fc2.bias.data = torch.Tensor(ctrl_points).view(-1)
def forward(self, x):
x = self.stn_convnet(x)
batch_size, _, h, w = x.size()
x = x.view(batch_size, -1)
img_feat = self.stn_fc1(x)
x = self.stn_fc2(0.1 * img_feat)
if self.activation == 'sigmoid':
x = F.sigmoid(x)
x = x.view(-1, self.num_ctrlpoints, 2)
return img_feat, x
if __name__ == "__main__":
in_planes = 3
num_ctrlpoints = 20
activation='none' # 'sigmoid'
stn_head = STNHead(in_planes, num_ctrlpoints, activation)
input = torch.randn(10, 3, 32, 64)
control_points = stn_head(input)
print(control_points.size()) |