448 lines
16 KiB
Python
448 lines
16 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from . import pointnet2_utils as pointutils
|
|
|
|
RADIUS = 2.5
|
|
|
|
|
|
def index_points_group(points, knn_idx):
|
|
"""
|
|
Input:
|
|
points: input points data, [B, N, C]
|
|
knn_idx: sample index data, [B, N, K]
|
|
Return:
|
|
new_points:, indexed points data, [B, N, K, C]
|
|
"""
|
|
points_flipped = points.permute(0, 2, 1).contiguous()
|
|
new_points = pointutils.grouping_operation(points_flipped,
|
|
knn_idx.int()).permute(
|
|
0, 2, 3, 1)
|
|
|
|
return new_points
|
|
|
|
|
|
def curvature(pc, nsample=10, radius=RADIUS):
|
|
# pc: B 3 N
|
|
assert pc.shape[1] == 3
|
|
pc = pc.permute(0, 2, 1)
|
|
|
|
dist, kidx = pointutils.knn(nsample, pc.contiguous(),
|
|
pc.contiguous()) # (B, N, 10)
|
|
|
|
if radius is not None:
|
|
tmp_idx = kidx[:, :, 0].unsqueeze(2).repeat(1, 1,
|
|
nsample).to(kidx.device)
|
|
kidx[dist > radius] = tmp_idx[dist > radius]
|
|
|
|
grouped_pc = index_points_group(pc, kidx) # B N 10 3
|
|
pc_curvature = torch.sum(grouped_pc - pc.unsqueeze(2), dim=2) / 9.0
|
|
return pc_curvature # B N 3
|
|
|
|
|
|
class PointNetSetAbstractionRatio(nn.Module):
|
|
|
|
def __init__(self,
|
|
ratio,
|
|
radius,
|
|
nsample,
|
|
in_channel,
|
|
mlp,
|
|
group_all,
|
|
return_fps=False,
|
|
use_xyz=True,
|
|
use_act=True,
|
|
act=F.relu,
|
|
mean_aggr=False,
|
|
use_instance_norm=False):
|
|
super(PointNetSetAbstractionRatio, self).__init__()
|
|
self.ratio = ratio
|
|
self.radius = radius
|
|
self.nsample = nsample
|
|
self.group_all = group_all
|
|
self.use_xyz = use_xyz
|
|
self.use_act = use_act
|
|
self.mean_aggr = mean_aggr
|
|
self.act = act
|
|
self.mlp_convs = nn.ModuleList()
|
|
self.mlp_bns = nn.ModuleList()
|
|
last_channel = (in_channel + 3) if use_xyz else in_channel
|
|
for out_channel in mlp:
|
|
self.mlp_convs.append(
|
|
nn.Conv2d(last_channel, out_channel, 1, bias=False))
|
|
if use_instance_norm:
|
|
self.mlp_bns.append(
|
|
nn.InstanceNorm2d(out_channel, affine=True))
|
|
else:
|
|
self.mlp_bns.append(nn.BatchNorm2d(out_channel))
|
|
|
|
last_channel = out_channel
|
|
|
|
if group_all:
|
|
self.queryandgroup = pointutils.GroupAll(self.use_xyz)
|
|
else:
|
|
self.queryandgroup = pointutils.QueryAndGroup(
|
|
radius, nsample, self.use_xyz)
|
|
self.return_fps = return_fps
|
|
|
|
def forward(self, xyz, points, fps_idx=None):
|
|
"""
|
|
Input:
|
|
xyz: input points position data, [B, C, N]
|
|
points: input points data, [B, D, N]
|
|
Return:
|
|
new_xyz: sampled points position data, [B, C, S]
|
|
new_points: sample points feature data, [B, D', S]
|
|
"""
|
|
B, C, N = xyz.shape
|
|
npoint = int(N * self.ratio)
|
|
|
|
xyz = xyz.contiguous()
|
|
xyz_t = xyz.permute(0, 2, 1).contiguous()
|
|
|
|
if (self.group_all is False) and (npoint != -1):
|
|
if fps_idx is None:
|
|
fps_idx = pointutils.furthest_point_sample(xyz_t,
|
|
npoint) # [B, N]
|
|
new_xyz = pointutils.gather_operation(xyz, fps_idx) # [B, C, N]
|
|
else:
|
|
new_xyz = xyz
|
|
new_points, _ = self.queryandgroup(xyz_t,
|
|
new_xyz.transpose(2,
|
|
1).contiguous(),
|
|
points) # [B, 3+C, N, S]
|
|
|
|
# new_xyz: sampled points position data, [B, C, npoint]
|
|
# new_points: sampled points data, [B, C+D, npoint, nsample]
|
|
for i, conv in enumerate(self.mlp_convs):
|
|
if self.use_act:
|
|
bn = self.mlp_bns[i]
|
|
new_points = self.act(bn(conv(new_points)))
|
|
else:
|
|
new_points = conv(new_points)
|
|
|
|
if self.mean_aggr:
|
|
new_points = torch.mean(new_points, -1)
|
|
else:
|
|
new_points = torch.max(new_points, -1)[0]
|
|
|
|
if self.return_fps:
|
|
return new_xyz, new_points, fps_idx
|
|
else:
|
|
return new_xyz, new_points
|
|
|
|
|
|
class PointNetSetAbstraction(nn.Module):
|
|
|
|
def __init__(self,
|
|
npoint,
|
|
radius,
|
|
nsample,
|
|
in_channel,
|
|
mlp,
|
|
group_all,
|
|
return_fps=False,
|
|
use_xyz=True,
|
|
use_act=True,
|
|
act=F.relu,
|
|
mean_aggr=False,
|
|
use_instance_norm=False):
|
|
super(PointNetSetAbstraction, self).__init__()
|
|
self.npoint = npoint
|
|
self.radius = radius
|
|
self.nsample = nsample
|
|
self.group_all = group_all
|
|
self.use_xyz = use_xyz
|
|
self.use_act = use_act
|
|
self.mean_aggr = mean_aggr
|
|
self.act = act
|
|
self.mlp_convs = nn.ModuleList()
|
|
self.mlp_bns = nn.ModuleList()
|
|
last_channel = (in_channel + 3) if use_xyz else in_channel
|
|
for out_channel in mlp:
|
|
self.mlp_convs.append(
|
|
nn.Conv2d(last_channel, out_channel, 1, bias=False))
|
|
if use_instance_norm:
|
|
self.mlp_bns.append(
|
|
nn.InstanceNorm2d(out_channel, affine=True))
|
|
else:
|
|
self.mlp_bns.append(nn.BatchNorm2d(out_channel))
|
|
|
|
last_channel = out_channel
|
|
|
|
if group_all:
|
|
self.queryandgroup = pointutils.GroupAll(self.use_xyz)
|
|
else:
|
|
self.queryandgroup = pointutils.QueryAndGroup(
|
|
radius, nsample, self.use_xyz)
|
|
self.return_fps = return_fps
|
|
|
|
def forward(self, xyz, points, fps_idx=None):
|
|
"""
|
|
Input:
|
|
xyz: input points position data, [B, C, N]
|
|
points: input points data, [B, D, N]
|
|
Return:
|
|
new_xyz: sampled points position data, [B, S, C]
|
|
new_points: sample points feature data, [B, S, D']
|
|
"""
|
|
# device = xyz.device
|
|
B, C, N = xyz.shape
|
|
xyz = xyz.contiguous()
|
|
xyz_t = xyz.permute(0, 2, 1).contiguous()
|
|
|
|
if (self.group_all is False) and (self.npoint != -1):
|
|
if fps_idx is None:
|
|
fps_idx = pointutils.furthest_point_sample(
|
|
xyz_t, self.npoint) # [B, N]
|
|
new_xyz = pointutils.gather_operation(xyz, fps_idx) # [B, C, N]
|
|
else:
|
|
new_xyz = xyz
|
|
new_points, _ = self.queryandgroup(xyz_t,
|
|
new_xyz.transpose(2,
|
|
1).contiguous(),
|
|
points) # [B, 3+C, N, S]
|
|
|
|
# new_xyz: sampled points position data, [B, C, npoint]
|
|
# new_points: sampled points data, [B, C+D, npoint, nsample]
|
|
for i, conv in enumerate(self.mlp_convs):
|
|
if self.use_act:
|
|
bn = self.mlp_bns[i]
|
|
new_points = self.act(bn(conv(new_points)))
|
|
else:
|
|
new_points = conv(new_points)
|
|
|
|
if self.mean_aggr:
|
|
new_points = torch.mean(new_points, -1)
|
|
else:
|
|
new_points = torch.max(new_points, -1)[0]
|
|
|
|
if self.return_fps:
|
|
return new_xyz, new_points, fps_idx
|
|
else:
|
|
return new_xyz, new_points
|
|
|
|
|
|
class PointNetFeaturePropogation(nn.Module):
|
|
|
|
def __init__(self, in_channel, mlp, learn_mask=False, nsample=3):
|
|
super(PointNetFeaturePropogation, self).__init__()
|
|
self.mlp_convs = nn.ModuleList()
|
|
self.mlp_bns = nn.ModuleList()
|
|
self.apply_mlp = mlp is not None
|
|
last_channel = in_channel
|
|
self.nsample = nsample
|
|
if self.apply_mlp:
|
|
for out_channel in mlp:
|
|
self.mlp_convs.append(nn.Conv1d(last_channel, out_channel, 1))
|
|
self.mlp_bns.append(nn.BatchNorm1d(out_channel))
|
|
last_channel = out_channel
|
|
|
|
if learn_mask:
|
|
self.queryandgroup = pointutils.QueryAndGroup(
|
|
None, 9, use_xyz=True)
|
|
last_channel = (128 + 3)
|
|
for out_channel in [32, 1]:
|
|
self.mlp_convs.append(
|
|
nn.Conv2d(last_channel, out_channel, 1, bias=False))
|
|
self.mlp_bns.append(nn.BatchNorm2d(out_channel))
|
|
last_channel = out_channel
|
|
|
|
def forward(self, pos1, pos2, feature1, feature2, hidden=None):
|
|
"""
|
|
Input:
|
|
pos1: input points position data, [B, C, N]
|
|
pos2: sampled input points position data, [B, C, S]
|
|
feature1: input points data, [B, D, N]
|
|
feature2: input points data, [B, D, S]
|
|
Return:
|
|
feat_new: upsampled points data, [B, D', N]
|
|
"""
|
|
pos1_t = pos1.permute(0, 2, 1).contiguous()
|
|
pos2_t = pos2.permute(0, 2, 1).contiguous()
|
|
B, C, N = pos1.shape
|
|
|
|
if hidden is None:
|
|
if self.nsample == 3:
|
|
dists, idx = pointutils.three_nn(pos1_t, pos2_t)
|
|
else:
|
|
dists, idx = pointutils.knn(self.nsample, pos1_t, pos2_t)
|
|
dists[dists < 1e-10] = 1e-10
|
|
weight = 1.0 / dists
|
|
weight = weight / torch.sum(weight, -1, keepdim=True) # [B,N,3]
|
|
interpolated_feat = torch.sum(
|
|
pointutils.grouping_operation(feature2, idx)
|
|
* weight.view(B, 1, N, self.nsample),
|
|
dim=-1) # [B,C,N,3]
|
|
else:
|
|
dist, idx = pointutils.knn(9, pos1_t, pos2_t)
|
|
|
|
new_feat, _ = self.queryandgroup(pos2_t, pos1_t,
|
|
hidden) # [B, 3+C, N, 9]
|
|
|
|
for i, conv in enumerate(self.mlp_convs):
|
|
new_feat = conv(new_feat)
|
|
weight = torch.softmax(new_feat, dim=-1) # [B, 1, N, 9]
|
|
interpolated_feat = torch.sum(
|
|
pointutils.grouping_operation(feature2, idx) * weight,
|
|
dim=-1) # [B, C, N]
|
|
|
|
if feature1 is not None:
|
|
feat_new = torch.cat([interpolated_feat, feature1], 1)
|
|
else:
|
|
feat_new = interpolated_feat
|
|
|
|
if self.apply_mlp:
|
|
for i, conv in enumerate(self.mlp_convs):
|
|
bn = self.mlp_bns[i]
|
|
feat_new = F.relu(bn(conv(feat_new)))
|
|
return feat_new
|
|
|
|
|
|
class Sinkhorn(nn.Module):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, corr, epsilon, gamma, max_iter):
|
|
# Early return if no iteration
|
|
if max_iter == 0:
|
|
return corr
|
|
|
|
# Init. of Sinkhorn algorithm
|
|
power = gamma / (gamma + epsilon)
|
|
a = (
|
|
torch.ones((corr.shape[0], corr.shape[1], 1),
|
|
device=corr.device,
|
|
dtype=corr.dtype) / corr.shape[1])
|
|
prob1 = (
|
|
torch.ones((corr.shape[0], corr.shape[1], 1),
|
|
device=corr.device,
|
|
dtype=corr.dtype) / corr.shape[1])
|
|
prob2 = (
|
|
torch.ones((corr.shape[0], corr.shape[2], 1),
|
|
device=corr.device,
|
|
dtype=corr.dtype) / corr.shape[2])
|
|
|
|
# Sinkhorn algorithm
|
|
for _ in range(max_iter):
|
|
# Update b
|
|
KTa = torch.bmm(corr.transpose(1, 2), a)
|
|
b = torch.pow(prob2 / (KTa + 1e-8), power)
|
|
# Update a
|
|
Kb = torch.bmm(corr, b)
|
|
a = torch.pow(prob1 / (Kb + 1e-8), power)
|
|
|
|
# Transportation map
|
|
T = torch.mul(torch.mul(a, corr), b.transpose(1, 2))
|
|
|
|
return T
|
|
|
|
|
|
class PointWiseOptimLayer(nn.Module):
|
|
|
|
def __init__(self, nsample, radius, in_channel, mlp, use_curvature=True):
|
|
super().__init__()
|
|
self.nsample = nsample
|
|
self.radius = radius
|
|
self.use_curvature = use_curvature
|
|
|
|
self.pos_embed = nn.Sequential(
|
|
nn.Conv1d(3, 32, 1), nn.ReLU(inplace=True), nn.Conv1d(32, 64, 1))
|
|
|
|
self.qk_net = nn.Sequential(
|
|
nn.Conv1d(in_channel + 64, in_channel + 64, 1),
|
|
nn.ReLU(inplace=True),
|
|
nn.Conv1d(in_channel + 64, in_channel + 64, 1))
|
|
if self.use_curvature:
|
|
self.curvate_net = nn.Sequential(
|
|
nn.Conv1d(3, 32, 1), nn.ReLU(inplace=True),
|
|
nn.Conv1d(32, 32, 1))
|
|
self.mlp_conv = nn.Conv1d(
|
|
in_channel + 64 + 32, mlp[-1], 1, bias=True)
|
|
else:
|
|
self.mlp_conv = nn.Conv1d(in_channel + 64, mlp[-1], 1, bias=True)
|
|
|
|
def forward(self,
|
|
pos1,
|
|
pos2,
|
|
feature1,
|
|
feature2,
|
|
nsample,
|
|
radius=None,
|
|
pos1_raw=None,
|
|
return_score=False):
|
|
"""
|
|
Input:
|
|
pos1: (batch_size, 3, npoint)
|
|
pos2: (batch_size, 3, npoint)
|
|
feature1: (batch_size, channel, npoint)
|
|
feature2: (batch_size, channel, npoint)
|
|
Output:
|
|
pos1: (batch_size, 3, npoint)
|
|
cost: (batch_size, channel, npoint)
|
|
"""
|
|
pos1_t = pos1.permute(0, 2, 1).contiguous()
|
|
pos2_t = pos2.permute(0, 2, 1).contiguous()
|
|
self.nsample = nsample
|
|
self.radius = radius
|
|
|
|
dist, idx = pointutils.knn(self.nsample, pos1_t, pos2_t) # [B, N, K]
|
|
if self.radius is not None:
|
|
tmp_idx = idx[:, :,
|
|
0].unsqueeze(2).repeat(1, 1,
|
|
self.nsample).to(idx.device)
|
|
idx[dist > self.radius] = tmp_idx[dist > self.radius]
|
|
|
|
pos1_embed_norm = self.pos_embed(pos1)
|
|
pos2_embed_norm = self.pos_embed(pos2) # [B, C1, N]
|
|
|
|
feat1_w_pos = torch.cat([feature1, pos1_embed_norm], dim=1)
|
|
feat2_w_pos = torch.cat([feature2, pos2_embed_norm],
|
|
dim=1) # [B, C1+C2, N]
|
|
|
|
feat1_w_pos = self.qk_net(feat1_w_pos)
|
|
feat2_w_pos = self.qk_net(feat2_w_pos) # [B, C1+C2, N]
|
|
|
|
feat2_grouped = pointutils.grouping_operation(feat2_w_pos,
|
|
idx) # [B, C1+C2, N, S]
|
|
|
|
score = torch.softmax(
|
|
feat1_w_pos.unsqueeze(-1) * feat2_grouped * 1.
|
|
/ math.sqrt(feat1_w_pos.shape[1]),
|
|
dim=-1) # [B, C1+C2, N, S]
|
|
cost = (score * (feat1_w_pos.unsqueeze(-1) - feat2_grouped)**2).sum(
|
|
dim=-1) # [B, C1+C2, N]
|
|
|
|
if self.use_curvature:
|
|
curvate1_raw = curvature(pos1_raw).permute(0, 2, 1) # [B, 3, N]
|
|
curvate1 = curvature(pos1).permute(0, 2, 1) # [B, 3, N]
|
|
curvate_cost = self.curvate_net(curvate1_raw) - self.curvate_net(
|
|
curvate1)
|
|
curvate_cost = curvate_cost**2
|
|
cost = self.mlp_conv(torch.cat([cost, curvate_cost],
|
|
dim=1)) # [B, C, N]
|
|
else:
|
|
cost = self.mlp_conv(cost) # [B, C, N]
|
|
|
|
if return_score:
|
|
pos2_grouped = pointutils.grouping_operation(pos2,
|
|
idx) # [B, 3, N, S]
|
|
# [B, N, K]
|
|
index = (dist > self.radius).sum(
|
|
dim=2, keepdim=True).float() > (dist.shape[2] - 0.1
|
|
) # [B, N, 1]
|
|
index = index.unsqueeze(1).repeat(1, score.shape[1], 1,
|
|
dist.shape[2]) # [B, N, K]
|
|
score_tmp = score.clone()
|
|
score_tmp[index] = 0.0
|
|
score = score_tmp
|
|
return pos1, cost, score, pos2_grouped
|
|
else:
|
|
return pos1, cost
|