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