Upload models/anchor_gen.py

models/anchor_gen.py

@@ -0,0 +1,107 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Function
+from models import basic, clusterkit
+import pdb
+
+class AnchorAnalysis:
+    def __init__(self, mode, colorLabeler):
+        ## anchor generating mode: 1.random; 2.clustering
+        self.mode = mode
+        self.colorLabeler = colorLabeler
+
+    def _detect_correlation(self, data_tensors, color_probs, hint_masks, thres=0.1):
+        N,C,H,W = data_tensors.shape
+        ## (N,C,HW)
+        data_vecs = data_tensors.flatten(2)
+        prob_vecs = color_probs.flatten(2)
+        mask_vecs = hint_masks.flatten(2)
+        #anchor_data = torch.masked_select(data_vecs, mask_vecs.bool()).view(N,C,-1)
+        #anchor_prob = torch.masked_select(prob_vecs, mask_vecs.bool()).view(N,313,-1)
+        #_,_,K = anchor_data.shape
+        anchor_mask = torch.matmul(mask_vecs.permute(0,2,1), mask_vecs)
+        cosine_sim = True
+        ## non-similarity matrix
+        if cosine_sim:
+            norm_data = F.normalize(data_vecs, p=2, dim=1)
+            ## (N,HW,HW) = (N,HW,C) X (N,C,HW)
+            corr_matrix = torch.matmul(norm_data.permute(0,2,1), norm_data)
+            ## remapping: [-1.0,1.0] to [0.0,1.0], and convert into dis-similarity
+            dist_matrix = 1.0 - 0.5*(corr_matrix + 1.0)
+        else:
+            ## (N,HW,HW) = (N,HW,C) X (N,C,HW)
+            XtX = torch.matmul(data_vecs.permute(0,2,1), data_vecs)
+            diag_vec = torch.diagonal(XtX, dim1=-2, dim2=-1)
+            A = diag_vec.unsqueeze(1).repeat(1,H*W,1)
+            At = diag_vec.unsqueeze(2).repeat(1,1,H*W)
+            dist_matrix = A - 2*XtX + At
+        #dist_matrix = dist_matrix + 1e7*torch.eye(K).to(data_tensors.device).repeat(N,1,1)
+        ## for debug use
+        K = 8
+        anchor_adj_matrix = torch.masked_select(dist_matrix, anchor_mask.bool()).view(N,K,K)
+        ## dectect connected nodes
+        adj_matrix = torch.where((dist_matrix < thres) & (anchor_mask > 0), torch.ones_like(dist_matrix), torch.zeros_like(dist_matrix))
+        adj_matrix = torch.matmul(adj_matrix, adj_matrix)
+        adj_matrix = adj_matrix / (1e-7+adj_matrix)
+        ## merge nodes
+        ## (N,K,C) = (N,K,K) X (N,K,C)
+        anchor_prob = torch.matmul(adj_matrix, prob_vecs.permute(0,2,1)) / torch.sum(adj_matrix, dim=2, keepdim=True)
+        updated_prob_vecs = anchor_prob.permute(0,2,1) * mask_vecs + (1-mask_vecs) * prob_vecs
+        color_probs = updated_prob_vecs.view(N,313,H,W)
+        return color_probs, anchor_adj_matrix
+
+    def _sample_anchor_colors(self, pred_prob, hint_mask, T=0):
+        N,C,H,W = pred_prob.shape
+        topk = 10
+        assert T < topk
+        sorted_probs, batch_indexs = torch.sort(pred_prob, dim=1, descending=True)
+        ## (N,topk,H,W,1)
+        topk_probs = torch.softmax(sorted_probs[:,:topk,:,:], dim=1).unsqueeze(4)
+        topk_indexs = batch_indexs[:,:topk,:,:]
+        topk_ABs = torch.stack([self.colorLabeler.q_to_ab.index_select(0, q_i.flatten()).reshape(topk,H,W,2)
+                    for q_i in topk_indexs])
+        ## (N,topk,H,W,2)
+        topk_ABs = topk_ABs / 110.0
+        ## choose the most distinctive 3 colors for each anchor
+        if T == 0:
+            sampled_ABs = topk_ABs[:,0,:,:,:]
+        elif T == 1:
+            sampled_AB0 = topk_ABs[:,[0],:,:,:]
+            internal_diff = torch.norm(topk_ABs-sampled_AB0, p=2, dim=4, keepdim=True)
+            _, batch_indexs = torch.sort(internal_diff, dim=1, descending=True)
+            ## (N,1,H,W,2)
+            selected_index = batch_indexs[:,[0],:,:,:].expand([-1,-1,-1,-1,2])
+            sampled_ABs = torch.gather(topk_ABs, 1, selected_index)
+            sampled_ABs = sampled_ABs.squeeze(1)
+        else:
+            sampled_AB0 = topk_ABs[:,[0],:,:,:]
+            internal_diff = torch.norm(topk_ABs-sampled_AB0, p=2, dim=4, keepdim=True)
+            _, batch_indexs = torch.sort(internal_diff, dim=1, descending=True)
+            selected_index = batch_indexs[:,[0],:,:,:].expand([-1,-1,-1,-1,2])
+            sampled_AB1 = torch.gather(topk_ABs, 1, selected_index)
+            internal_diff2 = torch.norm(topk_ABs-sampled_AB1, p=2, dim=4, keepdim=True)
+            _, batch_indexs = torch.sort(internal_diff+internal_diff2, dim=1, descending=True)
+            ## (N,1,H,W,2)
+            selected_index = batch_indexs[:,[T-2],:,:,:].expand([-1,-1,-1,-1,2])
+            sampled_ABs = torch.gather(topk_ABs, 1, selected_index)
+            sampled_ABs = sampled_ABs.squeeze(1)
+
+        return sampled_ABs.permute(0,3,1,2)
+
+    def __call__(self, data_tensors, n_anchors, spixel_sizes, use_sklearn_kmeans=False):
+        N,C,H,W = data_tensors.shape
+        if self.mode == 'clustering':
+            ## clusters map: (N,K,H,W)
+            cluster_mask = clusterkit.batch_kmeans_pytorch(data_tensors, n_anchors, 'euclidean', use_sklearn_kmeans)
+            #noises = torch.rand(N,1,H,W).to(cluster_mask.device)
+            perturb_factors = spixel_sizes
+            cluster_prob = cluster_mask + perturb_factors * 0.01
+            hint_mask_layers = F.one_hot(torch.argmax(cluster_prob.flatten(2), dim=-1), num_classes=H*W).float()
+            hint_mask = torch.sum(hint_mask_layers, dim=1, keepdim=True).view(N,1,H,W)
+        else:
+            #print('----------hello, random!')
+            cluster_mask = torch.zeros(N,n_anchors,H,W).to(data_tensors.device)
+            binary_mask = basic.get_random_mask(N, H, W, minNum=n_anchors, maxNum=n_anchors)
+            hint_mask = torch.from_numpy(binary_mask).to(data_tensors.device)
+        return hint_mask, cluster_mask