Vincentqyw
fix: cpu running
472119d
raw
history blame
6.2 kB
import torch
import torch.nn as nn
import torch.nn.functional as F
from lanet_utils import image_grid
class ConvBlock(nn.Module):
def __init__(self, in_channels, out_channels):
super(ConvBlock, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(
in_channels,
out_channels,
kernel_size=3,
stride=1,
padding=1,
bias=False,
),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.Conv2d(
out_channels,
out_channels,
kernel_size=3,
stride=1,
padding=1,
bias=False,
),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
)
def forward(self, x):
return self.conv(x)
class DilationConv3x3(nn.Module):
def __init__(self, in_channels, out_channels):
super(DilationConv3x3, self).__init__()
self.conv = nn.Conv2d(
in_channels,
out_channels,
kernel_size=3,
stride=1,
padding=2,
dilation=2,
bias=False,
)
self.bn = nn.BatchNorm2d(out_channels)
def forward(self, x):
x = self.conv(x)
x = self.bn(x)
return x
class InterestPointModule(nn.Module):
def __init__(self, is_test=False):
super(InterestPointModule, self).__init__()
self.is_test = is_test
self.conv1 = ConvBlock(3, 32)
self.conv2 = ConvBlock(32, 64)
self.conv3 = ConvBlock(64, 128)
self.conv4 = ConvBlock(128, 256)
self.maxpool2x2 = nn.MaxPool2d(2, 2)
# score head
self.score_conv = nn.Conv2d(
256, 256, kernel_size=3, stride=1, padding=1, bias=False
)
self.score_norm = nn.BatchNorm2d(256)
self.score_out = nn.Conv2d(256, 3, kernel_size=3, stride=1, padding=1)
self.softmax = nn.Softmax(dim=1)
# location head
self.loc_conv = nn.Conv2d(
256, 256, kernel_size=3, stride=1, padding=1, bias=False
)
self.loc_norm = nn.BatchNorm2d(256)
self.loc_out = nn.Conv2d(256, 2, kernel_size=3, stride=1, padding=1)
# descriptor out
self.des_conv2 = DilationConv3x3(64, 256)
self.des_conv3 = DilationConv3x3(128, 256)
# cross_head:
self.shift_out = nn.Conv2d(256, 1, kernel_size=3, stride=1, padding=1)
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
B, _, H, W = x.shape
x = self.conv1(x)
x = self.maxpool2x2(x)
x2 = self.conv2(x)
x = self.maxpool2x2(x2)
x3 = self.conv3(x)
x = self.maxpool2x2(x3)
x = self.conv4(x)
B, _, Hc, Wc = x.shape
# score head
score_x = self.score_out(self.relu(self.score_norm(self.score_conv(x))))
aware = self.softmax(score_x[:, 0:2, :, :])
score = score_x[:, 2, :, :].unsqueeze(1).sigmoid()
border_mask = torch.ones(B, Hc, Wc)
border_mask[:, 0] = 0
border_mask[:, Hc - 1] = 0
border_mask[:, :, 0] = 0
border_mask[:, :, Wc - 1] = 0
border_mask = border_mask.unsqueeze(1)
score = score * border_mask.to(score.device)
# location head
coord_x = self.relu(self.loc_norm(self.loc_conv(x)))
coord_cell = self.loc_out(coord_x).tanh()
shift_ratio = self.shift_out(coord_x).sigmoid() * 2.0
step = ((H / Hc) - 1) / 2.0
center_base = (
image_grid(
B,
Hc,
Wc,
dtype=coord_cell.dtype,
device=coord_cell.device,
ones=False,
normalized=False,
).mul(H / Hc)
+ step
)
coord_un = center_base.add(coord_cell.mul(shift_ratio * step))
coord = coord_un.clone()
coord[:, 0] = torch.clamp(coord_un[:, 0], min=0, max=W - 1)
coord[:, 1] = torch.clamp(coord_un[:, 1], min=0, max=H - 1)
# descriptor block
desc_block = []
desc_block.append(self.des_conv2(x2))
desc_block.append(self.des_conv3(x3))
desc_block.append(aware)
if self.is_test:
coord_norm = coord[:, :2].clone()
coord_norm[:, 0] = (coord_norm[:, 0] / (float(W - 1) / 2.0)) - 1.0
coord_norm[:, 1] = (coord_norm[:, 1] / (float(H - 1) / 2.0)) - 1.0
coord_norm = coord_norm.permute(0, 2, 3, 1)
desc2 = torch.nn.functional.grid_sample(desc_block[0], coord_norm)
desc3 = torch.nn.functional.grid_sample(desc_block[1], coord_norm)
aware = desc_block[2]
desc = torch.mul(desc2, aware[:, 0, :, :]) + torch.mul(
desc3, aware[:, 1, :, :]
)
desc = desc.div(
torch.unsqueeze(torch.norm(desc, p=2, dim=1), 1)
) # Divide by norm to normalize.
return score, coord, desc
return score, coord, desc_block
class CorrespondenceModule(nn.Module):
def __init__(self, match_type="dual_softmax"):
super(CorrespondenceModule, self).__init__()
self.match_type = match_type
if self.match_type == "dual_softmax":
self.temperature = 0.1
else:
raise NotImplementedError()
def forward(self, source_desc, target_desc):
b, c, h, w = source_desc.size()
source_desc = source_desc.div(
torch.unsqueeze(torch.norm(source_desc, p=2, dim=1), 1)
).view(b, -1, h * w)
target_desc = target_desc.div(
torch.unsqueeze(torch.norm(target_desc, p=2, dim=1), 1)
).view(b, -1, h * w)
if self.match_type == "dual_softmax":
sim_mat = (
torch.einsum("bcm, bcn -> bmn", source_desc, target_desc)
/ self.temperature
)
confidence_matrix = F.softmax(sim_mat, 1) * F.softmax(sim_mat, 2)
else:
raise NotImplementedError()
return confidence_matrix