Spaces:
Runtime error
Runtime error
| 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 |