├── samples ├── sample_A.png └── sample_B.png ├── .gitignore ├── README.md ├── args.py ├── utils.py ├── loss.py ├── train.py └── network.py /samples/sample_A.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kimphys/UnsuperPoint.pytorch/HEAD/samples/sample_A.png -------------------------------------------------------------------------------- /samples/sample_B.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kimphys/UnsuperPoint.pytorch/HEAD/samples/sample_B.png -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Directories 2 | __pycache__/ 3 | .idea/ 4 | 5 | # OS Files 6 | .DS_Store 7 | 8 | # Temp files 9 | ._* -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Note 2 | 3 | This repository is under construction. I do not guarantee the efficiency of these codes. Please use only for reference. I'll update them soon. 4 | 5 | # Introduction 6 | 7 | This repository is a PyTorch implementation of [UnsuperPoint: End-to-end Unsupervised Interest Point Detector and Descriptor](https://arxiv.org/abs/1907.04011). 8 | 9 | # Features 10 | 11 | This repository has several features as follows 12 | 13 | * PyTorch implementation 14 | * Multi-GPU training 15 | 16 | # Citation 17 | 18 | ``` 19 | @misc{christiansen2019unsuperpoint, 20 | title={UnsuperPoint: End-to-end Unsupervised Interest Point Detector and Descriptor}, 21 | author={Peter Hviid Christiansen and Mikkel Fly Kragh and Yury Brodskiy and Henrik Karstoft}, 22 | year={2019}, 23 | eprint={1907.04011}, 24 | archivePrefix={arXiv}, 25 | primaryClass={cs.CV} 26 | } 27 | ``` 28 | 29 | 30 | 31 | 32 | -------------------------------------------------------------------------------- /args.py: -------------------------------------------------------------------------------- 1 | class args(): 2 | 3 | # hyperparameters 4 | a_usp = 1 5 | a_pos = 1 6 | a_scr = 2 7 | a_xy = 100 8 | a_desc = 0.001 9 | a_decorr = 0.03 10 | lambda_d = 250 11 | margin_p = 1 12 | margin_n = 0.2 13 | 14 | # training args 15 | epochs = 10 # "number of training epochs, default is 2" 16 | save_per_epoch = 1 17 | batch_size = 160 # "batch size for training, default is 4" 18 | dataset = "/home/oem/shk/dataset/coco_train.txt" 19 | HEIGHT = 240 20 | WIDTH = 240 21 | lr = 1e-4 #"learning rate, default is 0.001" 22 | # resume = "models/rgb.pt" # if you have, please put the path of the model like "./models/densefuse_gray.model" 23 | resume = None 24 | save_model_dir = "./models/" #"path to folder where trained model with checkpoints will be saved." 25 | workers = 20 26 | 27 | # For GPU training 28 | world_size = 4 29 | rank = 0 30 | dist_backend = 'nccl' 31 | gpu = 0,1,2,3 32 | multiprocessing_distributed = True 33 | distributed = True 34 | 35 | # For testing 36 | test_save_dir = "./" 37 | test_img = "./test_rgb.txt" 38 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | import numpy as np 3 | import torch 4 | import torch.nn as nn 5 | import torchvision 6 | from tqdm import tqdm 7 | 8 | def merge(imgs, row): 9 | B, C, H, W = imgs.shape[0], imgs.shape[1], imgs.shape[2], imgs.shape[3] 10 | 11 | q, r = divmod(B, row) 12 | 13 | if not r == 0: 14 | new_image = Image.new("RGB", (W * row, (q + 1) * H), (256,256,256)) 15 | else: 16 | new_image = Image.new("RGB", (W * row, q * H), (256,256,256)) 17 | trans = torchvision.transforms.ToPILImage() 18 | 19 | i = 0 20 | j = 0 21 | 22 | for b in range(B): 23 | 24 | if not i < row: 25 | i = 0 26 | j = j + 1 27 | 28 | leftupper = (i * W, j * H) 29 | img = trans(imgs[b]) 30 | new_image.paste(img, leftupper) 31 | 32 | i = i + 1 33 | 34 | return new_image 35 | 36 | def save_train_samples(imgs_A, imgs_B, idx, row=3): 37 | 38 | imgs_A_sample = merge(imgs_A, row) 39 | imgs_B_sample = merge(imgs_B, row) 40 | 41 | imgs_A_sample.save("{}_A.png".format(idx)) 42 | imgs_B_sample.save("{}_B.png".format(idx)) 43 | 44 | return -------------------------------------------------------------------------------- /loss.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | import math 7 | from time import time 8 | 9 | import torchvision 10 | import torchvision.transforms as transforms # Apply homography (rotation, scale, skew, perspective transforms) 11 | 12 | from args import args 13 | a_usp =args.a_usp 14 | a_pos = args.a_pos 15 | a_scr = args.a_scr 16 | a_xy = args.a_xy 17 | a_desc = args.a_desc 18 | a_decorr = args.a_decorr 19 | lambda_d = args.lambda_d 20 | margin_p = args.margin_p 21 | margin_n = args.margin_n 22 | 23 | def ApplyHomography(pA, angle): 24 | 25 | # The angle as the input is written in the degree, NOT RADIAN 26 | # So, we need to convert it from degree to radian. 27 | 28 | B, C, H, W = pA.shape[0], pA.shape[1], pA.shape[2], pA.shape[3] 29 | 30 | for b in range(B): 31 | pA_x = pA[b,0,:,:] 32 | pA_y = pA[b,1,:,:] 33 | pA[b,0,:,:] = pA_x * torch.cos(angle[b]) - pA_y * torch.sin(angle[b]) 34 | pA[b,1,:,:] = pA_x * torch.sin(angle[b]) + pA_y * torch.cos(angle[b]) 35 | 36 | return pA 37 | 38 | def PointDistance(pA, pB, angle): 39 | 40 | pA = ApplyHomography(pA, angle) 41 | 42 | x_d = pA[:,0,:,:] - pB[:,0,:,:] 43 | y_d = pA[:,1,:,:] - pB[:,1,:,:] 44 | 45 | Distance = torch.sqrt(x_d * x_d + y_d * y_d) 46 | 47 | PositionPoint = torch.sum(Distance) 48 | 49 | return Distance, PositionPoint 50 | 51 | def PointLoss(angle, a_pos, a_scr, pos_A, pos_B, scr_A, scr_B): 52 | 53 | mse = nn.MSELoss() 54 | 55 | Distance, PositionPoint = PointDistance(pos_A, pos_B, angle) 56 | ScorePoint = mse(scr_A, scr_B) 57 | 58 | ## Repeatable point 59 | scr_mean = (scr_A + scr_B) / 2 60 | RepeatPoint = torch.sum(scr_mean * (Distance - PositionPoint)) 61 | 62 | return a_pos * PositionPoint + a_scr * ScorePoint + RepeatPoint 63 | 64 | def PointPredLoss(a_xy, pos_r): 65 | 66 | B, C, H, W = pos_r.shape[0], pos_r.shape[1], pos_r.shape[2], pos_r.shape[3] 67 | 68 | pos_X = pos_r[:,0,:,:].view(-1, H * W) 69 | pos_Y = pos_r[:,1,:,:].view(-1, H * W) 70 | 71 | pos_X_sorted, pos_X_indices = torch.sort(pos_X) 72 | pos_Y_sorted, pos_Y_indices = torch.sort(pos_Y) 73 | 74 | loss_x = torch.pow((pos_X - torch.true_divide((pos_X_indices - 1),(H * W - 1))),2) 75 | loss_y = torch.pow((pos_Y - torch.true_divide((pos_Y_indices - 1),(H * W - 1))),2) 76 | 77 | return torch.sum(loss_x + loss_y) 78 | 79 | def DescLoss(lambda_d, pos_A, pos_B, angle, desc_A, desc_B, margin_p, margin_n): 80 | 81 | B, C, H, W = desc_A.shape[0], desc_A.shape[1], desc_A.shape[2], desc_A.shape[3] 82 | 83 | pos_A = ApplyHomography(pos_A, angle) 84 | 85 | pos_A = pos_A.view(B, 2, H * W).transpose(1,2) 86 | pos_B = pos_B.view(B, 2, H * W).transpose(1,2) 87 | pos_dist = torch.cdist(pos_A,pos_B) 88 | 89 | 90 | desc_A = desc_A.view(B, C, H * W).transpose(1,2) 91 | desc_B = desc_B.view(B, C, H * W) 92 | 93 | product = torch.bmm(desc_A, desc_B) 94 | 95 | ones = torch.ones(pos_dist.shape) 96 | zeros_dist = torch.zeros(pos_dist.shape) 97 | zeros_p = torch.zeros(product.shape) 98 | 99 | if args.gpu is not None: 100 | ones = ones.cuda(args.gpu, non_blocking=True) 101 | zeros_dist = zeros_dist.cuda(args.gpu, non_blocking=True) 102 | zeros_p = zeros_p.cuda(args.gpu, non_blocking=True) 103 | 104 | ConstantMatrix = torch.where(pos_dist >= 8, ones, zeros_dist) 105 | product_p = torch.where(product < margin_p, product - margin_p, zeros_p) 106 | product_n = torch.where(product > margin_n, product - margin_n, zeros_p) 107 | 108 | loss = torch.sum(lambda_d * ConstantMatrix * product_p + (1 - ConstantMatrix) * product_n) 109 | 110 | return loss 111 | 112 | # Under construction 113 | def DecorreDescLoss(desc_A, desc_B): 114 | 115 | loss = 0 116 | 117 | return loss 118 | 119 | def ComputeLoss(angle, scr_A, scr_B, pos_A_r, pos_B_r, pos_A, pos_B, desc_A, desc_B): 120 | 121 | p_loss = PointLoss(angle, a_pos, a_scr, pos_A, pos_B, scr_A, scr_B) 122 | 123 | pred_loss = PointPredLoss(a_xy, pos_A_r) + PointPredLoss(a_xy, pos_B_r) 124 | 125 | d_loss = DescLoss(lambda_d, pos_A, pos_B, angle, desc_A, desc_B, margin_p, margin_n) 126 | 127 | deco_loss = DecorreDescLoss(desc_A, desc_B) 128 | 129 | return a_usp * p_loss + a_xy * pred_loss + a_desc * d_loss + a_decorr * deco_loss 130 | 131 | if __name__ == "__main__": 132 | angle = torch.FloatTensor([30, -10]) 133 | scr_A = torch.rand(2,1,32,32) 134 | scr_B = torch.rand(2,1,32,32) 135 | pos_A_r = torch.rand(2,2,32,32) 136 | pos_B_r = torch.rand(2,2,32,32) 137 | pos_A = torch.rand(2,2,32,32) 138 | pos_B = torch.rand(2,2,32,32) 139 | desc_A = torch.rand(2,256,32,32) 140 | desc_B = torch.rand(2,256,32,32) 141 | 142 | loss = ComputeLoss(angle, scr_A, scr_B, pos_A_r, pos_B_r, pos_A, pos_B, desc_A, desc_B) 143 | 144 | print(loss) -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | 4 | import torch.distributed as dist 5 | import torch.multiprocessing as mp 6 | from torch.nn.parallel import DistributedDataParallel 7 | 8 | import torch 9 | import torchvision 10 | import torchvision.transforms as transforms 11 | 12 | import torch.optim as optim 13 | from torch.autograd import Variable 14 | 15 | from torch.utils.data import Dataset, DataLoader 16 | 17 | import torch.nn as nn 18 | import torch.nn.functional as F 19 | 20 | import random 21 | import numpy as np 22 | from tqdm import tqdm 23 | from PIL import Image 24 | from time import time 25 | import math 26 | 27 | from network import Unsuperpoint 28 | from args import args 29 | from loss import ComputeLoss 30 | from utils import save_train_samples 31 | 32 | is_cuda = torch.cuda.is_available() 33 | 34 | os.environ['MASTER_ADDR'] = '127.0.0.1' 35 | os.environ['MASTER_PORT'] = '29500' 36 | 37 | class MyTrainDataset(Dataset): 38 | def __init__(self, img_path_file): 39 | f = open(img_path_file, 'r') 40 | img_list = f.read().splitlines() 41 | f.close() 42 | 43 | self.img_list = img_list 44 | 45 | def __getitem__(self, index): 46 | img_1 = Image.open(self.img_list[index]).convert('RGB') 47 | 48 | angle = random.randint(-30,30) 49 | img_2 = transforms.functional.rotate(img_1,angle) 50 | 51 | custom_transform = transforms.Compose([transforms.Resize((args.HEIGHT,args.WIDTH)), 52 | transforms.ToTensor()]) 53 | 54 | img_1 = custom_transform(img_1) 55 | img_2 = custom_transform(img_2) 56 | 57 | return img_1, img_2, angle * math.pi / 180 58 | 59 | def __len__(self): 60 | 61 | return len(self.img_list) 62 | 63 | def setup(rank, world_size): 64 | dist.init_process_group("nccl", rank=rank, world_size=world_size) 65 | 66 | def cleanup(): 67 | dist.destropy_process_group 68 | 69 | def main(): 70 | args.distributed = args.world_size > 1 or args.multiprocessing_distributed 71 | 72 | if args.gpu is not None: 73 | if not len(args.gpu) > torch.cuda.device_count(): 74 | ngpus_per_node = len(args.gpu) 75 | else: 76 | print("We will use all available GPUs") 77 | ngpus_per_node = torch.cuda.device_count() 78 | else: 79 | ngpus_per_node = torch.cuda.device_count() 80 | 81 | if args.multiprocessing_distributed: 82 | args.world_size = ngpus_per_node * args.world_size 83 | mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args)) 84 | else: 85 | main_worker(args.gpu, ngpus_per_node, args) 86 | 87 | def main_worker(gpu, ngpus_per_node, args): 88 | args.gpu = gpu 89 | 90 | if args.gpu is not None: 91 | print("Use GPU: {} for training".format(args.gpu)) 92 | 93 | if args.distributed: 94 | if args.multiprocessing_distributed: 95 | args.rank = args.rank * ngpus_per_node + gpu 96 | dist.init_process_group(backend=args.dist_backend, world_size=args.world_size, rank=args.rank) 97 | 98 | model = Unsuperpoint() 99 | optimizer = torch.optim.Adam(model.parameters(), args.lr) 100 | 101 | epoch = 0 102 | 103 | if not torch.cuda.is_available(): 104 | print('using CPU, this will be slow') 105 | elif args.distributed: 106 | if args.gpu is not None: 107 | torch.cuda.set_device(args.gpu) 108 | model.cuda(args.gpu) 109 | args.batch_size = int(args.batch_size / ngpus_per_node) 110 | args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node) 111 | model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu]) 112 | else: 113 | model.cuda() 114 | model = torch.nn.parallel.DistributedDataParallel(model) 115 | elif args.gpu is not None: 116 | torch.cuda.set_device(args.gpu) 117 | model = model.cuda(args.gpu) 118 | else: 119 | model = torch.nn.DataParallel(model).cuda() 120 | 121 | if args.resume: 122 | if args.gpu is None: 123 | checkpoint = torch.load(args.resume) 124 | else: 125 | loc = 'cuda:{}'.format(args.gpu) 126 | checkpoint = torch.load(args.resume, map_location=loc) 127 | model.load_state_dict(checkpoint['model_state_dict']) 128 | optimizer.load_state_dict(checkpoint['optimizer_state_dict']) 129 | epoch = checkpoint['epoch'] 130 | loss = checkpoint['loss'] 131 | else: 132 | # print("Training from scratch") 133 | pass 134 | 135 | img_path_file = args.dataset 136 | 137 | trainloader = DataLoader(MyTrainDataset(img_path_file), batch_size=args.batch_size, pin_memory=False, shuffle=True, num_workers=args.workers) 138 | 139 | for ep in range(epoch, args.epochs): 140 | 141 | pbar = tqdm(trainloader) 142 | idx = 0 143 | 144 | for imgs_A, imgs_B, angle in pbar: 145 | # for imgs_A, imgs_B, angle in trainloader: 146 | 147 | ''' 148 | if ep == 0 and idx ==0: 149 | save_train_samples(imgs_A,imgs_B,idx) 150 | ''' 151 | 152 | if args.gpu is not None: 153 | imgs_A = imgs_A.cuda(args.gpu, non_blocking=True) 154 | imgs_B = imgs_B.cuda(args.gpu, non_blocking=True) 155 | 156 | optimizer.zero_grad() 157 | 158 | preds_A, preds_B = model(imgs_A,imgs_B) 159 | 160 | scr_A, pos_A_r, pos_A, desc_A = preds_A[0], preds_A[1], preds_A[2], preds_A[3] 161 | scr_B, pos_B_r, pos_B, desc_B = preds_B[0], preds_B[1], preds_B[2], preds_B[3] 162 | 163 | loss = ComputeLoss(angle, scr_A, scr_B, pos_A_r, pos_B_r, pos_A, pos_B, desc_A, desc_B) 164 | 165 | loss.backward() 166 | optimizer.step() 167 | 168 | idx += 1 169 | if torch.cuda.current_device() == 0: 170 | print("GPU{} Total_loss:".format(torch.cuda.current_device()), loss) 171 | 172 | if (ep + 1) % args.save_per_epoch == 0: 173 | # Save model 174 | torch.save({ 175 | 'epoch': ep, 176 | 'model_state_dict': model.state_dict(), 177 | 'optimizer_state_dict': optimizer.state_dict(), 178 | 'loss': loss 179 | }, args.save_model_dir + 'ckpt_{}.pt'.format(ep)) 180 | 181 | 182 | print('Finished training') 183 | 184 | if __name__ == "__main__": 185 | main() -------------------------------------------------------------------------------- /network.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | class Mish(nn.Module): # https://github.com/digantamisra98/Mish 7 | def forward(self, x): 8 | return x * F.softplus(x).tanh() 9 | 10 | class Swish(nn.Module): 11 | def forward(self, x): 12 | return x * torch.sigmoid(x) 13 | 14 | # Convolution operation 15 | class ConvLayer(torch.nn.Module): 16 | def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, pad=1, bn=True, activation_type='leaky'): 17 | super(ConvLayer, self).__init__() 18 | self.conv2d = nn.Conv2d(in_channels=in_channels, 19 | out_channels=out_channels, 20 | kernel_size=kernel_size, 21 | stride=stride, 22 | padding=pad) 23 | if bn: 24 | self.batchNorm = nn.BatchNorm2d(out_channels, momentum=0.03, eps=1e-4) 25 | else: 26 | self.batchNorm = None 27 | 28 | self.activation_type = activation_type 29 | if activation_type == 'leaky': 30 | self.activation = nn.LeakyReLU(0.1, inplace=True) 31 | elif activation_type == 'sigmoid': 32 | self.activation = nn.Sigmoid() 33 | elif activation_type == 'mish': 34 | self.activation = Mish() 35 | elif activation_type == 'swish': 36 | self.activation = Swish() 37 | else: 38 | self.activation = None 39 | 40 | def forward(self, x): 41 | x = self.conv2d(x) 42 | 43 | if self.batchNorm is not None: 44 | x = self.batchNorm(x) 45 | 46 | if self.activation is not None: 47 | x = self.activation(x) 48 | 49 | return x 50 | 51 | class Backbone(nn.Module): 52 | def __init__(self, in_channels): 53 | super(Backbone, self).__init__() 54 | 55 | channels = [in_channels,32,32,64,64,128,128,256,256] 56 | 57 | self.conv2d_1 = ConvLayer(in_channels=channels[0],out_channels=channels[1]) 58 | self.conv2d_2 = ConvLayer(in_channels=channels[1],out_channels=channels[2]) 59 | 60 | self.conv2d_3 = ConvLayer(in_channels=channels[2],out_channels=channels[3]) 61 | self.conv2d_4 = ConvLayer(in_channels=channels[3],out_channels=channels[4]) 62 | 63 | self.conv2d_5 = ConvLayer(in_channels=channels[4],out_channels=channels[5]) 64 | self.conv2d_6 = ConvLayer(in_channels=channels[5],out_channels=channels[6]) 65 | 66 | self.conv2d_7 = ConvLayer(in_channels=channels[6],out_channels=channels[7]) 67 | self.conv2d_8 = ConvLayer(in_channels=channels[7],out_channels=channels[8]) 68 | 69 | self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2) 70 | 71 | def forward(self, x): 72 | 73 | x = self.conv2d_1(x) 74 | x = self.conv2d_2(x) 75 | x = self.maxpool(x) 76 | 77 | x = self.conv2d_3(x) 78 | x = self.conv2d_4(x) 79 | x = self.maxpool(x) 80 | 81 | x = self.conv2d_5(x) 82 | x = self.conv2d_6(x) 83 | x = self.maxpool(x) 84 | 85 | x = self.conv2d_7(x) 86 | x = self.conv2d_8(x) 87 | 88 | return x 89 | 90 | class Score(nn.Module): 91 | def __init__(self, in_channels): 92 | super(Score, self).__init__() 93 | self.conv_1 = ConvLayer(in_channels=in_channels,out_channels=in_channels) 94 | self.conv_2 = ConvLayer(in_channels=in_channels,out_channels=1, activation_type='sigmoid') 95 | 96 | def forward(self, x): 97 | x = self.conv_1(x) 98 | x = self.conv_2(x) 99 | return x 100 | 101 | def MappingXY(relative, downsampling): 102 | 103 | # relative -> B 2 H W 104 | H = relative.shape[2] 105 | W = relative.shape[3] 106 | 107 | pmap = relative.clone().detach() 108 | 109 | for i in range(H): 110 | for j in range(W): 111 | pmap[:,0,i,j] = (i + relative[:,0,i,j]) * downsampling 112 | pmap[:,1,i,j] = (j + relative[:,1,i,j]) * downsampling 113 | 114 | return pmap 115 | 116 | class PositionXY(nn.Module): 117 | def __init__(self, in_channels): 118 | super(PositionXY, self).__init__() 119 | self.conv_1 = ConvLayer(in_channels=in_channels,out_channels=in_channels) 120 | self.conv_2 = ConvLayer(in_channels=in_channels,out_channels=2, activation_type='sigmoid') 121 | 122 | def forward(self, x, downsampling=8): 123 | x = self.conv_1(x) 124 | x = self.conv_2(x) 125 | 126 | pmap = MappingXY(x, downsampling) 127 | 128 | return pmap, x 129 | 130 | class Descriptor(nn.Module): 131 | def __init__(self, in_channels): 132 | super(Descriptor, self).__init__() 133 | self.conv_1 = ConvLayer(in_channels=in_channels,out_channels=in_channels) 134 | self.conv_2 = ConvLayer(in_channels=in_channels,out_channels=in_channels, activation_type='None') 135 | 136 | def forward(self, x): 137 | x = self.conv_1(x) 138 | x = self.conv_2(x) 139 | 140 | return x 141 | 142 | # Under construction 143 | def interpolation(position, descriptor, h , w): 144 | 145 | interpolate = descriptor 146 | 147 | return interpolate 148 | 149 | 150 | class Unsuperpoint(nn.Module): 151 | def __init__(self, in_channels=3, s_channels=256, p_channels=256, d_channels=256): 152 | super(Unsuperpoint, self).__init__() 153 | 154 | self.backbone = Backbone(in_channels=in_channels) 155 | self.score = Score(in_channels=s_channels) 156 | self.position = PositionXY(in_channels=p_channels) 157 | self.descriptor = Descriptor(in_channels=d_channels) 158 | 159 | def forward_once(self, x): 160 | h, w = x.shape[2], x.shape[3] 161 | x = self.backbone(x) 162 | 163 | x1 = self.score(x) 164 | x2, x2_r = self.position(x) 165 | x3 = self.descriptor(x) 166 | x3 = interpolation(x2, x3, h, w) 167 | 168 | return x1, x2_r, x2, x3 169 | 170 | def forward(self, input1, input2): 171 | output1 = self.forward_once(input1) 172 | output2 = self.forward_once(input2) 173 | 174 | return output1, output2 175 | 176 | if __name__ == "__main__": 177 | inputs1 = torch.autograd.Variable(torch.rand(3,3,256,256)) 178 | inputs2 = torch.autograd.Variable(torch.rand(3,3,256,256)) 179 | 180 | print("Input1: {}, Input2: {}".format(inputs1.shape, inputs2.shape)) 181 | 182 | model = Unsuperpoint() 183 | 184 | preds_1, preds_2 = model(inputs1, inputs2) 185 | 186 | print("---For input 1---") 187 | print("Score: ", preds_1[0].shape) 188 | print("Relative position: ", preds_1[1].shape) 189 | print("Mapping position: ", preds_1[2].shape) 190 | print("Descriptor: ", preds_1[3].shape) 191 | print("---For input 2---") 192 | print("Score: ", preds_2[0].shape) 193 | print("Relative position: ", preds_2[1].shape) 194 | print("Mapping position: ", preds_2[2].shape) 195 | print("Descriptor: ", preds_2[3].shape) --------------------------------------------------------------------------------