Spaces:
Build error
Build error
File size: 2,414 Bytes
d444fe9 |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
class SurfaceClassifier(nn.Module):
def __init__(self, filter_channels, num_views=1, no_residual=True, last_op=None):
super(SurfaceClassifier, self).__init__()
self.filters = []
self.num_views = num_views
self.no_residual = no_residual
filter_channels = filter_channels
self.last_op = last_op
if self.no_residual:
for l in range(0, len(filter_channels) - 1):
self.filters.append(nn.Conv1d(
filter_channels[l],
filter_channels[l + 1],
1))
self.add_module("conv%d" % l, self.filters[l])
else:
for l in range(0, len(filter_channels) - 1):
if 0 != l:
self.filters.append(
nn.Conv1d(
filter_channels[l] + filter_channels[0],
filter_channels[l + 1],
1))
else:
self.filters.append(nn.Conv1d(
filter_channels[l],
filter_channels[l + 1],
1))
self.add_module("conv%d" % l, self.filters[l])
def forward(self, feature):
'''
:param feature: list of [BxC_inxHxW] tensors of image features
:param xy: [Bx3xN] tensor of (x,y) coodinates in the image plane
:return: [BxC_outxN] tensor of features extracted at the coordinates
'''
y = feature
tmpy = feature
for i, f in enumerate(self.filters):
if self.no_residual:
y = self._modules['conv' + str(i)](y)
else:
y = self._modules['conv' + str(i)](
y if i == 0
else torch.cat([y, tmpy], 1)
)
if i != len(self.filters) - 1:
y = F.leaky_relu(y)
if self.num_views > 1 and i == len(self.filters) // 2:
y = y.view(
-1, self.num_views, y.shape[1], y.shape[2]
).mean(dim=1)
tmpy = feature.view(
-1, self.num_views, feature.shape[1], feature.shape[2]
).mean(dim=1)
if self.last_op:
y = self.last_op(y)
return y
|