import cv2 import math import matplotlib.pyplot as plt import numpy as np import torch import torch._utils import torch.nn as nn import torch.nn.functional as F from uot import UOTHead from netvlad import NetVLAD, NetVLADLoupe from ALIKE.alike import configs from ALIKE.alnet import ALNet from BEVNet import RICNN, EncodePosition, RIAvgpool2d, RIMaxpool2d import tools def simple_nms(scores, nms_radius=2, itertation=2, mode='1'): """ Fast Non-maximum suppression to remove nearby points """ assert (nms_radius >= 0) if mode == 'ri': max_pool = RIMaxpool2d(kernel_size=nms_radius * 2 + 1, stride=1, padding=nms_radius) else: max_pool = nn.MaxPool2d(kernel_size=nms_radius * 2 + 1, stride=1, padding=nms_radius) zeros = torch.zeros_like(scores) max_mask = scores == max_pool(scores) for _ in range(itertation): 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) class BEVHead(nn.Module): def __init__(self, alnet='alike-n', iter=5, num_kpt=100, cluster_num=16, vlad_size=256): super(BEVHead, self).__init__() cfg = configs[alnet] self.feature_extractor = ALNet(c1=cfg['c1'], c2=cfg['c2'], c3=cfg['c3'], c4=cfg['c4'], dim=cfg['dim'], single_head=cfg['single_head']) self.feature_size = int(self.feature_extractor.feature_size) self.select = 'maxpool' self.num_kpt = num_kpt self.ep = EncodePosition(feature_size=self.feature_size) self.uot = UOTHead(nb_iter=iter,name='original') self.netvlad_bev = NetVLAD(self.feature_size, cluster_num) # state_dict=torch.load('/data4/caodanyang/results/FUSIONLCD/bev_07250/models/checkpoint_049.pth.tar', map_location='cpu')['model'] # state_dict_new={} # for k,v in state_dict.items(): # state_dict_new[k[4:]]=v # self.load_state_dict(state_dict_new) # for param in self.parameters(): # param.requires_grad = False def forward(self, batch_dict): assert type(batch_dict) is dict, 'Input should be a dict' bev = batch_dict['bev'] guider = (bev[:, 2:3] > 0).float() b, c, h_bev, w_bev = bev.shape x = bev[:, 0:3, :, :] points = bev[:, 3:7, :, :] # xyzi points[:, 2] = 0 points[:, 3] = 1 score_bev, feature_bev = self.feature_extractor(x) score_bev = score_bev * guider if self.select == 'avgpool': avgpool = RIAvgpool2d(kernel_size=5, stride=4, padding=1) grid = np.array(np.meshgrid(np.arange(h_bev), np.arange(w_bev))).swapaxes(0, 2) grid = torch.from_numpy(grid).to(x.device).permute(2, 0, 1).unsqueeze(0).repeat(b, 1, 1, 1) score_bev_avg = avgpool(score_bev) grid_avg = avgpool(grid.float() * score_bev) / (score_bev_avg + 1e-8) grid_avg = torch.round(grid_avg).long().permute(0, 2, 3, 1) points_avg = avgpool(score_bev * points) / (score_bev_avg + 1e-8) feature_bev_avg = avgpool(feature_bev * score_bev) / (score_bev_avg + 1e-8) score_bev = score_bev.view(b, h_bev, w_bev) score_bev_avg = score_bev_avg.squeeze(1) kpts = [] feas_kpt = [] pixels_kpt = [] # cnt=0 for i in range(b): uv = list(torch.where(score_bev_avg[i] > 0)) num_kpt = int(self.num_kpt) if num_kpt == 0: print('NO BEV key point') exit() while len(uv[0]) < num_kpt: uv[0] = torch.cat([uv[0], uv[0][:(num_kpt - len(uv[0]))]]) uv[1] = torch.cat([uv[1], uv[1][:(num_kpt - len(uv[1]))]]) score_bev0 = score_bev_avg[i, uv[0], uv[1]] score_bev1, idx = torch.topk(score_bev0, k=self.num_kpt) # cnt=max(cnt,len(uv[0])) # idx=torch.arange(len(uv[0])).to(x.device) pc = points_avg[i, :, uv[0], uv[1]].permute(1, 0) # pc = torch.cat([pc, pc * 0], dim=1) kpt = pc[idx] fea_kpt = feature_bev_avg[i, :, uv[0][idx], uv[1][idx]] pixel_kpt = grid_avg[i, uv[0][idx], uv[1][idx]] pixels_kpt.append(pixel_kpt) kpts.append(kpt.unsqueeze(0)) feas_kpt.append(fea_kpt.unsqueeze(0)) else: score_bev_max = simple_nms(score_bev, nms_radius=3) score_bev = score_bev.view(b, h_bev, w_bev) score_bev_max = score_bev_max.view(b, h_bev, w_bev) kpts = [] feas_kpt = [] pixels_kpt = [] for i in range(b): uv = list(torch.where((score_bev[i] == score_bev_max[i]) & (score_bev[i] > 0))) num_kpt = int(self.num_kpt) if num_kpt == 0: print('NO BEV key point') exit() while len(uv[0]) < num_kpt: uv[0] = torch.cat([uv[0], uv[0][:(num_kpt - len(uv[0]))]]) uv[1] = torch.cat([uv[1], uv[1][:(num_kpt - len(uv[1]))]]) score_bev0 = score_bev[i, uv[0], uv[1]] # sc0 = score_bev0.cpu().detach().numpy() score_bev1, idx = torch.topk(score_bev0, k=self.num_kpt) pc = points[i, :, uv[0], uv[1]].permute(1, 0) # pc = torch.cat([pc, pc * 0], dim=1) kpt = pc[idx] fea_kpt = feature_bev[i, :, uv[0][idx], uv[1][idx]] pixel_kpt = torch.cat([uv[0][idx], uv[1][idx]]).view(2, -1).T pixels_kpt.append(pixel_kpt.unsqueeze(0)) kpts.append(kpt.unsqueeze(0)) feas_kpt.append(fea_kpt.unsqueeze(0)) # kpts1=torch.zeros((b,cnt,kpt.shape[1])).to(x.device) # feas_kpt1=torch.zeros((b,fea_kpt.shape[0],cnt)).to(x.device) # for i in range(b): # kpts1[i,:kpts[i].shape[1]]=kpts[i].squeeze(0) # feas_kpt1[i,:,:feas_kpt[i].shape[2]]=feas_kpt[i].squeeze(0) kpts = torch.cat(kpts) feas_kpt = torch.cat(feas_kpt) pixels_kpt = torch.cat(pixels_kpt) if hasattr(self, 'ep'): feas_kpt = self.ep(kpts, feas_kpt) batch_dict['pixels_kpt'] = pixels_kpt batch_dict['score_bev'] = score_bev batch_dict['fea_kpt_original'] = feas_kpt batch_dict['fea_bev'] = feature_bev batch_dict['key_points'] = kpts if hasattr(self, 'netvlad_bev'): try: vlad_bev = self.netvlad_bev(feas_kpt.transpose(1, 2).contiguous()) except: vlad_bev = self.netvlad_bev(feas_kpt.unsqueeze(3)) batch_dict['vlad_bev'] = vlad_bev if ('pose_to_frame' in batch_dict.keys()) and (hasattr(self, 'uot')): self.uot(batch_dict) #################################### show bev and kpt ############################################ if 0: for i in range(b): bevshow = x[i].permute(1, 2, 0).cpu().detach().numpy() bevshow = np.ascontiguousarray(bevshow[:, :, 0:3] * 255, dtype=np.uint8) bevshow1 = bevshow.copy() bevshow1[:, 1] = [255, 255, 255] for j in range(kpt.shape[0]): center = (int(uv[1][idx[j]].cpu().detach().numpy()), int(uv[0][idx[j]].cpu().detach().numpy())) cv2.circle(bevshow1, center, 2, (0, 0, 255), -1, cv2.LINE_AA) bevshow2 = np.hstack((bevshow, bevshow1)) # cv2.namedWindow('2', cv2.WINDOW_NORMAL) # cv2.imshow('2', bevshow2) # cv2.waitKey(0) fig = plt.figure() plt.imshow(bevshow2) plt.show() ######################################################################################################### #################################### show match ############################################ if 0: for i in range(b // 2): kpt1 = kpts[i] pose_to_frame = batch_dict['pose_to_frame'][i] # pose_to_frame = batch_dict['transformation'][i] # pose_to_frame = torch.cat((pose_to_frame, torch.tensor([0, 0, 0, 1]).view(1, 4).to(pose_to_frame.device))) kpt1 = (pose_to_frame @ kpt1.permute(1, 0)).permute(1, 0) kpt2 = kpts[i + b // 2] bev1 = batch_dict['bev'][i][0:3].permute(1, 2, 0) bev1 = np.ascontiguousarray(bev1.cpu().detach().numpy() * 255, dtype=np.uint8) bev2 = batch_dict['bev'][i + b // 2][0:3].permute(1, 2, 0) bev2 = np.ascontiguousarray(bev2.cpu().detach().numpy() * 255, dtype=np.uint8) pixel1 = pixels_kpt[i].cpu().detach().numpy() pixel2 = pixels_kpt[i + b // 2].cpu().detach().numpy() fea1 = feas_kpt[i].permute(1, 0).cpu().detach().numpy() fea2 = feas_kpt[i + b // 2].permute(1, 0).cpu().detach().numpy() idx1, idx2, dis = tools.nn_match(fea1, fea2, 'cosine') # idx11, idx21, dis1 = tools.nn_match(kpt1, kpt2, 'euclidean') # idx1 = idx1[dis < 0.1] # idx2 = idx2[dis < 0.1] h, w, _ = bev1.shape img = np.hstack((bev1, bev2)) img[:, w] = [255, 255, 255] tp = 0 img1 = img.copy() for j in range(len(pixel1)): center1 = (int(pixel1[j, 1]), int(pixel1[j, 0])) center2 = (int(pixel2[j, 1]) + w, int(pixel2[j, 0])) cv2.circle(img, center1, 2, (155, 155, 155), -1, cv2.LINE_AA) cv2.circle(img, center2, 2, (155, 155, 155), -1, cv2.LINE_AA) for j in range(len(idx1)): center1 = (int(pixel1[idx1[j], 1]), int(pixel1[idx1[j], 0])) center2 = (int(pixel2[idx2[j], 1]) + w, int(pixel2[idx2[j], 0])) dis_kpt = (kpt1[idx1[j]] - kpt2[idx2[j]]).norm(p=2) if dis_kpt < 2: tp = tp + 1 cv2.line(img, center1, center2, (0, 166, 0), 1, cv2.LINE_AA) else: cv2.line(img, center1, center2, (0, 0, 188), 1, cv2.LINE_AA) cv2.circle(img, center1, 2, (255, 255, 255), -1, cv2.LINE_AA) cv2.circle(img, center2, 2, (255, 255, 255), -1, cv2.LINE_AA) # print(np.arccos(pose_to_frame.cpu().detach().numpy()[0, 0]) / np.pi * 180, (tp / len(idx1))) img2 = np.vstack((img1, img)) img2[h, :] = [255, 255, 255] cv2.namedWindow('bev match %.3f,%.1fdeg' % (tp / len(idx1), np.arccos(pose_to_frame.cpu().detach().numpy()[0, 0]) / np.pi * 180)) cv2.imshow('bev match %.3f,%.1fdeg' % (tp / len(idx1), np.arccos(pose_to_frame.cpu().detach().numpy()[0, 0]) / np.pi * 180), img2) cv2.waitKey(0) ##################################################################################################### ############################################ ICP ################################################## if 0: import open3d as o3d for i in range(b // 2): pose_to_frame = batch_dict['pose_to_frame'][i].cpu().detach().numpy() print('angle', np.arccos(pose_to_frame[0, 0]) / 3.14 * 180) transformation = batch_dict['transformation'][i].cpu().detach().numpy() transformation = np.vstack((transformation, [0, 0, 0, 1])) scan1 = batch_dict['scan_query'][i].cpu().detach().numpy() scan2 = batch_dict['scan_positive'][i].cpu().detach().numpy() pcd1 = o3d.geometry.PointCloud() pcd1.points = o3d.utility.Vector3dVector(scan1[:, :3]) pcd1.colors = o3d.utility.Vector3dVector([[0, 0, 1] for i in range(len(pcd1.points))]) pcd11 = o3d.geometry.PointCloud() pcd11.points = o3d.utility.Vector3dVector(scan1[:, :3]) pcd11.colors = o3d.utility.Vector3dVector([[0, 0, 1] for i in range(len(pcd1.points))]) pcd2 = o3d.geometry.PointCloud() pcd2.points = o3d.utility.Vector3dVector(scan2[:, :3]) pcd2.colors = o3d.utility.Vector3dVector([[0, 1, 0] for i in range(len(pcd2.points))]) icp_config = o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=200, relative_fitness=1e-6, relative_rmse=1e-6) trans_init = transformation threshold = 2 estimation_method = o3d.pipelines.registration.TransformationEstimationPointToPoint() registration_result = o3d.pipelines.registration.registration_icp(pcd1, pcd2, threshold, trans_init, estimation_method, icp_config) # 将待配准点云应用变换 pcd1.transform(registration_result.transformation) vis1 = o3d.visualization.Visualizer() vis1.create_window(window_name='registration', width=600, height=600) # 创建窗口 render_option: o3d.visualization.RenderOption = vis1.get_render_option() # 设置点云渲染参数 render_option.background_color = np.array([1, 1, 1]) # 设置背景色（这里为黑色） render_option.point_size = 2 # 设置渲染点的大小 vis1.add_geometry(pcd11) vis1.run() vis2 = o3d.visualization.Visualizer() vis2.create_window(window_name='registration', width=600, height=600) # 创建窗口 render_option: o3d.visualization.RenderOption = vis2.get_render_option() # 设置点云渲染参数 render_option.background_color = np.array([1, 1, 1]) # 设置背景色（这里为黑色） render_option.point_size = 2 # 设置渲染点的大小 vis2.add_geometry(pcd2) vis2.run() vis = o3d.visualization.Visualizer() vis.create_window(window_name='registration', width=600, height=600) # 创建窗口 render_option: o3d.visualization.RenderOption = vis.get_render_option() # 设置点云渲染参数 render_option.background_color = np.array([1, 1, 1]) # 设置背景色（这里为黑色） render_option.point_size = 2 # 设置渲染点的大小 vis.add_geometry(pcd1) vis.add_geometry(pcd2) vis.run() ####################################################################################################### return batch_dict class ImgHead(nn.Module): def __init__(self, alnet='alike-n', num_kpt=150, cluster_num=0,vlad_size=256): super(ImgHead, self).__init__() cfg = configs[alnet] self.feature_extractor = ALNet(c1=cfg['c1'], c2=cfg['c2'], c3=cfg['c3'], c4=cfg['c4'], dim=cfg['dim'], single_head=cfg['single_head']) self.feature_size = int(self.feature_extractor.feature_size) # try: # model_path = cfg['model_path'] # except: # model_path = '' # if model_path != '': # state_dict = torch.load(model_path) # self.feature_extractor.load_state_dict(state_dict) # for param in self.feature_extractor.parameters(): # param.requires_grad = False if num_kpt>0: self.num_kpt = num_kpt def forward(self, batch_dict): x = batch_dict['img'][:, 0:3].float() / 255.0 # x=x[:,:,:,384:768,] # pixels = batch_dict['img'][:, 3:5] b, c, h, w = x.shape pixel_features = [] kpts = [] scores = [] score_img, feature_img = self.feature_extractor(x) # feature_img=feature_img*0 if hasattr(self,'num_kpt') : score_img = simple_nms(score_img, 2, 2) s_thr = 0.1 for i in range(b): score_global1 = score_img[i, 0] values, indices = torch.topk(score_global1.view(-1), k=self.num_kpt, dim=0, largest=True) if torch.max(values) < s_thr: print('0 pixel') exit() num_low_value = torch.sum(values < s_thr) if num_low_value > 0: indices1 = indices.clone() indices1[(self.num_kpt - num_low_value):] = indices[:num_low_value] indices = indices1 row = torch.div(indices, score_global1.shape[1], rounding_mode='trunc') col = indices % score_global1.shape[1] pixel_feature = feature_img[i:i + 1, :, row, col] pixel_features.append(pixel_feature) kpts.append(torch.cat([row.view(1, -1, 1), col.view(1, -1, 1)], dim=2)) scores.append(values.view(1, -1)) pixel_features = torch.cat(pixel_features) kpts = torch.cat(kpts) scores = torch.cat(scores) #################################### show match ############################################ if 0: for i in range(b // 2): img1 = batch_dict['img'][i][0:3].permute(1, 2, 0) img1 = np.ascontiguousarray(img1.cpu().detach().numpy(), dtype=np.uint8) img2 = batch_dict['img'][i + b // 2][0:3].permute(1, 2, 0) img2 = np.ascontiguousarray(img2.cpu().detach().numpy(), dtype=np.uint8) pixel1 = kpts[i].cpu().detach().numpy() pixel2 = kpts[i + b // 2].cpu().detach().numpy() fea1 = pixel_features[i].permute(1, 0).cpu().detach().numpy() fea2 = pixel_features[i + b // 2].permute(1, 0).cpu().detach().numpy() idx1, idx2, dis = tools.nn_match(fea1, fea2, 'euclidean') idx1 = idx1[dis < 10] idx2 = idx2[dis < 10] h, w, _ = img1.shape img = np.vstack((img1, img2)) img[h, :] = [255, 255, 255] for i in range(len(idx1)): center1 = (int(pixel1[idx1[i], 1]), int(pixel1[idx1[i], 0])) center2 = (int(pixel2[idx2[i], 1]), int(pixel2[idx2[i], 0] + h)) cv2.line(img, center1, center2, (0, 188, 0), 1, cv2.LINE_AA) cv2.circle(img, center1, 2, (0, 0, 255), -1, cv2.LINE_AA) cv2.circle(img, center2, 2, (0, 0, 255), -1, cv2.LINE_AA) fig = plt.figure() plt.imshow(img[:, :, [2, 1, 0]]) plt.show() # cv2.namedWindow('img match') # cv2.imshow('img match', img) # cv2.waitKey(0) ######################################################################################################### batch_dict['key_pixels'] = kpts batch_dict['fea_kpl'] = pixel_features batch_dict['fea_img'] = feature_img batch_dict['score_img'] = score_img if hasattr(self, 'netvlad_img'): vlad = self.netvlad_img(pixel_features.transpose(1, 2).contiguous()) batch_dict['vlad_img'] = vlad return batch_dict class LocalPool(nn.Module): def __init__(self, in_c): super().__init__() self.conv1 = nn.Conv2d(100, 10, 1, 1, 0, bias=True) self.mp=nn.MaxPool2d((1, 10)) def forward(self, x): b, c, n, k = x.shape #k=100 x1 = x.permute(0, 3, 2, 1) # b,k,n,c x2=self.conv1(x1) x3=x2.permute(0,3,2,1) x4=self.mp(x3) return x4 # bcn1 class TransformerEncoder(nn.Module): def __init__(self, in_c=128, num_heads=4, dropout=0.1, num_layers=2): super().__init__() encoder_layer = nn.TransformerEncoderLayer(d_model=in_c, nhead=num_heads, dropout=dropout, batch_first=True) self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers) def forward(self, x): y = self.encoder(x) return y class Attention(nn.Module): def __init__(self, d_model): super(Attention, self).__init__() self.d_model = d_model self.w_q = nn.Linear(d_model, d_model) self.w_k = nn.Linear(d_model, d_model) self.w_v = nn.Linear(d_model, d_model) self.fnn = nn.Linear(d_model, d_model) # self.dp=nn.Dropout(0.1) def forward(self, q, k=None, v=None): proj_q = self.w_q(q) # BNC proj_k = self.w_k(k) proj_w = self.w_v(v) # proj_q=self.dp(proj_q) # proj_k=self.dp(proj_k) # proj_w=self.dp(proj_w) weights = nn.functional.softmax(torch.matmul(proj_q, proj_k.transpose(-2, -1)) / (self.d_model ** 0.5), dim=-1) attn_output = torch.matmul(weights, proj_w).contiguous() output = self.fnn(attn_output) return output, weights class Generator(nn.Module): def __init__(self, in_c=128, num=150): super().__init__() self.mha = Attention(in_c) self.conv1 = nn.Sequential( nn.ConvTranspose1d(in_c, in_c, kernel_size=3, stride=3, padding=0), nn.AdaptiveMaxPool1d(num) ) def forward(self, x): b, c, n = x.shape # x=x.detach() x1 = x.permute(0, 2, 1) # BNC x2, _ = self.mha(x1, x1, x1) x2 = x2.permute(0, 2, 1) x3 = self.conv1(x2) return x3 class Converter(nn.Module): def __init__(self, in_c=128): super().__init__() self.mha = Attention(in_c) self.conv1 = nn.Sequential( nn.Conv1d(in_c, in_c, kernel_size=1, stride=1, padding=0),# nn.BatchNorm1d(in_c), nn.ReLU(), nn.Conv1d(in_c, in_c // 4, kernel_size=1, stride=1, padding=0),# nn.BatchNorm1d(in_c // 4), nn.ReLU(), nn.Conv1d(in_c // 4, in_c // 8, kernel_size=1, stride=1, padding=0),# nn.BatchNorm1d(in_c // 8), nn.ReLU(), nn.Conv1d(in_c // 8, in_c // 4, kernel_size=1, stride=1, padding=0),# nn.BatchNorm1d(in_c // 4), nn.ReLU(), nn.Conv1d(in_c // 4, in_c, kernel_size=1, stride=1, padding=0),# nn.BatchNorm1d(in_c), nn.ReLU(), nn.Conv1d(in_c, in_c, kernel_size=1, stride=1, padding=0) ) self.conv2 = nn.Conv1d(in_c * 2, in_c, 1, 1, 0, bias=False) def forward(self, x): # return x b, c, n = x.shape # x=x.detach() mask = (x == 0).all(dim=1) x1 = x.permute(0, 2, 1) # BNC x2, _ = self.mha(x1, x1, x1) x2 = x2.permute(0, 2, 1) x3 = self.conv1(x) x4=torch.cat([x2,x3],dim=1) x5=self.conv2(x4) x5 = x5.masked_fill(mask.unsqueeze(1), 0) return x5 class FusionHead(nn.Module): def __init__(self, in_c=128): super().__init__() self.mha1 = Attention(in_c) self.mha2 = Attention(in_c) self.conv1 = nn.Conv1d(in_c * 2, in_c, 1) def forward(self, x): fea_kpt = x[:, :, 0] fea_kpl_gen = x[:, :, 3] B, C, K, N = x.shape x1 = x[:, :, :3] # BC3N x2 = x1.permute(0, 3, 2, 1).contiguous()#BN3C x3 = x2.view(B * N, 3, C) x4, _ = self.mha1(x3, x3, x3) x5 = torch.max(x4, dim=1)[0]#B*N 3 C x6=x5.view(B,N,C) x7, _ = self.mha2(x6, fea_kpl_gen.permute(0, 2, 1), fea_kpl_gen.permute(0, 2, 1)) x7 = x7.permute(0, 2, 1) x8 = torch.cat([fea_kpt, x7] ,dim=1) x9 = self.conv1(x8) return x9 def cosine_similarity(feature1, feature2): # BNC feature1 = feature1 / torch.sqrt(torch.sum(feature1 ** 2, -1, keepdim=True) + 1e-8) feature2 = feature2 / torch.sqrt(torch.sum(feature2 ** 2, -1, keepdim=True) + 1e-8) C = torch.bmm(feature1, feature2.transpose(1, 2)) # distance_matrix = torch.sum(feature1 ** 2, -1, keepdim=True) # distance_matrix = distance_matrix + torch.sum(feature2 ** 2, -1, keepdim=True).transpose(1, 2) # distance_matrix = distance_matrix - 2 * torch.bmm(feature1, feature2.transpose(1, 2)) # c^2=a^2+b^2-2abcos # C = distance_matrix ** 0.5 return C class Fusion(nn.Module): def __init__(self, cfg): super().__init__() flag = cfg['flag'] self.flag = flag if flag == 'fusion': self.img = ImgHead(alnet='alike-n', num_kpt=cfg['kpts_number_img'], cluster_num=cfg['cluster_num_img'], vlad_size=cfg['vlad_size']) self.bev = BEVHead(alnet='alike-n', iter=cfg['sinkhorn_iter'], num_kpt=cfg['kpts_number_bev'], cluster_num=cfg['cluster_num_bev'], vlad_size=cfg['vlad_size']) assert self.img.feature_size == self.bev.feature_size, 'img feature and image feature should be the same size' feature_size = self.img.feature_size self.localpool = LocalPool(feature_size) self.cvt_img = Converter(feature_size) self.cvt_bev = Converter(feature_size) self.gen_pan = Generator(feature_size, cfg['kpts_number_bev']) self.att_fusion = FusionHead(feature_size) # self.netvlad_fusion = NetVLADLoupe(feature_size, cfg['cluster_num_fusion'], cfg['vlad_size']) self.netvlad_fusion = NetVLAD(feature_size, cfg['cluster_num_fusion']) self.uot = UOTHead(nb_iter=cfg['sinkhorn_iter'],name='fusion') self.vlad='fusion' self.w= torch.nn.Parameter(torch.zeros(1)) if flag == 'bev': self.bev = BEVHead(alnet='alike-n',iter=cfg['sinkhorn_iter'], num_kpt=cfg['kpts_number_bev'], cluster_num=cfg['cluster_num_bev'], vlad_size=256) if flag == 'img': self.img = ImgHead(alnet='alike-n', num_kpt=cfg['kpts_number_img'], cluster_num=cfg['cluster_num_img'], vlad_size=cfg['vlad_size']) def forward(self, batch_dict): if self.flag == 'fusion': batch_dict = self.img(batch_dict) batch_dict = self.bev(batch_dict) fea_img = batch_dict['fea_img'] fea_bev = batch_dict['fea_bev'] relation = batch_dict['relation'] fea_kpt_original = batch_dict['fea_kpt_original'] # fea_kpl = batch_dict['fea_kpl'] # pixel_kpt = batch_dict['pixels_kpt'] b, n1, n2, _ = relation.shape n2 = n2 - 1 # ns=torch.sum((relation[:,:,-1]>0).all(dim=2),dim=1) # n_least=torch.min(ns) # n_least=min(n_least,256) # relation1=[] # for i in range(b): # idx=torch.randperm(ns[i])[:n_least].to(relation.device) # relation1.append(relation[i:i+1,idx]) # relation1=torch.cat(relation1) # relation=relation1 pixel_img = relation[:, :, 0:n2].clone() grid_img = pixel_img[:, :, :, [1, 0]].float() / torch.tensor([fea_img.shape[3] - 1, fea_img.shape[2] - 1]).to(fea_img.device).float() * 2 - 1 fea_pl_dual = F.grid_sample(fea_img, grid_img, align_corners=True, mode='bilinear', padding_mode='zeros') fea_pl_dual = self.localpool(fea_pl_dual).squeeze(3) fea_pt_dual_gen = self.cvt_bev(fea_pl_dual) if 'pose_to_frame' in batch_dict.keys() and hasattr(self, 'uot'): pixel_bev = relation[:, :, n2:n2 + 1, 0:2].clone() grid_bev = pixel_bev[:, :, :, [1, 0]].float() / torch.tensor([fea_bev.shape[3] - 1, fea_bev.shape[2] - 1]).to(fea_bev.device).float() * 2 - 1 fea_pt_dual = (F.grid_sample(fea_bev, grid_bev, align_corners=True, mode='bilinear', padding_mode='zeros')).squeeze(3) fea_pl_dual_gen = self.cvt_img(fea_pt_dual) batch_dict['fea_pt_dual_gen'] = fea_pt_dual_gen batch_dict['fea_pl_dual_gen'] = fea_pl_dual_gen batch_dict['fea_pt_dual'] = fea_pt_dual batch_dict['fea_pl_dual'] = fea_pl_dual fea_kpt_original_gen = self.gen_pan(fea_pt_dual_gen) batch_dict['fea_kpt_original_gen'] = fea_kpt_original_gen fea_kpl_gen = self.cvt_img(fea_kpt_original) fea_kpt_gen_gen = self.cvt_bev(fea_kpl_gen) batch_dict['fea_kpt_gen_gen'] = fea_kpt_gen_gen batch_dict['fea_kpl_gen']=fea_kpl_gen fea_kpts = torch.cat([fea_kpt_original.unsqueeze(2), fea_kpt_original_gen.unsqueeze(2), fea_kpt_gen_gen.unsqueeze(2), fea_kpl_gen.unsqueeze(2)], dim=2) fea_kpt_fusion = self.att_fusion(fea_kpts) batch_dict['fea_kpt_fusion'] = fea_kpt_original # sim10 = cosine_similarity(fea_pt_dual.permute(0, 2, 1), fea_pt_dual.permute(0, 2, 1))[0].cpu().detach().numpy() # sim11 = cosine_similarity(fea_pt_dual_gen.permute(0, 2, 1), fea_pt_dual_gen.permute(0, 2, 1))[0].cpu().detach().numpy() # sim20 = cosine_similarity(fea_pl_dual.permute(0, 2, 1), fea_pl_dual.permute(0, 2, 1))[0].cpu().detach().numpy() # sim21 = cosine_similarity(fea_pl_dual_gen.permute(0, 2, 1), fea_pl_dual_gen.permute(0, 2, 1))[0].cpu().detach().numpy() # sim30 = cosine_similarity(fea_kpt_original.permute(0, 2, 1), fea_kpt_original.permute(0, 2, 1))[0].cpu().detach().numpy() # sim31 = cosine_similarity(fea_kpt_original_gen.permute(0, 2, 1), fea_kpt_original_gen.permute(0, 2, 1))[0].cpu().detach().numpy() # sim32 = cosine_similarity(fea_kpt_gen_gen.permute(0, 2, 1), fea_kpt_gen_gen.permute(0, 2, 1))[0].cpu().detach().numpy() # fig=plt.figure() # plt.subplot(2, 4, 1), plt.imshow(sim10), plt.title('points') # plt.subplot(2, 4, 5), plt.imshow(sim11), plt.title('gen points') # plt.subplot(2, 4, 2), plt.imshow(sim20), plt.title('pixel') # plt.subplot(2, 4, 6), plt.imshow(sim21), plt.title('gen pixel') # plt.subplot(2, 4, 3), plt.imshow(sim30), plt.title('kpt orig') # plt.subplot(2, 4, 7), plt.imshow(sim31), plt.title('pan kpt') # plt.subplot(2, 4, 4), plt.imshow(sim30), plt.title('kpt orig') # plt.subplot(2, 4, 8), plt.imshow(sim32), plt.title('kpt gen gen') # plt.show() if 'pose_to_frame' in batch_dict.keys() and hasattr(self, 'uot'): self.uot(batch_dict) vlad_fusion = self.netvlad_fusion(fea_kpt_fusion.unsqueeze(3)) if self.vlad=='bev': batch_dict['vlads']=batch_dict['vlad_bev'] if self.vlad=='fusion': if 'vlad_bev' in batch_dict.keys(): batch_dict['vlads']=torch.sigmoid(self.w)*vlad_fusion + (1-torch.sigmoid(self.w))*batch_dict['vlad_bev'] else: batch_dict['vlads']=vlad_fusion if self.flag == 'bev': batch_dict = self.bev(batch_dict) batch_dict['vlads'] = batch_dict['vlad_bev'] if self.flag == 'img': batch_dict = self.img(batch_dict) batch_dict['vlads'] = batch_dict['vlad_img'] return batch_dict if __name__ == '__main__': b=BEVHead()