vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/imatch-toolbox/third_party/patch2pix/networks/ncn/model.py

@@ -0,0 +1,333 @@
+from __future__ import print_function, division
+from collections import OrderedDict
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+import torchvision.models as models
+from .conv4d import Conv4d
+
+def Softmax1D(x,dim):
+    x_k = torch.max(x,dim)[0].unsqueeze(dim)
+    x -= x_k.expand_as(x)
+    exp_x = torch.exp(x)
+    return torch.div(exp_x,torch.sum(exp_x,dim).unsqueeze(dim).expand_as(x))
+
+def featureL2Norm(feature):
+    epsilon = 1e-6
+    norm = torch.pow(torch.sum(torch.pow(feature,2),1)+epsilon,0.5).unsqueeze(1).expand_as(feature)
+    feat_norm = torch.div(feature,norm)
+    return feat_norm
+
+class FeatureExtraction(torch.nn.Module):
+    def __init__(self, train_fe=False, feature_extraction_cnn='resnet101', feature_extraction_model_file='', normalization=True, last_layer='', use_cuda=True):
+        super(FeatureExtraction, self).__init__()
+        self.normalization = normalization
+        self.feature_extraction_cnn=feature_extraction_cnn
+        if feature_extraction_cnn == 'vgg':
+            self.model = models.vgg16(pretrained=True)
+            # keep feature extraction network up to indicated layer
+            vgg_feature_layers=['conv1_1','relu1_1','conv1_2','relu1_2','pool1','conv2_1',
+                         'relu2_1','conv2_2','relu2_2','pool2','conv3_1','relu3_1',
+                         'conv3_2','relu3_2','conv3_3','relu3_3','pool3','conv4_1',
+                         'relu4_1','conv4_2','relu4_2','conv4_3','relu4_3','pool4',
+                         'conv5_1','relu5_1','conv5_2','relu5_2','conv5_3','relu5_3','pool5']
+            if last_layer=='':
+                last_layer = 'pool4'
+            last_layer_idx = vgg_feature_layers.index(last_layer)
+            self.model = nn.Sequential(*list(self.model.features.children())[:last_layer_idx+1])
+        # for resnet below
+        resnet_feature_layers = ['conv1','bn1','relu','maxpool','layer1','layer2','layer3','layer4']
+        if feature_extraction_cnn=='resnet101':
+            self.model = models.resnet101(pretrained=True)            
+            if last_layer=='':
+                last_layer = 'layer3'                            
+            resnet_module_list = [getattr(self.model,l) for l in resnet_feature_layers]
+            last_layer_idx = resnet_feature_layers.index(last_layer)
+            self.model = nn.Sequential(*resnet_module_list[:last_layer_idx+1])
+
+        if feature_extraction_cnn=='resnet101fpn':
+            if feature_extraction_model_file!='':
+                resnet = models.resnet101(pretrained=True) 
+                # swap stride (2,2) and (1,1) in first layers (PyTorch ResNet is slightly different to caffe2 ResNet)
+                # this is required for compatibility with caffe2 models
+                resnet.layer2[0].conv1.stride=(2,2)
+                resnet.layer2[0].conv2.stride=(1,1)
+                resnet.layer3[0].conv1.stride=(2,2)
+                resnet.layer3[0].conv2.stride=(1,1)
+                resnet.layer4[0].conv1.stride=(2,2)
+                resnet.layer4[0].conv2.stride=(1,1)
+            else:
+                resnet = models.resnet101(pretrained=True) 
+            resnet_module_list = [getattr(resnet,l) for l in resnet_feature_layers]
+            conv_body = nn.Sequential(*resnet_module_list)
+            self.model = fpn_body(conv_body,
+                                  resnet_feature_layers,
+                                  fpn_layers=['layer1','layer2','layer3'],
+                                  normalize=normalization,
+                                  hypercols=True)
+            if feature_extraction_model_file!='':
+                self.model.load_pretrained_weights(feature_extraction_model_file)
+
+        if feature_extraction_cnn == 'densenet201':
+            self.model = models.densenet201(pretrained=True)
+            # keep feature extraction network up to denseblock3
+            # self.model = nn.Sequential(*list(self.model.features.children())[:-3])
+            # keep feature extraction network up to transitionlayer2
+            self.model = nn.Sequential(*list(self.model.features.children())[:-4])
+        if train_fe==False:
+            # freeze parameters
+            for param in self.model.parameters():
+                param.requires_grad = False
+        # move to GPU
+        if use_cuda:
+            self.model = self.model.cuda()
+        
+    def forward(self, image_batch):
+        features = self.model(image_batch)
+        if self.normalization and not self.feature_extraction_cnn=='resnet101fpn':
+            features = featureL2Norm(features)
+        return features
+    
+class FeatureCorrelation(torch.nn.Module):
+    def __init__(self,shape='3D',normalization=True):
+        super(FeatureCorrelation, self).__init__()
+        self.normalization = normalization
+        self.shape=shape
+        self.ReLU = nn.ReLU()
+    
+    def forward(self, feature_A, feature_B):        
+        if self.shape=='3D':
+            b,c,h,w = feature_A.size()
+            # reshape features for matrix multiplication
+            feature_A = feature_A.transpose(2,3).contiguous().view(b,c,h*w)
+            feature_B = feature_B.view(b,c,h*w).transpose(1,2)
+            # perform matrix mult.
+            feature_mul = torch.bmm(feature_B,feature_A)
+            # indexed [batch,idx_A=row_A+h*col_A,row_B,col_B]
+            correlation_tensor = feature_mul.view(b,h,w,h*w).transpose(2,3).transpose(1,2)
+        elif self.shape=='4D':
+            b,c,hA,wA = feature_A.size()
+            b,c,hB,wB = feature_B.size()
+            # reshape features for matrix multiplication
+            feature_A = feature_A.view(b,c,hA*wA).transpose(1,2) # size [b,c,h*w]
+            feature_B = feature_B.view(b,c,hB*wB) # size [b,c,h*w]
+            # perform matrix mult.
+            feature_mul = torch.bmm(feature_A,feature_B)
+            # indexed [batch,row_A,col_A,row_B,col_B]
+            correlation_tensor = feature_mul.view(b,hA,wA,hB,wB).unsqueeze(1)
+        
+        if self.normalization:
+            correlation_tensor = featureL2Norm(self.ReLU(correlation_tensor))
+            
+        return correlation_tensor
+
+class NeighConsensus(torch.nn.Module):
+    def __init__(self, use_cuda=True, kernel_sizes=[3,3,3], channels=[10,10,1], symmetric_mode=True):
+        super(NeighConsensus, self).__init__()
+        self.symmetric_mode = symmetric_mode
+        self.kernel_sizes = kernel_sizes
+        self.channels = channels
+        num_layers = len(kernel_sizes)
+        nn_modules = list()
+        for i in range(num_layers):
+            if i==0:
+                ch_in = 1
+            else:
+                ch_in = channels[i-1]
+            ch_out = channels[i]
+            k_size = kernel_sizes[i]
+            nn_modules.append(Conv4d(in_channels=ch_in,out_channels=ch_out,kernel_size=k_size,bias=True))
+            nn_modules.append(nn.ReLU(inplace=True))
+        self.conv = nn.Sequential(*nn_modules)        
+        if use_cuda:
+            self.conv.cuda()
+
+    def forward(self, x):
+        if self.symmetric_mode:
+            # apply network on the input and its "transpose" (swapping A-B to B-A ordering of the correlation tensor),
+            # this second result is "transposed back" to the A-B ordering to match the first result and be able to add together
+            x = self.conv(x)+self.conv(x.permute(0,1,4,5,2,3)).permute(0,1,4,5,2,3)
+            # because of the ReLU layers in between linear layers, 
+            # this operation is different than convolving a single time with the filters+filters^T
+            # and therefore it makes sense to do this.
+        else:
+            x = self.conv(x)
+        return x
+
+def MutualMatching(corr4d):
+    # mutual matching
+    batch_size,ch,fs1,fs2,fs3,fs4 = corr4d.size()
+
+    corr4d_B=corr4d.view(batch_size,fs1*fs2,fs3,fs4) # [batch_idx,k_A,i_B,j_B]
+    corr4d_A=corr4d.view(batch_size,fs1,fs2,fs3*fs4)
+
+    # get max
+    corr4d_B_max,_=torch.max(corr4d_B,dim=1,keepdim=True)
+    corr4d_A_max,_=torch.max(corr4d_A,dim=3,keepdim=True)
+
+    eps = 1e-5
+    corr4d_B=corr4d_B/(corr4d_B_max+eps)
+    corr4d_A=corr4d_A/(corr4d_A_max+eps)
+
+    corr4d_B=corr4d_B.view(batch_size,1,fs1,fs2,fs3,fs4)
+    corr4d_A=corr4d_A.view(batch_size,1,fs1,fs2,fs3,fs4)
+
+    corr4d=corr4d*(corr4d_A*corr4d_B) # parenthesis are important for symmetric output 
+    return corr4d
+
+def MutualNorm(corr4d):
+    # mutual matching
+    batch_size,ch,fs1,fs2,fs3,fs4 = corr4d.size()
+
+    corr4d_B=corr4d.view(batch_size,fs1*fs2,fs3,fs4) # [batch_idx,k_A,i_B,j_B]
+    corr4d_A=corr4d.view(batch_size,fs1,fs2,fs3*fs4)
+
+    # get max
+    corr4d_B_max,_=torch.max(corr4d_B,dim=1,keepdim=True)
+    corr4d_A_max,_=torch.max(corr4d_A,dim=3,keepdim=True)
+
+    eps = 1e-5
+    corr4d_B=corr4d_B/(corr4d_B_max+eps)
+    corr4d_A=corr4d_A/(corr4d_A_max+eps)
+
+    corr4d_B=corr4d_B.view(batch_size,1,fs1,fs2,fs3,fs4)
+    corr4d_A=corr4d_A.view(batch_size,1,fs1,fs2,fs3,fs4)
+    return (corr4d_A*corr4d_B)
+
+def maxpool4d(corr4d_hres,k_size=4):
+    slices=[]
+    for i in range(k_size):
+        for j in range(k_size):
+            for k in range(k_size):
+                for l in range(k_size):
+                    sl = corr4d_hres[:,0,i::k_size,j::k_size,k::k_size,l::k_size].unsqueeze(0)
+                    slices.append(sl)
+
+    slices=torch.cat(tuple(slices),dim=1)
+    corr4d,max_idx=torch.max(slices,dim=1,keepdim=True)
+    max_l=torch.fmod(max_idx,k_size)
+    max_k=torch.fmod(max_idx.sub(max_l).div(k_size),k_size)
+    max_j=torch.fmod(max_idx.sub(max_l).div(k_size).sub(max_k).div(k_size),k_size)
+    max_i=max_idx.sub(max_l).div(k_size).sub(max_k).div(k_size).sub(max_j).div(k_size)
+    # i,j,k,l represent the *relative* coords of the max point in the box of size k_size*k_size*k_size*k_size
+    return (corr4d,max_i,max_j,max_k,max_l)
+
+class ImMatchNet(nn.Module):
+    def __init__(self, 
+                 feature_extraction_cnn='resnet101', 
+                 feature_extraction_last_layer='',
+                 feature_extraction_model_file=None,
+                 return_correlation=False,  
+                 ncons_kernel_sizes=[3,3,3],
+                 ncons_channels=[10,10,1],
+                 normalize_features=True,
+                 train_fe=False,
+                 use_cuda=True,
+                 relocalization_k_size=0,
+                 half_precision=False,
+                 checkpoint=None,
+                 ):
+        
+        super(ImMatchNet, self).__init__()
+        # Load checkpoint
+        if checkpoint is not None and checkpoint is not '':
+            print('Loading checkpoint...')
+            checkpoint = torch.load(checkpoint, map_location=lambda storage, loc: storage)
+            checkpoint['state_dict'] = OrderedDict([(k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items()])
+            # override relevant parameters
+            print('Using checkpoint parameters: ')
+            ncons_channels=checkpoint['args'].ncons_channels
+            print('  ncons_channels: '+str(ncons_channels))
+            ncons_kernel_sizes=checkpoint['args'].ncons_kernel_sizes
+            print('  ncons_kernel_sizes: '+str(ncons_kernel_sizes))            
+
+        self.use_cuda = use_cuda
+        self.normalize_features = normalize_features
+        self.return_correlation = return_correlation
+        self.relocalization_k_size = relocalization_k_size
+        self.half_precision = half_precision
+        
+        self.FeatureExtraction = FeatureExtraction(train_fe=train_fe,
+                                                   feature_extraction_cnn=feature_extraction_cnn,
+                                                   feature_extraction_model_file=feature_extraction_model_file,
+                                                   last_layer=feature_extraction_last_layer,
+                                                   normalization=normalize_features,
+                                                   use_cuda=self.use_cuda)
+        
+        self.FeatureCorrelation = FeatureCorrelation(shape='4D',normalization=False)
+
+        self.NeighConsensus = NeighConsensus(use_cuda=self.use_cuda,
+                                             kernel_sizes=ncons_kernel_sizes,
+                                             channels=ncons_channels)
+
+        # Load weights
+        if checkpoint is not None and checkpoint is not '':
+            print('Copying weights...')
+            for name, param in self.FeatureExtraction.state_dict().items():
+                if 'num_batches_tracked' not in name:
+                    self.FeatureExtraction.state_dict()[name].copy_(checkpoint['state_dict']['FeatureExtraction.' + name])    
+            for name, param in self.NeighConsensus.state_dict().items():
+                self.NeighConsensus.state_dict()[name].copy_(checkpoint['state_dict']['NeighConsensus.' + name])
+            print('Done!')
+        
+        self.FeatureExtraction.eval()
+
+        if self.half_precision:
+            for p in self.NeighConsensus.parameters():
+                p.data=p.data.half()
+            for l in self.NeighConsensus.conv:
+                if isinstance(l,Conv4d):
+                    l.use_half=True
+                    
+    # used only for foward pass at eval and for training with strong supervision
+    def forward(self, tnf_batch): 
+        # feature extraction
+        feature_A = self.FeatureExtraction(tnf_batch['source_image'])
+        feature_B = self.FeatureExtraction(tnf_batch['target_image'])
+        if self.half_precision:
+            feature_A=feature_A.half()
+            feature_B=feature_B.half()
+            
+        # feature correlation
+        corr4d = self.FeatureCorrelation(feature_A,feature_B)
+
+        # do 4d maxpooling for relocalization
+        if self.relocalization_k_size>1:
+            corr4d,max_i,max_j,max_k,max_l=maxpool4d(corr4d,k_size=self.relocalization_k_size)
+
+        # run match processing model
+        corr4d = MutualMatching(corr4d)
+        corr4d = self.NeighConsensus(corr4d)
+        corr4d = MutualMatching(corr4d)
+        
+        if self.relocalization_k_size>1:
+            delta4d=(max_i,max_j,max_k,max_l)
+            return (corr4d,delta4d)
+        else:
+            return corr4d
+ 
+    def forward_feat(self, featA, featB, normalize=True): 
+        # feature normalization
+        if normalize:
+            feature_A = featureL2Norm(featA)
+            feature_B = featureL2Norm(featB)
+        else:
+            feature_A = featA
+            feature_B = featB
+        if self.half_precision:
+            feature_A=feature_A.half()
+            feature_B=feature_B.half()
+
+        # feature correlation
+        corr4d = self.FeatureCorrelation(feature_A,feature_B)
+        # do 4d maxpooling for relocalization
+        if self.relocalization_k_size>1:
+            corr4d,max_i,max_j,max_k,max_l=maxpool4d(corr4d,k_size=self.relocalization_k_size)
+        corr4d = MutualMatching(corr4d)
+        corr4d = self.NeighConsensus(corr4d)
+        corr4d = MutualMatching(corr4d)
+        if self.relocalization_k_size>1:
+            delta4d=(max_i,max_j,max_k,max_l)
+            return (corr4d,delta4d)
+        else:
+            return corr4d