Spaces:
Running
Running
| """ | |
| Implementation of the line matching methods. | |
| """ | |
| import numpy as np | |
| import cv2 | |
| import torch | |
| import torch.nn.functional as F | |
| from ..misc.geometry_utils import keypoints_to_grid | |
| class WunschLineMatcher(object): | |
| """Class matching two sets of line segments | |
| with the Needleman-Wunsch algorithm.""" | |
| def __init__( | |
| self, | |
| cross_check=True, | |
| num_samples=10, | |
| min_dist_pts=8, | |
| top_k_candidates=10, | |
| grid_size=8, | |
| sampling="regular", | |
| line_score=False, | |
| ): | |
| self.cross_check = cross_check | |
| self.num_samples = num_samples | |
| self.min_dist_pts = min_dist_pts | |
| self.top_k_candidates = top_k_candidates | |
| self.grid_size = grid_size | |
| self.line_score = line_score # True to compute saliency on a line | |
| self.sampling_mode = sampling | |
| if sampling not in ["regular", "d2_net", "asl_feat"]: | |
| raise ValueError("Wrong sampling mode: " + sampling) | |
| def forward(self, line_seg1, line_seg2, desc1, desc2): | |
| """ | |
| Find the best matches between two sets of line segments | |
| and their corresponding descriptors. | |
| """ | |
| img_size1 = (desc1.shape[2] * self.grid_size, desc1.shape[3] * self.grid_size) | |
| img_size2 = (desc2.shape[2] * self.grid_size, desc2.shape[3] * self.grid_size) | |
| device = desc1.device | |
| # Default case when an image has no lines | |
| if len(line_seg1) == 0: | |
| return np.empty((0), dtype=int) | |
| if len(line_seg2) == 0: | |
| return -np.ones(len(line_seg1), dtype=int) | |
| # Sample points regularly along each line | |
| if self.sampling_mode == "regular": | |
| line_points1, valid_points1 = self.sample_line_points(line_seg1) | |
| line_points2, valid_points2 = self.sample_line_points(line_seg2) | |
| else: | |
| line_points1, valid_points1 = self.sample_salient_points( | |
| line_seg1, desc1, img_size1, self.sampling_mode | |
| ) | |
| line_points2, valid_points2 = self.sample_salient_points( | |
| line_seg2, desc2, img_size2, self.sampling_mode | |
| ) | |
| line_points1 = torch.tensor( | |
| line_points1.reshape(-1, 2), dtype=torch.float, device=device | |
| ) | |
| line_points2 = torch.tensor( | |
| line_points2.reshape(-1, 2), dtype=torch.float, device=device | |
| ) | |
| # Extract the descriptors for each point | |
| grid1 = keypoints_to_grid(line_points1, img_size1) | |
| grid2 = keypoints_to_grid(line_points2, img_size2) | |
| desc1 = F.normalize(F.grid_sample(desc1, grid1)[0, :, :, 0], dim=0) | |
| desc2 = F.normalize(F.grid_sample(desc2, grid2)[0, :, :, 0], dim=0) | |
| # Precompute the distance between line points for every pair of lines | |
| # Assign a score of -1 for unvalid points | |
| scores = desc1.t() @ desc2 | |
| scores[~valid_points1.flatten()] = -1 | |
| scores[:, ~valid_points2.flatten()] = -1 | |
| scores = scores.reshape( | |
| len(line_seg1), self.num_samples, len(line_seg2), self.num_samples | |
| ) | |
| scores = scores.permute(0, 2, 1, 3) | |
| # scores.shape = (n_lines1, n_lines2, num_samples, num_samples) | |
| # Pre-filter the line candidates and find the best match for each line | |
| matches = self.filter_and_match_lines(scores) | |
| # [Optionally] filter matches with mutual nearest neighbor filtering | |
| if self.cross_check: | |
| matches2 = self.filter_and_match_lines(scores.permute(1, 0, 3, 2)) | |
| mutual = matches2[matches] == np.arange(len(line_seg1)) | |
| matches[~mutual] = -1 | |
| return matches | |
| def d2_net_saliency_score(self, desc): | |
| """Compute the D2-Net saliency score | |
| on a 3D or 4D descriptor.""" | |
| is_3d = len(desc.shape) == 3 | |
| b_size = len(desc) | |
| feat = F.relu(desc) | |
| # Compute the soft local max | |
| exp = torch.exp(feat) | |
| if is_3d: | |
| sum_exp = 3 * F.avg_pool1d(exp, kernel_size=3, stride=1, padding=1) | |
| else: | |
| sum_exp = 9 * F.avg_pool2d(exp, kernel_size=3, stride=1, padding=1) | |
| soft_local_max = exp / sum_exp | |
| # Compute the depth-wise maximum | |
| depth_wise_max = torch.max(feat, dim=1)[0] | |
| depth_wise_max = feat / depth_wise_max.unsqueeze(1) | |
| # Total saliency score | |
| score = torch.max(soft_local_max * depth_wise_max, dim=1)[0] | |
| normalization = torch.sum(score.reshape(b_size, -1), dim=1) | |
| if is_3d: | |
| normalization = normalization.reshape(b_size, 1) | |
| else: | |
| normalization = normalization.reshape(b_size, 1, 1) | |
| score = score / normalization | |
| return score | |
| def asl_feat_saliency_score(self, desc): | |
| """Compute the ASLFeat saliency score on a 3D or 4D descriptor.""" | |
| is_3d = len(desc.shape) == 3 | |
| b_size = len(desc) | |
| # Compute the soft local peakiness | |
| if is_3d: | |
| local_avg = F.avg_pool1d(desc, kernel_size=3, stride=1, padding=1) | |
| else: | |
| local_avg = F.avg_pool2d(desc, kernel_size=3, stride=1, padding=1) | |
| soft_local_score = F.softplus(desc - local_avg) | |
| # Compute the depth-wise peakiness | |
| depth_wise_mean = torch.mean(desc, dim=1).unsqueeze(1) | |
| depth_wise_score = F.softplus(desc - depth_wise_mean) | |
| # Total saliency score | |
| score = torch.max(soft_local_score * depth_wise_score, dim=1)[0] | |
| normalization = torch.sum(score.reshape(b_size, -1), dim=1) | |
| if is_3d: | |
| normalization = normalization.reshape(b_size, 1) | |
| else: | |
| normalization = normalization.reshape(b_size, 1, 1) | |
| score = score / normalization | |
| return score | |
| def sample_salient_points(self, line_seg, desc, img_size, saliency_type="d2_net"): | |
| """ | |
| Sample the most salient points along each line segments, with a | |
| minimal distance between each point. Pad the remaining points. | |
| Inputs: | |
| line_seg: an Nx2x2 torch.Tensor. | |
| desc: a NxDxHxW torch.Tensor. | |
| image_size: the original image size. | |
| saliency_type: 'd2_net' or 'asl_feat'. | |
| Outputs: | |
| line_points: an Nxnum_samplesx2 np.array. | |
| valid_points: a boolean Nxnum_samples np.array. | |
| """ | |
| device = desc.device | |
| if not self.line_score: | |
| # Compute the score map | |
| if saliency_type == "d2_net": | |
| score = self.d2_net_saliency_score(desc) | |
| else: | |
| score = self.asl_feat_saliency_score(desc) | |
| num_lines = len(line_seg) | |
| line_lengths = np.linalg.norm(line_seg[:, 0] - line_seg[:, 1], axis=1) | |
| # The number of samples depends on the length of the line | |
| num_samples_lst = np.clip( | |
| line_lengths // self.min_dist_pts, 2, self.num_samples | |
| ) | |
| line_points = np.empty((num_lines, self.num_samples, 2), dtype=float) | |
| valid_points = np.empty((num_lines, self.num_samples), dtype=bool) | |
| # Sample the score on a fixed number of points of each line | |
| n_samples_per_region = 4 | |
| for n in np.arange(2, self.num_samples + 1): | |
| sample_rate = n * n_samples_per_region | |
| # Consider all lines where we can fit up to n points | |
| cur_mask = num_samples_lst == n | |
| cur_line_seg = line_seg[cur_mask] | |
| cur_num_lines = len(cur_line_seg) | |
| if cur_num_lines == 0: | |
| continue | |
| line_points_x = np.linspace( | |
| cur_line_seg[:, 0, 0], cur_line_seg[:, 1, 0], sample_rate, axis=-1 | |
| ) | |
| line_points_y = np.linspace( | |
| cur_line_seg[:, 0, 1], cur_line_seg[:, 1, 1], sample_rate, axis=-1 | |
| ) | |
| cur_line_points = np.stack([line_points_x, line_points_y], axis=-1).reshape( | |
| -1, 2 | |
| ) | |
| # cur_line_points is of shape (n_cur_lines * sample_rate, 2) | |
| cur_line_points = torch.tensor( | |
| cur_line_points, dtype=torch.float, device=device | |
| ) | |
| grid_points = keypoints_to_grid(cur_line_points, img_size) | |
| if self.line_score: | |
| # The saliency score is high when the activation are locally | |
| # maximal along the line (and not in a square neigborhood) | |
| line_desc = F.grid_sample(desc, grid_points).squeeze() | |
| line_desc = line_desc.reshape(-1, cur_num_lines, sample_rate) | |
| line_desc = line_desc.permute(1, 0, 2) | |
| if saliency_type == "d2_net": | |
| scores = self.d2_net_saliency_score(line_desc) | |
| else: | |
| scores = self.asl_feat_saliency_score(line_desc) | |
| else: | |
| scores = F.grid_sample(score.unsqueeze(1), grid_points).squeeze() | |
| # Take the most salient point in n distinct regions | |
| scores = scores.reshape(-1, n, n_samples_per_region) | |
| best = torch.max(scores, dim=2, keepdim=True)[1].cpu().numpy() | |
| cur_line_points = cur_line_points.reshape(-1, n, n_samples_per_region, 2) | |
| cur_line_points = np.take_along_axis( | |
| cur_line_points, best[..., None], axis=2 | |
| )[:, :, 0] | |
| # Pad | |
| cur_valid_points = np.ones((cur_num_lines, self.num_samples), dtype=bool) | |
| cur_valid_points[:, n:] = False | |
| cur_line_points = np.concatenate( | |
| [ | |
| cur_line_points, | |
| np.zeros((cur_num_lines, self.num_samples - n, 2), dtype=float), | |
| ], | |
| axis=1, | |
| ) | |
| line_points[cur_mask] = cur_line_points | |
| valid_points[cur_mask] = cur_valid_points | |
| return line_points, valid_points | |
| def sample_line_points(self, line_seg): | |
| """ | |
| Regularly sample points along each line segments, with a minimal | |
| distance between each point. Pad the remaining points. | |
| Inputs: | |
| line_seg: an Nx2x2 torch.Tensor. | |
| Outputs: | |
| line_points: an Nxnum_samplesx2 np.array. | |
| valid_points: a boolean Nxnum_samples np.array. | |
| """ | |
| num_lines = len(line_seg) | |
| line_lengths = np.linalg.norm(line_seg[:, 0] - line_seg[:, 1], axis=1) | |
| # Sample the points separated by at least min_dist_pts along each line | |
| # The number of samples depends on the length of the line | |
| num_samples_lst = np.clip( | |
| line_lengths // self.min_dist_pts, 2, self.num_samples | |
| ) | |
| line_points = np.empty((num_lines, self.num_samples, 2), dtype=float) | |
| valid_points = np.empty((num_lines, self.num_samples), dtype=bool) | |
| for n in np.arange(2, self.num_samples + 1): | |
| # Consider all lines where we can fit up to n points | |
| cur_mask = num_samples_lst == n | |
| cur_line_seg = line_seg[cur_mask] | |
| line_points_x = np.linspace( | |
| cur_line_seg[:, 0, 0], cur_line_seg[:, 1, 0], n, axis=-1 | |
| ) | |
| line_points_y = np.linspace( | |
| cur_line_seg[:, 0, 1], cur_line_seg[:, 1, 1], n, axis=-1 | |
| ) | |
| cur_line_points = np.stack([line_points_x, line_points_y], axis=-1) | |
| # Pad | |
| cur_num_lines = len(cur_line_seg) | |
| cur_valid_points = np.ones((cur_num_lines, self.num_samples), dtype=bool) | |
| cur_valid_points[:, n:] = False | |
| cur_line_points = np.concatenate( | |
| [ | |
| cur_line_points, | |
| np.zeros((cur_num_lines, self.num_samples - n, 2), dtype=float), | |
| ], | |
| axis=1, | |
| ) | |
| line_points[cur_mask] = cur_line_points | |
| valid_points[cur_mask] = cur_valid_points | |
| return line_points, valid_points | |
| def filter_and_match_lines(self, scores): | |
| """ | |
| Use the scores to keep the top k best lines, compute the Needleman- | |
| Wunsch algorithm on each candidate pairs, and keep the highest score. | |
| Inputs: | |
| scores: a (N, M, n, n) torch.Tensor containing the pairwise scores | |
| of the elements to match. | |
| Outputs: | |
| matches: a (N) np.array containing the indices of the best match | |
| """ | |
| # Pre-filter the pairs and keep the top k best candidate lines | |
| line_scores1 = scores.max(3)[0] | |
| valid_scores1 = line_scores1 != -1 | |
| line_scores1 = (line_scores1 * valid_scores1).sum(2) / valid_scores1.sum(2) | |
| line_scores2 = scores.max(2)[0] | |
| valid_scores2 = line_scores2 != -1 | |
| line_scores2 = (line_scores2 * valid_scores2).sum(2) / valid_scores2.sum(2) | |
| line_scores = (line_scores1 + line_scores2) / 2 | |
| topk_lines = torch.argsort(line_scores, dim=1)[:, -self.top_k_candidates :] | |
| scores, topk_lines = scores.cpu().numpy(), topk_lines.cpu().numpy() | |
| # topk_lines.shape = (n_lines1, top_k_candidates) | |
| top_scores = np.take_along_axis(scores, topk_lines[:, :, None, None], axis=1) | |
| # Consider the reversed line segments as well | |
| top_scores = np.concatenate([top_scores, top_scores[..., ::-1]], axis=1) | |
| # Compute the line distance matrix with Needleman-Wunsch algo and | |
| # retrieve the closest line neighbor | |
| n_lines1, top2k, n, m = top_scores.shape | |
| top_scores = top_scores.reshape(n_lines1 * top2k, n, m) | |
| nw_scores = self.needleman_wunsch(top_scores) | |
| nw_scores = nw_scores.reshape(n_lines1, top2k) | |
| matches = np.mod(np.argmax(nw_scores, axis=1), top2k // 2) | |
| matches = topk_lines[np.arange(n_lines1), matches] | |
| return matches | |
| def needleman_wunsch(self, scores): | |
| """ | |
| Batched implementation of the Needleman-Wunsch algorithm. | |
| The cost of the InDel operation is set to 0 by subtracting the gap | |
| penalty to the scores. | |
| Inputs: | |
| scores: a (B, N, M) np.array containing the pairwise scores | |
| of the elements to match. | |
| """ | |
| b, n, m = scores.shape | |
| # Recalibrate the scores to get a gap score of 0 | |
| gap = 0.1 | |
| nw_scores = scores - gap | |
| # Run the dynamic programming algorithm | |
| nw_grid = np.zeros((b, n + 1, m + 1), dtype=float) | |
| for i in range(n): | |
| for j in range(m): | |
| nw_grid[:, i + 1, j + 1] = np.maximum( | |
| np.maximum(nw_grid[:, i + 1, j], nw_grid[:, i, j + 1]), | |
| nw_grid[:, i, j] + nw_scores[:, i, j], | |
| ) | |
| return nw_grid[:, -1, -1] | |
| def get_pairwise_distance(self, line_seg1, line_seg2, desc1, desc2): | |
| """ | |
| Compute the OPPOSITE of the NW score for pairs of line segments | |
| and their corresponding descriptors. | |
| """ | |
| num_lines = len(line_seg1) | |
| assert num_lines == len( | |
| line_seg2 | |
| ), "The same number of lines is required in pairwise score." | |
| img_size1 = (desc1.shape[2] * self.grid_size, desc1.shape[3] * self.grid_size) | |
| img_size2 = (desc2.shape[2] * self.grid_size, desc2.shape[3] * self.grid_size) | |
| device = desc1.device | |
| # Sample points regularly along each line | |
| line_points1, valid_points1 = self.sample_line_points(line_seg1) | |
| line_points2, valid_points2 = self.sample_line_points(line_seg2) | |
| line_points1 = torch.tensor( | |
| line_points1.reshape(-1, 2), dtype=torch.float, device=device | |
| ) | |
| line_points2 = torch.tensor( | |
| line_points2.reshape(-1, 2), dtype=torch.float, device=device | |
| ) | |
| # Extract the descriptors for each point | |
| grid1 = keypoints_to_grid(line_points1, img_size1) | |
| grid2 = keypoints_to_grid(line_points2, img_size2) | |
| desc1 = F.normalize(F.grid_sample(desc1, grid1)[0, :, :, 0], dim=0) | |
| desc1 = desc1.reshape(-1, num_lines, self.num_samples) | |
| desc2 = F.normalize(F.grid_sample(desc2, grid2)[0, :, :, 0], dim=0) | |
| desc2 = desc2.reshape(-1, num_lines, self.num_samples) | |
| # Compute the distance between line points for every pair of lines | |
| # Assign a score of -1 for unvalid points | |
| scores = torch.einsum("dns,dnt->nst", desc1, desc2).cpu().numpy() | |
| scores = scores.reshape(num_lines * self.num_samples, self.num_samples) | |
| scores[~valid_points1.flatten()] = -1 | |
| scores = scores.reshape(num_lines, self.num_samples, self.num_samples) | |
| scores = scores.transpose(1, 0, 2).reshape(self.num_samples, -1) | |
| scores[:, ~valid_points2.flatten()] = -1 | |
| scores = scores.reshape(self.num_samples, num_lines, self.num_samples) | |
| scores = scores.transpose(1, 0, 2) | |
| # scores.shape = (num_lines, num_samples, num_samples) | |
| # Compute the NW score for each pair of lines | |
| pairwise_scores = np.array([self.needleman_wunsch(s) for s in scores]) | |
| return -pairwise_scores | |