def execute(self, x, cls_label):
batch_size, _, N = x.size()
x = self.relu(self.bn1(self.conv1(x))) # B, D, N
x = self.relu(self.bn2(self.conv2(x)))
x1 = self.sa1(x)
x2 = self.sa2(x1)
x3 = self.sa3(x2)
x4 = self.sa4(x3)
x = concat((x1, x2, x3, x4), dim=1)
x = self.conv_fuse(x)
x_max = jt.max(x, 2)
x_avg = jt.mean(x, 2)
x_max_feature = x_max.view(batch_size,
-1).unsqueeze(-1).repeat(1, 1, N)
x_avg_feature = x_avg.view(batch_size,
-1).unsqueeze(-1).repeat(1, 1, N)
cls_label_one_hot = cls_label.view(batch_size, 16, 1)
cls_label_feature = self.label_conv(cls_label_one_hot).repeat(1, 1, N)
x_global_feature = concat(
(x_max_feature, x_avg_feature, cls_label_feature), 1) # 1024 + 64
x = concat((x, x_global_feature), 1) # 1024 * 3 + 64
x = self.relu(self.bns1(self.convs1(x)))
x = self.dp1(x)
x = self.relu(self.bns2(self.convs2(x)))
x = self.convs3(x)
return x
def execute(self, point_cloud, label):
B, D, N = point_cloud.size()
trans = self.stn(point_cloud)
point_cloud = point_cloud.transpose(0, 2, 1)
point_cloud = nn.bmm(point_cloud, trans)
point_cloud = point_cloud.transpose(0, 2, 1)
out1 = self.relu(self.bn1(self.conv1(point_cloud)))
out2 = self.relu(self.bn2(self.conv2(out1)))
out3 = self.relu(self.bn3(self.conv3(out2)))
trans_feat = self.fstn(out3)
x = out3.transpose(0, 2, 1)
net_transformed = nn.bmm(x, trans_feat)
net_transformed = net_transformed.transpose(0, 2, 1)
out4 = self.relu(self.bn4(self.conv4(net_transformed)))
out5 = self.bn5(self.conv5(out4))
out_max = jt.argmax(out5, 2, keepdims=True)[1]
out_max = out_max.view(-1, 2048)
out_max = concat((out_max, label), 1)
expand = out_max.view(-1, 2048 + 16, 1).repeat(1, 1, N)
concat_feature = concat([expand, out1, out2, out3, out4, out5], 1)
net = self.relu(self.bns1(self.convs1(concat_feature)))
net = self.relu(self.bns2(self.convs2(net)))
net = self.relu(self.bns3(self.convs3(net)))
net = self.convs4(net)
return net
def execute(self, x):
imsize = x.shape
_, _, _, x = self.backbone(x)
b, c, h, w = x.shape
x_k = self.conv_0(x)
x1 = self.layer5a(x_k).reshape(b, c, -1)
x2 = self.layer5b(x_k).reshape(b, c, -1)
x3 = self.layer5c(x_k).reshape(b, c, -1)
x4 = self.layer5d(x_k).reshape(b, c, -1)
x_k = concat([x1, x2, x3, x4], 2).transpose(0, 2, 1) # b 110 c
x_q = self.conv_1(x)
x_q = x_q.reshape(b, c, -1)
x_attention = nn.bmm(x_k, x_q) # b 110 N
x_v = self.conv_2(x)
x1 = self.layer5a(x_v).reshape(b, c, -1)
x2 = self.layer5b(x_v).reshape(b, c, -1)
x3 = self.layer5c(x_v).reshape(b, c, -1)
x4 = self.layer5d(x_v).reshape(b, c, -1)
x_v = concat([x1, x2, x3, x4], 2) # b c 110
x = nn.bmm(x_v, x_attention).reshape(b, c, h, w)
x = self.final_conv(x)
x = nn.resize(x,
size=(imsize[2], imsize[3]),
mode='bilinear',
align_corners=True)
return x
def sample_and_group(num_points, num_neighbors, points, features=None):
""" sample `num-points` from points via furthest point sampling algorithm
and consider them sampled points as centroids to construct KNN graphs
on points where k=`num-neighbors`
"""
B, _, C = points.shape
S = num_points
fps_idx = farthest_point_sample(points, num_points)
# _, fps_idx = sampler(points) # [B, npoint]
sampled_points = index_points(points, fps_idx)
point_knn_idx = knn_points(num_neighbors, points, sampled_points)
grouped_points = index_points(
points, point_knn_idx) # [B, num-point, num-neighbos, C]
if features is None:
reshaped_points = sampled_points.view(B, S, 1, C)
grouped_points_relative = grouped_points - reshaped_points
return sampled_points, concat([
grouped_points_relative,
reshaped_points.repeat(1, 1, num_neighbors, 1)
],
dim=-1)
sampled_features = index_points(features, fps_idx)
grouped_features = index_points(features, point_knn_idx)
reshaped_features = sampled_features.view(B, S, 1, -1)
grouped_feature_relative = grouped_features - reshaped_features
return sampled_points, concat([
grouped_feature_relative,
reshaped_features.repeat(1, 1, num_neighbors, 1)
],
dim=-1)
def execute(self, x):
#print (x.shape)
# x = x.transpose(0, 2, 1) # b, n, c -> b, c, n
batch_size = x.shape[0]
x = get_graph_feature(x, knn=self.knn, k=self.k)
x = self.conv1(x)
x1 = x.max(dim=-1, keepdims=False)
x = get_graph_feature(x1, knn=self.knn, k=self.k)
x = self.conv2(x)
x2 = x.max(dim=-1, keepdims=False)
x = get_graph_feature(x2, knn=self.knn, k=self.k)
x = self.conv3(x)
x3 = x.max(dim=-1, keepdims=False)
x = get_graph_feature(x3, knn=self.knn, k=self.k)
x = self.conv4(x)
x4 = x.max(dim=-1, keepdims=False)
x = concat((x1, x2, x3, x4), dim=1)
x = self.conv5(x) # ############ this line
x1 = x.max(dim=2).reshape(batch_size, -1)
x2 = x.mean(dim=2).reshape(batch_size, -1)
x = concat((x1, x2), 1)
x = nn.leaky_relu(self.bn6(self.linear1(x)), scale=0.2)
x = self.dp1(x)
x = nn.leaky_relu(self.bn7(self.linear2(x)), scale=0.2)
x = self.dp2(x)
x = self.linear3(x)
return x
def execute(self, pcd, prev_s):
"""
Args:
pcd: b, n, 3
prev_s: b, c, n
"""
b, n, _ = pcd.shape
pcd_bcn = pcd.transpose(0, 2, 1)
l0_xyz = pcd
l0_points = pcd_bcn
if self.if_noise:
noise_points = init.gauss([b, 3, n], 'float', mean=0.0, std=self.noise_stdv)
l0_points = jittor.concat([l0_points, noise_points], 1)
l1_xyz, l1_points = self.sa_module_1(l0_xyz, l0_points) # b, 512, 128 (bnc)
l2_xyz, l2_points = self.sa_module_2(l1_xyz, l1_points)
l3_xyz, l3_points = self.sa_module_3(l2_xyz, l2_points)
l2_points = self.fp_module_3(l2_xyz, l3_xyz, l2_points, l3_points)
l2_points, prev_s['l2'] = self.unit_3(l2_points, prev_s['l2'])
l1_points = self.fp_module_2(l1_xyz, l2_xyz, l1_points, l2_points)
l1_points, prev_s['l1'] = self.unit_2(l1_points, prev_s['l1'])
l0_points = self.fp_module_1(l0_xyz, l1_xyz, concat([pcd_bcn, pcd_bcn], dim=1), l1_points)
l0_points, prev_s['l0'] = self.unit_1(l0_points, prev_s['l0']) # (B, 128, 2048)
noise = init.gauss([b, 32, n], 'float', mean=0.0, std=1.0)
feat = concat([l0_points, noise], dim=1)
delta_xyz = self.tanh(self.mlp_conv(feat)) * 1.0 / 10 ** (self.step - 1)
point_cloud = (pcd_bcn + delta_xyz).transpose(0, 2, 1)
return point_cloud, delta_xyz
def execute(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
spx = out
outs = []
for i in range(self.nums):
if i==0 or self.stype=='stage':
sp = spx[:, i*self.width: (i+1)*self.width]
else:
sp = sp + spx[:, i*self.width: (i+1)*self.width]
sp = self.convs[i](sp)
sp = self.relu(self.bns[i](sp))
outs.append(sp)
if self.stype=='normal' or self.stride==1:
outs.append(spx[:, self.nums*self.width: (self.nums+1)*self.width])
elif self.stype=='stage':
outs.append(self.pool(spx[:, self.nums*self.width: (self.nums+1)*self.width]))
out = concat(outs, 1)
out = self.conv3(out)
out = self.bn3(out)
if self.downsample is not None:
residual = self.downsample(x)
out += residual
out = self.relu(out)
return out
def execute(self, x):
xyz = x.permute(0, 2, 1)
batch_size, _, _ = x.size()
x = self.relu(self.bn1(self.conv1(x))) # B, D, N
x = self.relu(self.bn2(self.conv2(x))) # B, D, N
x = x.permute(0, 2, 1)
new_xyz, new_feature = sample_and_group(npoint=512,
nsample=32,
xyz=xyz,
points=x)
feature_0 = self.gather_local_0(new_feature)
feature = feature_0.permute(0, 2, 1)
new_xyz, new_feature = sample_and_group(npoint=256,
nsample=32,
xyz=new_xyz,
points=feature)
feature_1 = self.gather_local_1(new_feature)
x = self.pt_last(feature_1)
x = concat([x, feature_1], dim=1)
x = self.conv_fuse(x)
x = jt.max(x, 2)
x = x.view(batch_size, -1)
x = self.relu(self.bn6(self.linear1(x)))
x = self.dp1(x)
x = self.relu(self.bn7(self.linear2(x)))
x = self.dp2(x)
x = self.linear3(x)
return x
def sample_and_group(npoint, nsample, xyz, points):
B, N, C = xyz.shape
S = npoint
# xyz = xyz.contiguous()
sampler = FurthestPointSampler(npoint)
_, fps_idx = sampler(xyz) # [B, npoint]
# print ('fps size=', fps_idx.size())
# fps_idx = sampler(xyz).long() # [B, npoint]
new_xyz = index_points(xyz, fps_idx)
new_points = index_points(points, fps_idx)
# new_xyz = xyz[:]
# new_points = points[:]
idx = knn_point(nsample, xyz, new_xyz)
#idx = query_ball_point(radius, nsample, xyz, new_xyz)
grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C]
grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)
grouped_points = index_points(points, idx)
grouped_points_norm = grouped_points - new_points.view(B, S, 1, -1)
new_points = concat([
grouped_points_norm,
new_points.view(B, S, 1, -1).repeat(1, 1, nsample, 1)
],
dim=-1)
return new_xyz, new_points
def execute(self, cur_x, prev_s):
"""
Args:
cur_x: Tensor, (B, in_channel, N)
prev_s: Tensor, (B, in_channel, N)
Returns:
h: Tensor, (B, in_channel, N)
h: Tensor, (B, in_channel, N)
"""
if self.step == 1:
return cur_x, cur_x
z = self.conv_z(concat([cur_x, prev_s], dim=1))
r = self.conv_r(concat([cur_x, prev_s], dim=1))
h_hat = self.conv_h(concat([cur_x, r * prev_s], dim=1))
h = (1 - z) * cur_x + z * h_hat
return h, h
def execute(self, x, l):
batch_size = x.size(0)
num_points = x.size(2)
# print (x.size())
# x0 = get_graph_feature(x, k=self.k) # (batch_size, 3, num_points) -> (batch_size, 3*2, num_points, k)
x = get_graph_feature(x, knn=self.knn, k=self.k)
x = self.conv1(x) # (batch_size, 3*2, num_points, k) -> (batch_size, 64, num_points, k)
x = self.conv2(x) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points, k)
x1 = x.max(dim=-1, keepdims=False) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points)
x = get_graph_feature(x1, knn=self.knn, k=self.k)
x = self.conv3(x) # (batch_size, 64*2, num_points, k) -> (batch_size, 64, num_points, k)
x = self.conv4(x) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points, k)
x2 = x.max(dim=-1, keepdims=False) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points)
x = get_graph_feature(x2, knn=self.knn, k=self.k)
x = self.conv5(x) # (batch_size, 64*2, num_points, k) -> (batch_size, 64, num_points, k)
x3 = x.max(dim=-1, keepdims=False) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points)
x = concat((x1, x2, x3), dim=1) # (batch_size, 64*3, num_points)
x = self.conv6(x) # (batch_size, 64*3, num_points) -> (batch_size, emb_dims, num_points)
x = x.max(dim=-1, keepdims=True) # (batch_size, emb_dims, num_points) -> (batch_size, emb_dims, 1)
l = l.view(batch_size, -1, 1) # (batch_size, num_categoties, 1)
l = self.conv7(l) # (batch_size, num_categoties, 1) -> (batch_size, 64, 1)
x = concat((x, l), dim=1) # (batch_size, 1088, 1)
x = x.repeat(1, 1, num_points) # (batch_size, 1088, num_points)
x = concat((x, x1, x2, x3), dim=1) # (batch_size, 1088+64*3, num_points)
x = self.conv8(x) # (batch_size, 1088+64*3, num_points) -> (batch_size, 256, num_points)
x = self.dp1(x)
x = self.conv9(x) # (batch_size, 256, num_points) -> (batch_size, 256, num_points)
x = self.dp2(x)
x = self.conv10(x) # (batch_size, 256, num_points) -> (batch_size, 128, num_points)
x = self.conv11(x) # (batch_size, 256, num_points) -> (batch_size, seg_num_all, num_points)
return x
def execute(self, x):
x1 = self.aspp1(x)
x2 = self.aspp2(x)
x3 = self.aspp3(x)
x4 = self.aspp4(x)
x5 = self.global_avg_pool(x)
x5 = x5.broadcast((1,1,x4.shape[2],x4.shape[3]))
x = concat((x1, x2, x3, x4, x5), dim=1)
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
return x
def execute(self, x, low_level_feat):
low_level_feat = self.conv1(low_level_feat)
low_level_feat = self.bn1(low_level_feat)
low_level_feat = self.relu(low_level_feat)
#print (low_level_feat.shape)
x = nn.resize(x,
size=(low_level_feat.shape[2], low_level_feat.shape[3]),
mode='bilinear')
x = concat((x, low_level_feat), dim=1)
x = self.last_conv(x)
return x
def execute(self, x):
size = x.shape
_, _, _, x = self.backbone(x)
x = self.final_conv(
concat([
x,
self.layer5a(x),
self.layer5b(x),
self.layer5c(x),
self.layer5d(x),
], 1))
return nn.resize(x, size=(size[2], size[3]), mode='bilinear')
def execute(self, xyz1, xyz2, points1, points2):
"""
Input:
xyz1: input points position data, [B, C, N]
xyz2: sampled input points position data, [B, C, S]
points1: input points data, [B, N, D]
points2: input points data, [B, S, D]
Return:
new_points: upsampled points data, [B, N, D']
"""
# xyz1 = xyz1.permute(0, 2, 1)
# xyz2 = xyz2.permute(0, 2, 1)
# points2 = points2.permute(0, 2, 1)
B, N, C = xyz1.shape
_, S, _ = xyz2.shape
points2 = points2.transpose(0, 2, 1) # b, n, c
if S == 1:
interpolated_points = points2.repeat(1, N, 1)
else:
# dists = square_distance(xyz1, xyz2)
# idx, dists = jt.argsort(dists, dim=-1)
# dists, idx = dists[:, :, :3], idx[:, :, :3] # [B, N, 3]
_, idx = three_nn(xyz1, xyz2)
dists = ((xyz1.unsqueeze(-2) - xyz2['i0', idx]) ** 2).sum(-1)
# dists = jt.ones([B, N, 3])
# idx = jt.index([B, N, 3], 2)
dists = jt.clamp(dists, min_v=1e-10)
dist_recip = 1.0 / dists
norm = jt.sum(dist_recip, dim=2, keepdims=True)
weight = dist_recip / norm
interpolated_points = jt.sum(index_points(points2, idx) * weight.view(B, N, 3, 1), dim=2)
interpolated_points = interpolated_points.transpose(0, 2, 1) # b, c, n
if points1 is not None:
# points1 = points1.permute(0, 2, 1)
new_points = concat([interpolated_points, points1], dim=1)
else:
new_points = interpolated_points
# new_points = new_points.permute(0, 2, 1) # b, c, n
# l = len(self.mlp_convs)
# for i, conv in self.mlp_convs.layers.items():
# conv = self.mlp_convs[i]
# bn = self.mlp_bns[i]
# new_points = self.relu(bn(conv(new_points)))
new_points = self.mlp(new_points)
return new_points # .permute(0, 2, 1) # b, n, c
def execute(self, features, points=None):
# batch-size, num-points, num-dims -> batch-size, num-dims, num-points
batch_size = features.size(0)
if points is not None:
# add position embedding
# points: batch-size, num-points, num-dims -> batch-size, num-dims, num-points
points = self.pos_conv(points)
# end
# features: batch-size, num-points, num-dims -> batch-size, num-dims, num-points
features = features.permute(0, 2, 1)
x = self.conv(features) # B, D, N
x1 = self.oa1(x, points)
x2 = self.oa2(x1, points)
x3 = self.oa3(x2, points)
x4 = self.oa4(x3, points)
return self.fuse_conv(concat((features, x1, x2, x3, x4), dim=1))
def execute(self, context, feature):
batch_size, c, h, w = feature.shape
origin_feature = feature
feature = feature.reshape(batch_size, c, -1).transpose(0, 2,
1) # b, h*w, c
context = context.reshape(batch_size, context.shape[1],
-1) # b, n_cls, h*w
attention = self.softmax(context)
ocr_context = nn.bmm(attention, feature).transpose(0, 2,
1) # b, c, n_cls
relation = nn.bmm(feature, ocr_context).transpose(0, 2,
1) # b, n_cls, h*w
attention = self.softmax(relation) #b , n_cls, h*w
result = nn.bmm(ocr_context, attention).reshape(batch_size, c, h, w)
result = self.conv_1x1(result)
result = concat([result, origin_feature], dim=1)
result = self.last_conv(result)
return result
def execute(self, x_l, x_h): # (N, P, C_out)
"""
Given a point cloud, and its corresponding features, return a new set
of randomly-sampled representative points with features projected from
the point cloud.
:param x: (pts, fts) where
- pts: Regional point cloud such that fts[:,p_idx,:] is the
feature associated with pts[:,p_idx,:].
- fts: Regional features such that pts[:,p_idx,:] is the feature
associated with fts[:,p_idx,:].
:return: Randomly subsampled points and their features.
"""
pts_l, fts_l = x_l
pts_h, fts_h = x_h
rep_pts_fts = self.pointcnn((pts_h, pts_l, fts_l))
concat_feature = concat ((rep_pts_fts, fts_h), dim=2)
rep_pts_fts = self.conv_fuse(concat_feature)
return pts_h, rep_pts_fts
def execute(self, xyz1, xyz2, points1, points2):
"""
Input:
xyz1: input points position data, [B, C, N]
xyz2: sampled input points position data, [B, C, S]
points1: input points data, [B, D, N]
points2: input points data, [B, D, S]
Return:
new_points: upsampled points data, [B, D', N]
"""
# xyz1 = xyz1.permute(0, 2, 1)
# xyz2 = xyz2.permute(0, 2, 1)
# points2 = points2.permute(0, 2, 1)
B, N, C = xyz1.shape
_, S, _ = xyz2.shape
if S == 1:
interpolated_points = points2.repeat(1, N, 1)
else:
dists = square_distance(xyz1, xyz2)
idx, dists = jt.argsort(dists, dim=-1)
dists, idx = dists[:, :, :3], idx[:, :, :3] # [B, N, 3]
dist_recip = 1.0 / (dists + 1e-8)
norm = jt.sum(dist_recip, dim=2, keepdims=True)
weight = dist_recip / norm
interpolated_points = jt.sum(index_points(points2, idx) *
weight.view(B, N, 3, 1),
dim=2)
if points1 is not None:
# points1 = points1.permute(0, 2, 1)
new_points = concat([points1, interpolated_points], dim=-1)
else:
new_points = interpolated_points
new_points = new_points.permute(0, 2, 1)
# l = len(self.mlp_convs)
for i, conv in self.mlp_convs.layers.items():
# conv = self.mlp_convs[i]
bn = self.mlp_bns[i]
new_points = self.relu(bn(conv(new_points)))
return new_points.permute(0, 2, 1)
def sample_and_group_with_density_scale(npoint,
nsample,
xyz,
points,
density_scale=None):
"""
Input:
npoint:
nsample:
xyz: input points position data, [B, N, C]
points: input points data, [B, N, D]
Return:
new_xyz: sampled points position data, [B, 1, C]
new_points: sampled points data, [B, 1, N, C+D]
"""
B, N, C = xyz.shape
S = npoint
fps_idx = farthest_point_sample(xyz, npoint) # [B, npoint, C]
# jt.sync_all(True)
# print ('11111111111111111')
new_xyz = index_points(xyz, fps_idx)
# jt.sync_all(True)
# print ('2222222222222222222')
idx = knn_points(nsample, xyz, new_xyz)
# jt.sync_all(True)
# print ('333333333333333333')
grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C]
grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)
if points is not None:
grouped_points = index_points(points, idx)
new_points = concat([grouped_xyz_norm, grouped_points],
dim=-1) # [B, npoint, nsample, C+D]
else:
new_points = grouped_xyz_norm
# jt.sync_all(True)
# print ('44444444444444444444444')
if density_scale is None:
return new_xyz, new_points, grouped_xyz_norm, idx
else:
grouped_density = index_points(density_scale, idx)
return new_xyz, new_points, grouped_xyz_norm, idx, grouped_density
def execute(self, x):
#
# b, 3, npoint, nsample
# conv2d 3 -> 128 channels 1, 1
# b * npoint, c, nsample
# permute reshape
batch_size, _, N = x.size()
x = self.relu(self.bn1(self.conv1(x))) # B, D, N
x = self.relu(self.bn2(self.conv2(x)))
x1 = self.sa1(x)
x2 = self.sa2(x1)
x3 = self.sa3(x2)
x4 = self.sa4(x3)
x = concat((x1, x2, x3, x4), dim=1)
return x
def execute(self, x, xyz):
#
# b, 3, npoint, nsample
# conv2d 3 -> 128 channels 1, 1
# b * npoint, c, nsample
# permute reshape
batch_size, _, N = x.size()
# add position embedding
xyz = xyz.permute(0, 2, 1)
xyz = self.pos_xyz(xyz)
# end
x = self.relu(self.bn1(self.conv1(x))) # B, D, N
x1 = self.sa1(x, xyz)
x2 = self.sa2(x1, xyz)
x3 = self.sa3(x2, xyz)
x4 = self.sa4(x3, xyz)
x = concat((x1, x2, x3, x4), dim=1)
return x
def execute(self, xyz, feature, cls_label):
# for module in self.pointnet_modules:
# xyz, feature = module(xyz, feature)
B, N, _ = xyz.shape
l1_xyz, l1_feature = self.pointnet_modules[0](xyz, feature)
l2_xyz, l2_feature = self.pointnet_modules[1](l1_xyz, l1_feature)
l3_xyz, l3_feature = self.pointnet_modules[2](l2_xyz, l2_feature)
# print ('before interpolate shape')
# print (l2_xyz.shape, l2_feature.shape, l3_xyz.shape, l3_feature.shape)
l2_feature = self.fp3(l2_xyz, l3_xyz, l2_feature, l3_feature)
l1_feature = self.fp2(l1_xyz, l2_xyz, l1_feature, l2_feature)
cls_label_one_hot = cls_label.view(B, 16,
1).repeat(1, 1, N).permute(0, 2, 1)
# print ('before concat size ')
# print (cls_label_one_hot.size(),xyz.size(),feature.size())
feature = self.fp1(xyz, l1_xyz,
concat([cls_label_one_hot, xyz, feature], 2),
l1_feature)
feature = feature.permute(0, 2, 1)
# print (feature.shape)
return self.fc_layer(feature)
def get_graph_feature(x, knn=None, k=None, idx=None):
batch_size = x.shape[0]
num_points = x.shape[2]
x = x.reshape(batch_size, -1, num_points)
if idx is None:
idx = knn(x,x) # (batch_size, num_points, k)
idx = idx.permute(0, 2, 1)
idx_base = jt.array(np.arange(0, batch_size)).reshape(-1, 1, 1)*num_points
idx = idx + idx_base
idx = idx.reshape(-1)
_, num_dims, _ = x.shape
x = x.transpose(0, 2, 1) # (batch_size, num_points, num_dims) -> (batch_size*num_points, num_dims) # batch_size * num_points * k + range(0, batch_size*num_points)
feature = x.reshape(batch_size*num_points, -1)[idx, :]
feature = feature.reshape(batch_size, num_points, k, num_dims)
x = x.reshape(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1)
feature = concat((feature-x, x), dim=3).transpose(0, 3, 1, 2)
return feature
def sample_and_group_all(xyz, points, density_scale=None):
"""
Input:
xyz: input points position data, [B, N, C]
points: input points data, [B, N, D]
Return:
new_xyz: sampled points position data, [B, 1, C]
new_points: sampled points data, [B, 1, N, C+D]
"""
#device = xyz.device
B, N, C = xyz.shape
#new_xyz = torch.zeros(B, 1, C).to(device)
new_xyz = xyz.mean(dim=1, keepdims=True)
grouped_xyz = xyz.reshape(B, 1, N, C) - new_xyz.reshape(B, 1, 1, C)
if points is not None:
new_points = concat([grouped_xyz, points.view(B, 1, N, -1)], dim=3)
else:
new_points = grouped_xyz
if density_scale is None:
return new_xyz, new_points, grouped_xyz
else:
grouped_density = density_scale.reshape(B, 1, N, 1)
return new_xyz, new_points, grouped_xyz, grouped_density
def execute(self, x):
"""
inputs :
x : input feature maps( B X C X H X W)
returns :
out : attention value + input feature
attention: B X (HxW) X (HxW)
"""
m_batchsize, C, height, width = x.size()
proj_query = self.query_conv(x).reshape(m_batchsize, -1,
width * height).transpose(
0, 2, 1)
proj_key = self.key_conv(x).reshape(m_batchsize, -1, width * height)
energy = nn.bmm(proj_query, proj_key)
attention = self.softmax(energy)
proj_value = self.value_conv(x).reshape(m_batchsize, -1,
width * height)
out = nn.bmm(proj_value, attention.transpose(0, 2, 1))
out = out.reshape(m_batchsize, C, height, width)
out = self.scale_conv(out)
out = concat([out, x], 1)
return out
def execute(self, x):
batch_size, _, N = x.size()
# print (x.size())
x = self.relu(self.bn1(self.conv1(x))) # B, D, N
x = self.relu(self.bn2(self.conv2(x)))
x1 = self.sa1(x)
x2 = self.sa2(x1)
x3 = self.sa3(x2)
x4 = self.sa4(x3)
x = concat((x1, x2, x3, x4), dim=1)
x = self.conv_fuse(x)
# x = F.adaptive_max_pool1d(x, 1).view(batch_size, -1)
x = jt.max(x, 2)
x = x.view(batch_size, -1)
x = self.relu(self.bn6(self.linear1(x)))
x = self.dp1(x)
x = self.relu(self.bn7(self.linear2(x)))
x = self.dp2(x)
x = self.linear3(x)
return x
def execute(self, x): # (N, K, C_out)
"""
Applies XConv to the input data.
:param x: (rep_pt, pts, fts) where
- rep_pt: Representative point.
- pts: Regional point cloud such that fts[:,p_idx,:] is the feature
associated with pts[:,p_idx,:].
- fts: Regional features such that pts[:,p_idx,:] is the feature
associated with fts[:,p_idx,:].
:return: Features aggregated into point rep_pt.
"""
rep_pt, pts, fts = x # b, n, c // b ,n k, c // b, n, k, d
if fts is not None:
assert(rep_pt.size()[0] == pts.size()[0] == fts.size()[0]) # Check N is equal.
assert(rep_pt.size()[1] == pts.size()[1] == fts.size()[1]) # Check P is equal.
assert(pts.size()[2] == fts.size()[2] == self.K) # Check K is equal.
assert(fts.size()[3] == self.cin) # Check C_in is equal.
else:
assert(rep_pt.size()[0] == pts.size()[0]) # Check N is equal.
assert(rep_pt.size()[1] == pts.size()[1]) # Check P is equal.
assert(pts.size()[2] == self.K) # Check K is equal.
assert(rep_pt.size()[2] == pts.size()[3] == self.dims) # Check dims is equal.
N = pts.size()[0]
P = rep_pt.size()[1] # (N, P, K, dims)
p_center = jt.unsqueeze(rep_pt, dim = 2) # (N, P, 1, dims)
# print (p_center.size()) #
# Move pts to local coordinate system of rep_pt.
pts_local = pts - p_center.repeat(1, 1, self.K, 1) # (N, P, K, dims)
# pts_local = self.pts_layernorm(pts - p_center)
# Individually lift each point into C_mid space.
# print (pts_local.size(), 'before size')
pts_local = pts_local.permute(0, 3, 1, 2) # N, dim, P, K
fts_lifted0 = self.dense1(pts_local) # ?
# print (.size(), 'after size')
fts_lifted = self.dense2(fts_lifted0) # N, C_mid, P, K
fts = fts.permute(0, 3, 1, 2)
if fts is None:
fts_cat = fts_lifted
else:
fts_cat = concat((fts_lifted, fts), 1) # (N, C_mid + C_in, P, K)
# Learn the (N, K, K) X-transformation matrix.
X_shape = (N, P, self.K, self.K)
# X = self.x_trans(pts_local) # N, K*K, 1, P
x = self.x_trans_0(pts_local)
x = self.x_trans_1(x)
X = self.x_trans_2(x)
# print ('X size ', X.size())
X = X.permute(0, 2, 3, 1) # n p 1 k
X = X.view(X_shape) # N, P, K, K
# print (fts_cat.shape)
fts_cat = fts_cat.permute(0, 2, 3, 1)
fts_X = jt.matmul(X, fts_cat) #
# print ('fts X size =', fts_X.shape)
fts_p = self.end_conv(fts_X).squeeze(dim = 2)
# print ('xxxxxxxxxxx')
# print ('result size')
# print (fts_X.size(), fts_p.size())
return fts_p
