Spaces:
Running
Running
File size: 1,016 Bytes
a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 |
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 |
import torch
import torch.nn as nn
class SuperpointDecoder(nn.Module):
"""Junction decoder based on the SuperPoint architecture."""
def __init__(self, input_feat_dim=128, backbone_name="lcnn"):
super(SuperpointDecoder, self).__init__()
self.relu = torch.nn.ReLU(inplace=True)
# Perform strided convolution when using lcnn backbone.
if backbone_name == "lcnn":
self.convPa = torch.nn.Conv2d(
input_feat_dim, 256, kernel_size=3, stride=2, padding=1
)
elif backbone_name == "superpoint":
self.convPa = torch.nn.Conv2d(
input_feat_dim, 256, kernel_size=3, stride=1, padding=1
)
else:
raise ValueError("[Error] Unknown backbone option.")
self.convPb = torch.nn.Conv2d(256, 65, kernel_size=1, stride=1, padding=0)
def forward(self, input_features):
feat = self.relu(self.convPa(input_features))
semi = self.convPb(feat)
return semi
|