""" "LiftFeat: 3D Geometry-Aware Local Feature Matching" """ import numpy as np import os import torch from torch import nn import torch.nn.functional as F import tqdm import math import cv2 import sys sys.path.append('/home/yepeng_liu/code_python/laiwenpeng/LiftFeat') from utils.featurebooster import FeatureBooster from utils.config import featureboost_config # from models.model_dfb import LiftFeatModel # from models.interpolator import InterpolateSparse2d # from third_party.config import featureboost_config """ foundational functions """ def simple_nms(scores, radius): """Perform non maximum suppression on the heatmap using max-pooling. This method does not suppress contiguous points that have the same score. Args: scores: the score heatmap of size `(B, H, W)`. radius: an integer scalar, the radius of the NMS window. """ def max_pool(x): return torch.nn.functional.max_pool2d( x, kernel_size=radius * 2 + 1, stride=1, padding=radius ) zeros = torch.zeros_like(scores) max_mask = scores == max_pool(scores) for _ in range(2): supp_mask = max_pool(max_mask.float()) > 0 supp_scores = torch.where(supp_mask, zeros, scores) new_max_mask = supp_scores == max_pool(supp_scores) max_mask = max_mask | (new_max_mask & (~supp_mask)) return torch.where(max_mask, scores, zeros) def top_k_keypoints(keypoints, scores, k): if k >= len(keypoints): return keypoints, scores scores, indices = torch.topk(scores, k, dim=0, sorted=True) return keypoints[indices], scores def sample_k_keypoints(keypoints, scores, k): if k >= len(keypoints): return keypoints, scores indices = torch.multinomial(scores, k, replacement=False) return keypoints[indices], scores[indices] def soft_argmax_refinement(keypoints, scores, radius: int): width = 2 * radius + 1 sum_ = torch.nn.functional.avg_pool2d( scores[:, None], width, 1, radius, divisor_override=1 ) ar = torch.arange(-radius, radius + 1).to(scores) kernel_x = ar[None].expand(width, -1)[None, None] dx = torch.nn.functional.conv2d(scores[:, None], kernel_x, padding=radius) dy = torch.nn.functional.conv2d( scores[:, None], kernel_x.transpose(2, 3), padding=radius ) dydx = torch.stack([dy[:, 0], dx[:, 0]], -1) / sum_[:, 0, :, :, None] refined_keypoints = [] for i, kpts in enumerate(keypoints): delta = dydx[i][tuple(kpts.t())] refined_keypoints.append(kpts.float() + delta) return refined_keypoints # Legacy (broken) sampling of the descriptors def sample_descriptors(keypoints, descriptors, s): b, c, h, w = descriptors.shape keypoints = keypoints - s / 2 + 0.5 keypoints /= torch.tensor( [(w * s - s / 2 - 0.5), (h * s - s / 2 - 0.5)], ).to( keypoints )[None] keypoints = keypoints * 2 - 1 # normalize to (-1, 1) args = {"align_corners": True} if torch.__version__ >= "1.3" else {} descriptors = torch.nn.functional.grid_sample( descriptors, keypoints.view(b, 1, -1, 2), mode="bilinear", **args ) descriptors = torch.nn.functional.normalize( descriptors.reshape(b, c, -1), p=2, dim=1 ) return descriptors # The original keypoint sampling is incorrect. We patch it here but # keep the original one above for legacy. def sample_descriptors_fix_sampling(keypoints, descriptors, s: int = 8): """Interpolate descriptors at keypoint locations""" b, c, h, w = descriptors.shape keypoints = keypoints / (keypoints.new_tensor([w, h]) * s) keypoints = keypoints * 2 - 1 # normalize to (-1, 1) descriptors = torch.nn.functional.grid_sample( descriptors, keypoints.view(b, 1, -1, 2), mode="bilinear", align_corners=False ) descriptors = torch.nn.functional.normalize( descriptors.reshape(b, c, -1), p=2, dim=1 ) return descriptors class UpsampleLayer(nn.Module): def __init__(self, in_channels): super().__init__() # 定义特征提取层，减少通道数同时增加特征提取能力 self.conv = nn.Conv2d(in_channels, in_channels//2, kernel_size=3, stride=1, padding=1) # 使用BN层 self.bn = nn.BatchNorm2d(in_channels//2) # 使用LeakyReLU激活函数 self.leaky_relu = nn.LeakyReLU(0.1) def forward(self, x): x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=False) x = self.leaky_relu(self.bn(self.conv(x))) return x class KeypointHead(nn.Module): def __init__(self,in_channels,out_channels): super().__init__() self.layer1=BaseLayer(in_channels,32) self.layer2=BaseLayer(32,32) self.layer3=BaseLayer(32,64) self.layer4=BaseLayer(64,64) self.layer5=BaseLayer(64,128) self.conv=nn.Conv2d(128,out_channels,kernel_size=3,stride=1,padding=1) self.bn=nn.BatchNorm2d(65) def forward(self,x): x=self.layer1(x) x=self.layer2(x) x=self.layer3(x) x=self.layer4(x) x=self.layer5(x) x=self.bn(self.conv(x)) return x class DescriptorHead(nn.Module): def __init__(self,in_channels,out_channels): super().__init__() self.layer=nn.Sequential( BaseLayer(in_channels,32), BaseLayer(32,32,activation=False), BaseLayer(32,64,activation=False), BaseLayer(64,out_channels,activation=False) ) def forward(self,x): x=self.layer(x) # x=nn.functional.softmax(x,dim=1) return x class HeatmapHead(nn.Module): def __init__(self,in_channels,mid_channels,out_channels): super().__init__() self.convHa = nn.Conv2d(in_channels, mid_channels, kernel_size=3, stride=1, padding=1) self.bnHa = nn.BatchNorm2d(mid_channels) self.convHb = nn.Conv2d(mid_channels, out_channels, kernel_size=3, stride=1, padding=1) self.bnHb = nn.BatchNorm2d(out_channels) self.leaky_relu = nn.LeakyReLU(0.1) def forward(self,x): x = self.leaky_relu(self.bnHa(self.convHa(x))) x = self.leaky_relu(self.bnHb(self.convHb(x))) x = torch.sigmoid(x) return x class DepthHead(nn.Module): def __init__(self, in_channels): super().__init__() self.upsampleDa = UpsampleLayer(in_channels) self.upsampleDb = UpsampleLayer(in_channels//2) self.upsampleDc = UpsampleLayer(in_channels//4) self.convDepa = nn.Conv2d(in_channels//2+in_channels, in_channels//2, kernel_size=3, stride=1, padding=1) self.bnDepa = nn.BatchNorm2d(in_channels//2) self.convDepb = nn.Conv2d(in_channels//4+in_channels//2, in_channels//4, kernel_size=3, stride=1, padding=1) self.bnDepb = nn.BatchNorm2d(in_channels//4) self.convDepc = nn.Conv2d(in_channels//8+in_channels//4, 3, kernel_size=3, stride=1, padding=1) self.bnDepc = nn.BatchNorm2d(3) self.leaky_relu = nn.LeakyReLU(0.1) def forward(self, x): x0 = F.interpolate(x, scale_factor=2,mode='bilinear',align_corners=False) x1 = self.upsampleDa(x) x1 = torch.cat([x0,x1],dim=1) x1 = self.leaky_relu(self.bnDepa(self.convDepa(x1))) x1_0 = F.interpolate(x1,scale_factor=2,mode='bilinear',align_corners=False) x2 = self.upsampleDb(x1) x2 = torch.cat([x1_0,x2],dim=1) x2 = self.leaky_relu(self.bnDepb(self.convDepb(x2))) x2_0 = F.interpolate(x2,scale_factor=2,mode='bilinear',align_corners=False) x3 = self.upsampleDc(x2) x3 = torch.cat([x2_0,x3],dim=1) x = self.leaky_relu(self.bnDepc(self.convDepc(x3))) x = F.normalize(x,p=2,dim=1) return x class BaseLayer(nn.Module): def __init__(self,in_channels,out_channels,kernel_size=3,stride=1,padding=1,bias=False,activation=True): super().__init__() if activation: self.layer=nn.Sequential( nn.Conv2d(in_channels,out_channels,kernel_size,stride,padding,bias=bias), nn.BatchNorm2d(out_channels,affine=False), nn.ReLU(inplace=True) ) else: self.layer=nn.Sequential( nn.Conv2d(in_channels,out_channels,kernel_size,stride,padding,bias=bias), nn.BatchNorm2d(out_channels,affine=False) ) def forward(self,x): return self.layer(x) class LiftFeatSPModel(nn.Module): default_conf = { "has_detector": True, "has_descriptor": True, "descriptor_dim": 64, # Inference "sparse_outputs": True, "dense_outputs": False, "nms_radius": 4, "refinement_radius": 0, "detection_threshold": 0.005, "max_num_keypoints": -1, "max_num_keypoints_val": None, "force_num_keypoints": False, "randomize_keypoints_training": False, "remove_borders": 4, "legacy_sampling": True, # True to use the old broken sampling } def __init__(self, featureboost_config, use_kenc=False, use_normal=True, use_cross=True): super().__init__() self.device=torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.descriptor_dim = 64 self.norm = nn.InstanceNorm2d(1) self.relu = nn.ReLU(inplace=True) self.pool = nn.MaxPool2d(kernel_size=2, stride=2) c1,c2,c3,c4,c5 = 24,24,64,64,128 self.conv1a = nn.Conv2d(1, c1, kernel_size=3, stride=1, padding=1) self.conv1b = nn.Conv2d(c1, c1, kernel_size=3, stride=1, padding=1) self.conv2a = nn.Conv2d(c1, c2, kernel_size=3, stride=1, padding=1) self.conv2b = nn.Conv2d(c2, c2, kernel_size=3, stride=1, padding=1) self.conv3a = nn.Conv2d(c2, c3, kernel_size=3, stride=1, padding=1) self.conv3b = nn.Conv2d(c3, c3, kernel_size=3, stride=1, padding=1) self.conv4a = nn.Conv2d(c3, c4, kernel_size=3, stride=1, padding=1) self.conv4b = nn.Conv2d(c4, c4, kernel_size=3, stride=1, padding=1) self.conv5a = nn.Conv2d(c4, c5, kernel_size=3, stride=1, padding=1) self.conv5b = nn.Conv2d(c5, c5, kernel_size=3, stride=1, padding=1) self.upsample4 = UpsampleLayer(c4) self.upsample5 = UpsampleLayer(c5) self.conv_fusion45 = nn.Conv2d(c5//2+c4,c4,kernel_size=3,stride=1,padding=1) self.conv_fusion34 = nn.Conv2d(c4//2+c3,c3,kernel_size=3,stride=1,padding=1) # detector self.keypoint_head = KeypointHead(in_channels=c3,out_channels=65) # descriptor self.descriptor_head = DescriptorHead(in_channels=c3,out_channels=self.descriptor_dim) # # heatmap # self.heatmap_head = HeatmapHead(in_channels=c3,mid_channels=c3,out_channels=1) # depth self.depth_head = DepthHead(c3) self.fine_matcher = nn.Sequential( nn.Linear(128, 512), nn.BatchNorm1d(512, affine=False), nn.ReLU(inplace = True), nn.Linear(512, 512), nn.BatchNorm1d(512, affine=False), nn.ReLU(inplace = True), nn.Linear(512, 512), nn.BatchNorm1d(512, affine=False), nn.ReLU(inplace = True), nn.Linear(512, 512), nn.BatchNorm1d(512, affine=False), nn.ReLU(inplace = True), nn.Linear(512, 64), ) # feature_booster self.feature_boost = FeatureBooster(featureboost_config, use_kenc=use_kenc, use_cross=use_cross, use_normal=use_normal) def feature_extract(self, x): x1 = self.relu(self.conv1a(x)) x1 = self.relu(self.conv1b(x1)) x1 = self.pool(x1) x2 = self.relu(self.conv2a(x1)) x2 = self.relu(self.conv2b(x2)) x2 = self.pool(x2) x3 = self.relu(self.conv3a(x2)) x3 = self.relu(self.conv3b(x3)) x3 = self.pool(x3) x4 = self.relu(self.conv4a(x3)) x4 = self.relu(self.conv4b(x4)) x4 = self.pool(x4) x5 = self.relu(self.conv5a(x4)) x5 = self.relu(self.conv5b(x5)) x5 = self.pool(x5) return x3,x4,x5 def fuse_multi_features(self,x3,x4,x5): # upsample x5 feature x5 = self.upsample5(x5) x4 = torch.cat([x4,x5],dim=1) x4 = self.conv_fusion45(x4) # upsample x4 feature x4 = self.upsample4(x4) x3 = torch.cat([x3,x4],dim=1) x = self.conv_fusion34(x3) return x def _unfold2d(self, x, ws = 2): """ Unfolds tensor in 2D with desired ws (window size) and concat the channels """ B, C, H, W = x.shape x = x.unfold(2, ws , ws).unfold(3, ws,ws).reshape(B, C, H//ws, W//ws, ws**2) return x.permute(0, 1, 4, 2, 3).reshape(B, -1, H//ws, W//ws) def forward1(self, x): """ input: x -> torch.Tensor(B, C, H, W) grayscale or rgb images return: feats -> torch.Tensor(B, 64, H/8, W/8) dense local features keypoints -> torch.Tensor(B, 65, H/8, W/8) keypoint logit map heatmap -> torch.Tensor(B, 1, H/8, W/8) reliability map """ with torch.no_grad(): x = x.mean(dim=1, keepdim = True) x = self.norm(x) x3,x4,x5 = self.feature_extract(x) # features fusion x = self.fuse_multi_features(x3,x4,x5) # keypoint keypoint_map = self.keypoint_head(x) # descriptor des_map = self.descriptor_head(x) # # heatmap # heatmap = self.heatmap_head(x) # import pdb;pdb.set_trace() # depth d_feats = self.depth_head(x) return des_map, keypoint_map, d_feats # return des_map, keypoint_map, heatmap, d_feats def forward2(self, descs, kpts, normals): # import pdb;pdb.set_trace() normals_feat=self._unfold2d(normals, ws=8) normals_v=normals_feat.squeeze(0).permute(1,2,0).reshape(-1,normals_feat.shape[1]) descs_v=descs.squeeze(0).permute(1,2,0).reshape(-1,descs.shape[1]) kpts_v=kpts.squeeze(0).permute(1,2,0).reshape(-1,kpts.shape[1]) descs_refine = self.feature_boost(descs_v, kpts_v, normals_v) return descs_refine def forward(self,x): M1,K1,D1=self.forward1(x) descs_refine=self.forward2(M1,K1,D1) return descs_refine,M1,K1,D1 if __name__ == "__main__": img_path=os.path.join(os.path.dirname(__file__),'../assert/ref.jpg') img=cv2.imread(img_path,cv2.IMREAD_GRAYSCALE) img=cv2.resize(img,(800,608)) import pdb;pdb.set_trace() img=torch.from_numpy(img).unsqueeze(0).unsqueeze(0).float()/255.0 img=img.cuda() if torch.cuda.is_available() else img liftfeat_sp=LiftFeatSPModel(featureboost_config).to(torch.device('cuda' if torch.cuda.is_available() else 'cpu')) des_map, keypoint_map, d_feats=liftfeat_sp.forward1(img) des_fine=liftfeat_sp.forward2(des_map,keypoint_map,d_feats) print(des_map.shape)