Spaces:
Runtime error
Runtime error
File size: 5,577 Bytes
98c5805 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 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 |