import numpy as np import torch from pytlsd import lsd from sklearn.cluster import DBSCAN from .base_model import BaseModel from .superpoint import SuperPoint, sample_descriptors from ..geometry import warp_lines_torch def lines_to_wireframe(lines, line_scores, all_descs, conf): """Given a set of lines, their score and dense descriptors, merge close-by endpoints and compute a wireframe defined by its junctions and connectivity. Returns: junctions: list of [num_junc, 2] tensors listing all wireframe junctions junc_scores: list of [num_junc] tensors with the junction score junc_descs: list of [dim, num_junc] tensors with the junction descriptors connectivity: list of [num_junc, num_junc] bool arrays with True when 2 junctions are connected new_lines: the new set of [b_size, num_lines, 2, 2] lines lines_junc_idx: a [b_size, num_lines, 2] tensor with the indices of the junctions of each endpoint num_true_junctions: a list of the number of valid junctions for each image in the batch, i.e. before filling with random ones """ b_size, _, _, _ = all_descs.shape device = lines.device endpoints = lines.reshape(b_size, -1, 2) ( junctions, junc_scores, junc_descs, connectivity, new_lines, lines_junc_idx, num_true_junctions, ) = ([], [], [], [], [], [], []) for bs in range(b_size): # Cluster the junctions that are close-by db = DBSCAN(eps=conf.nms_radius, min_samples=1).fit(endpoints[bs].cpu().numpy()) clusters = db.labels_ n_clusters = len(set(clusters)) num_true_junctions.append(n_clusters) # Compute the average junction and score for each cluster clusters = torch.tensor(clusters, dtype=torch.long, device=device) new_junc = torch.zeros(n_clusters, 2, dtype=torch.float, device=device) new_junc.scatter_reduce_( 0, clusters[:, None].repeat(1, 2), endpoints[bs], reduce="mean", include_self=False, ) junctions.append(new_junc) new_scores = torch.zeros(n_clusters, dtype=torch.float, device=device) new_scores.scatter_reduce_( 0, clusters, torch.repeat_interleave(line_scores[bs], 2), reduce="mean", include_self=False, ) junc_scores.append(new_scores) # Compute the new lines new_lines.append(junctions[-1][clusters].reshape(-1, 2, 2)) lines_junc_idx.append(clusters.reshape(-1, 2)) # Compute the junction connectivity junc_connect = torch.eye(n_clusters, dtype=torch.bool, device=device) pairs = clusters.reshape(-1, 2) # these pairs are connected by a line junc_connect[pairs[:, 0], pairs[:, 1]] = True junc_connect[pairs[:, 1], pairs[:, 0]] = True connectivity.append(junc_connect) # Interpolate the new junction descriptors junc_descs.append( sample_descriptors(junctions[-1][None], all_descs[bs : (bs + 1)], 8)[0] ) new_lines = torch.stack(new_lines, dim=0) lines_junc_idx = torch.stack(lines_junc_idx, dim=0) return ( junctions, junc_scores, junc_descs, connectivity, new_lines, lines_junc_idx, num_true_junctions, ) class SPWireframeDescriptor(BaseModel): default_conf = { "sp_params": { "has_detector": True, "has_descriptor": True, "descriptor_dim": 256, "trainable": False, # Inference "return_all": True, "sparse_outputs": True, "nms_radius": 4, "detection_threshold": 0.005, "max_num_keypoints": 1000, "force_num_keypoints": True, "remove_borders": 4, }, "wireframe_params": { "merge_points": True, "merge_line_endpoints": True, "nms_radius": 3, "max_n_junctions": 500, }, "max_n_lines": 250, "min_length": 15, } required_data_keys = ["image"] def _init(self, conf): self.conf = conf self.sp = SuperPoint(conf.sp_params) def detect_lsd_lines(self, x, max_n_lines=None): if max_n_lines is None: max_n_lines = self.conf.max_n_lines lines, scores, valid_lines = [], [], [] for b in range(len(x)): # For each image on batch img = (x[b].squeeze().cpu().numpy() * 255).astype(np.uint8) if max_n_lines is None: b_segs = lsd(img) else: for s in [0.3, 0.4, 0.5, 0.7, 0.8, 1.0]: b_segs = lsd(img, scale=s) if len(b_segs) >= max_n_lines: break segs_length = np.linalg.norm(b_segs[:, 2:4] - b_segs[:, 0:2], axis=1) # Remove short lines b_segs = b_segs[segs_length >= self.conf.min_length] segs_length = segs_length[segs_length >= self.conf.min_length] b_scores = b_segs[:, -1] * np.sqrt(segs_length) # Take the most relevant segments with indices = np.argsort(-b_scores) if max_n_lines is not None: indices = indices[:max_n_lines] lines.append(torch.from_numpy(b_segs[indices, :4].reshape(-1, 2, 2))) scores.append(torch.from_numpy(b_scores[indices])) valid_lines.append(torch.ones_like(scores[-1], dtype=torch.bool)) lines = torch.stack(lines).to(x) scores = torch.stack(scores).to(x) valid_lines = torch.stack(valid_lines).to(x.device) return lines, scores, valid_lines def _forward(self, data): b_size, _, h, w = data["image"].shape device = data["image"].device if not self.conf.sp_params.force_num_keypoints: assert b_size == 1, "Only batch size of 1 accepted for non padded inputs" # Line detection if "lines" not in data or "line_scores" not in data: if "original_img" in data: # Detect more lines, because when projecting them to the image most of them will be discarded lines, line_scores, valid_lines = self.detect_lsd_lines( data["original_img"], self.conf.max_n_lines * 3 ) # Apply the same transformation that is applied in homography_adaptation lines, valid_lines2 = warp_lines_torch( lines, data["H"], False, data["image"].shape[-2:] ) valid_lines = valid_lines & valid_lines2 lines[~valid_lines] = -1 line_scores[~valid_lines] = 0 # Re-sort the line segments to pick the ones that are inside the image and have bigger score sorted_scores, sorting_indices = torch.sort( line_scores, dim=-1, descending=True ) line_scores = sorted_scores[:, : self.conf.max_n_lines] sorting_indices = sorting_indices[:, : self.conf.max_n_lines] lines = torch.take_along_dim(lines, sorting_indices[..., None, None], 1) valid_lines = torch.take_along_dim(valid_lines, sorting_indices, 1) else: lines, line_scores, valid_lines = self.detect_lsd_lines(data["image"]) else: lines, line_scores, valid_lines = ( data["lines"], data["line_scores"], data["valid_lines"], ) if line_scores.shape[-1] != 0: line_scores /= ( line_scores.new_tensor(1e-8) + line_scores.max(dim=1).values[:, None] ) # SuperPoint prediction pred = self.sp(data) # Remove keypoints that are too close to line endpoints if self.conf.wireframe_params.merge_points: kp = pred["keypoints"] line_endpts = lines.reshape(b_size, -1, 2) dist_pt_lines = torch.norm(kp[:, :, None] - line_endpts[:, None], dim=-1) # For each keypoint, mark it as valid or to remove pts_to_remove = torch.any( dist_pt_lines < self.conf.sp_params.nms_radius, dim=2 ) # Simply remove them (we assume batch_size = 1 here) assert len(kp) == 1 pred["keypoints"] = pred["keypoints"][0][~pts_to_remove[0]][None] pred["keypoint_scores"] = pred["keypoint_scores"][0][~pts_to_remove[0]][ None ] pred["descriptors"] = pred["descriptors"][0].T[~pts_to_remove[0]].T[None] # Connect the lines together to form a wireframe orig_lines = lines.clone() if self.conf.wireframe_params.merge_line_endpoints and len(lines[0]) > 0: # Merge first close-by endpoints to connect lines ( line_points, line_pts_scores, line_descs, line_association, lines, lines_junc_idx, num_true_junctions, ) = lines_to_wireframe( lines, line_scores, pred["all_descriptors"], conf=self.conf.wireframe_params, ) # Add the keypoints to the junctions and fill the rest with random keypoints (all_points, all_scores, all_descs, pl_associativity) = [], [], [], [] for bs in range(b_size): all_points.append( torch.cat([line_points[bs], pred["keypoints"][bs]], dim=0) ) all_scores.append( torch.cat([line_pts_scores[bs], pred["keypoint_scores"][bs]], dim=0) ) all_descs.append( torch.cat([line_descs[bs], pred["descriptors"][bs]], dim=1) ) associativity = torch.eye( len(all_points[-1]), dtype=torch.bool, device=device ) associativity[ : num_true_junctions[bs], : num_true_junctions[bs] ] = line_association[bs][ : num_true_junctions[bs], : num_true_junctions[bs] ] pl_associativity.append(associativity) all_points = torch.stack(all_points, dim=0) all_scores = torch.stack(all_scores, dim=0) all_descs = torch.stack(all_descs, dim=0) pl_associativity = torch.stack(pl_associativity, dim=0) else: # Lines are independent all_points = torch.cat( [lines.reshape(b_size, -1, 2), pred["keypoints"]], dim=1 ) n_pts = all_points.shape[1] num_lines = lines.shape[1] num_true_junctions = [num_lines * 2] * b_size all_scores = torch.cat( [ torch.repeat_interleave(line_scores, 2, dim=1), pred["keypoint_scores"], ], dim=1, ) pred["line_descriptors"] = self.endpoints_pooling( lines, pred["all_descriptors"], (h, w) ) all_descs = torch.cat( [ pred["line_descriptors"].reshape( b_size, self.conf.sp_params.descriptor_dim, -1 ), pred["descriptors"], ], dim=2, ) pl_associativity = torch.eye(n_pts, dtype=torch.bool, device=device)[ None ].repeat(b_size, 1, 1) lines_junc_idx = ( torch.arange(num_lines * 2, device=device) .reshape(1, -1, 2) .repeat(b_size, 1, 1) ) del pred["all_descriptors"] # Remove dense descriptors to save memory torch.cuda.empty_cache() return { "keypoints": all_points, "keypoint_scores": all_scores, "descriptors": all_descs, "pl_associativity": pl_associativity, "num_junctions": torch.tensor(num_true_junctions), "lines": lines, "orig_lines": orig_lines, "lines_junc_idx": lines_junc_idx, "line_scores": line_scores, "valid_lines": valid_lines, } @staticmethod def endpoints_pooling(segs, all_descriptors, img_shape): assert segs.ndim == 4 and segs.shape[-2:] == (2, 2) filter_shape = all_descriptors.shape[-2:] scale_x = filter_shape[1] / img_shape[1] scale_y = filter_shape[0] / img_shape[0] scaled_segs = torch.round( segs * torch.tensor([scale_x, scale_y]).to(segs) ).long() scaled_segs[..., 0] = torch.clip(scaled_segs[..., 0], 0, filter_shape[1] - 1) scaled_segs[..., 1] = torch.clip(scaled_segs[..., 1], 0, filter_shape[0] - 1) line_descriptors = [ all_descriptors[ None, b, ..., torch.squeeze(b_segs[..., 1]), torch.squeeze(b_segs[..., 0]), ] for b, b_segs in enumerate(scaled_segs) ] line_descriptors = torch.cat(line_descriptors) return line_descriptors # Shape (1, 256, 308, 2) def loss(self, pred, data): raise NotImplementedError def metrics(self, pred, data): return {}