def corr_and_add(feature_A,
feature_B,
k=10,
coords_A=None,
coords_B=None,
Npts=None):
# compute sparse correlation from A to B
scorr = sparse_corr(feature_A,
feature_B,
k=k,
ratio=False,
sparse_type='raw')
# compute sparse correlation from B to A
scorr2 = sparse_corr(feature_B,
feature_A,
k=k,
ratio=False,
reverse=True,
sparse_type='raw')
scorr = ME.SparseTensor(scorr[0], scorr[1].cuda())
scorr2 = ME.SparseTensor(scorr2[0],
scorr2[1].cuda(),
coordinate_manager=scorr.coordinate_manager)
scorr = ME.MinkowskiUnion()(scorr, scorr2)
return scorr
def scalespace_correlation(self, src_feats, trg_feats, scales, conv2d):
r""" *_feats: list of features from intermediate layers
[(bsz, ch_1, h, w), (bsz, ch_2, h, w), ..., (bsz, ch_n, h, w)]
"""
bsz, ch, side1, side2 = src_feats[0].size()
bsz, ch, side3, side4 = trg_feats[0].size()
# new_src_feats = []
# new_trg_feats = []
# for scale, conv in zip(scales, conv2d[0]):
# new_side = round(side * math.sqrt(scale))
# new_src_feat = F.interpolate(src_feat, (new_side, new_side), mode='bilinear', align_corners=True)
# new_trg_feat = F.interpolate(trg_feat, (new_side, new_side), mode='bilinear', align_corners=True)
# new_src_feats.append(conv(new_src_feat))
# new_trg_feats.append(conv(new_trg_feat))
# correlations = []
length1 = bsz * side1 * side2 * self.k
length2 = bsz * side3 * side4 * self.k
#import pdb; pdb.set_trace()
scalewise_feat_st = torch.zeros(length1 * len(scales) * len(scales),
1,
device=torch.device('cuda'))
scalewise_coord_st = torch.zeros(length1 * len(scales) * len(scales),
7,
device=torch.device('cuda'),
dtype=torch.int)
scalewise_feat_ts = torch.zeros(length2 * len(scales) * len(scales),
1,
device=torch.device('cuda'))
scalewise_coord_ts = torch.zeros(length2 * len(scales) * len(scales),
7,
device=torch.device('cuda'),
dtype=torch.int)
for s_1, src_feat in enumerate(src_feats):
new_ch = src_feat.size(1)
#src_feat = F.interpolate(src_feat, (side, side),mode='bilinear', align_corners=True)
#src_feat = src_feat.view(bsz, new_ch, -1)
for s_2, trg_feat in enumerate(trg_feats):
#trg_feat = F.interpolate(trg_feat, (side, side),mode='bilinear', align_corners=True)
#trg_feat = trg_feat.view(bsz, new_ch, -1)
corr_coords_st = sparse_corr_6d([s_1, s_2],
src_feat,
trg_feat,
k=self.k,
ratio=False,
sparse_type='raw')
corr_coords_ts = sparse_corr_6d([s_1, s_2],
trg_feat,
src_feat,
k=self.k,
ratio=False,
sparse_type='raw',
reverse=True)
start1 = (s_2 + s_1 * len(scales)) * length1
end1 = start1 + length1
start2 = (s_2 + s_1 * len(scales)) * length2
end2 = start2 + length2
scalewise_feat_st[start1:end1] = corr_coords_st[0]
scalewise_coord_st[start1:end1] = corr_coords_st[1]
scalewise_feat_ts[start2:end2] = corr_coords_ts[0]
scalewise_coord_ts[start2:end2] = corr_coords_ts[1]
#import pdb; pdb.set_trace()
s2t = ME.SparseTensor(scalewise_feat_st, scalewise_coord_st)
t2s = ME.SparseTensor(scalewise_feat_ts,
scalewise_coord_ts,
coordinate_manager=s2t.coordinate_manager)
correlations = ME.MinkowskiUnion()(s2t, t2s)
return correlations
