import bisect import numpy as np import cv2 def _compute_conf_thresh(data): dataset_name = data["dataset_name"][0].lower() if dataset_name == "scannet": thr = 5e-4 elif dataset_name == "megadepth": thr = 1e-4 else: raise ValueError(f"Unknown dataset: {dataset_name}") return thr # --- VISUALIZATION --- # def make_matching_figure( img0, img1, mkpts0, mkpts1, color, kpts0=None, kpts1=None, text=[], path=None, ): """ 使用OpenCV绘制匹配点可视化图像 参数: img0: 第一张图像 (BGR格式) img1: 第二张图像 (BGR格式) mkpts0: 第一张图像中的匹配点 (Nx2数组) mkpts1: 第二张图像中的匹配点 (Nx2数组) color: 每个匹配点的颜色 kpts0: 第一张图像中的所有关键点 (可选) kpts1: 第二张图像中的所有关键点 (可选) text: 要添加的文本 (可选) path: 保存图像的路径 (可选) 返回: 绘制好的OpenCV图像 (BGR格式) """ # 确保匹配点数量一致 assert mkpts0.shape[0] == mkpts1.shape[0], \ f"mkpts0: {mkpts0.shape[0]} v.s. mkpts1: {mkpts1.shape[0]}" # 确保图像有相同的高度，如果不同则调整 h0, w0 = img0.shape[:2] h1, w1 = img1.shape[:2] max_height = max(h0, h1) # 创建画布，两张图像并排显示 canvas = np.ones((max_height, w0 + w1, 3), dtype=np.uint8) * 255 # 将图像放置到画布上 canvas[:h0, :w0] = img0 canvas[:h1, w0:w0+w1] = img1 # 绘制所有关键点（如果提供） if kpts0 is not None and kpts1 is not None: for (x, y) in kpts0.astype(np.int32): cv2.circle(canvas, (x, y), 1, (255, 255, 255), -1) for (x, y) in kpts1.astype(np.int32): cv2.circle(canvas, (x + w0, y), 1, (255, 255, 255), -1) # 绘制匹配点和连接线 if mkpts0.shape[0] > 0 and mkpts1.shape[0] > 0: # 转换为整数坐标 mkpts0_int = mkpts0.astype(np.int32) mkpts1_int = mkpts1.astype(np.int32) # 绘制连接线 for i in range(len(mkpts0_int)): x0, y0 = mkpts0_int[i] x1, y1 = mkpts1_int[i] # 第二张图的x坐标需要加上第一张图的宽度 x1 += w0 # 将颜色从0-1范围转换为0-255 line_color = tuple(int(c * 255) for c in color[i][:3]) # 转换为BGR格式（因为OpenCV使用BGR） # line_color = (line_color[2], line_color[1], line_color[0]) cv2.line(canvas, (x0, y0), (x1, y1), line_color, 1) # 绘制匹配点 for i in range(len(mkpts0_int)): x0, y0 = mkpts0_int[i] x1, y1 = mkpts1_int[i] x1 += w0 pt_color = tuple(int(c * 255) for c in color[i][:3]) # pt_color = (pt_color[2], pt_color[1], pt_color[0]) cv2.circle(canvas, (x0, y0), 2, pt_color, -1) cv2.circle(canvas, (x1, y1), 2, pt_color, -1) # 添加文本 if text: # 确定文本颜色（基于图像亮度） roi = img0[:100, :200] if h0 > 100 and w0 > 200 else img0 brightness = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY).mean() text_color = (0, 0, 0) if brightness > 200 else (255, 255, 255) # 绘制文本 y_pos = 30 for i, line in enumerate(text): cv2.putText( canvas, line, (10, y_pos + i * 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, text_color, 2 ) # 保存图像（如果指定了路径） if path: cv2.imwrite(path, canvas) return canvas def _make_evaluation_figure(data, b_id, alpha="dynamic"): b_mask = data["m_bids"] == b_id conf_thr = _compute_conf_thresh(data) img0 = (data["image0"][b_id][0].cpu().numpy() * 255).round().astype(np.int32) img1 = (data["image1"][b_id][0].cpu().numpy() * 255).round().astype(np.int32) kpts0 = data["mkpts0_f"][b_mask].clone().detach().cpu().numpy() kpts1 = data["mkpts1_f"][b_mask].clone().detach().cpu().numpy() # for megadepth, we visualize matches on the resized image if "scale0" in data: kpts0 = kpts0 / data["scale0"][b_id].cpu().numpy()[[1, 0]] kpts1 = kpts1 / data["scale1"][b_id].cpu().numpy()[[1, 0]] epi_errs = data["epi_errs"][b_mask].cpu().numpy() correct_mask = epi_errs < conf_thr precision = np.mean(correct_mask) if len(correct_mask) > 0 else 0 n_correct = np.sum(correct_mask) n_gt_matches = int(data["conf_matrix_gt"][b_id].sum().cpu()) recall = 0 if n_gt_matches == 0 else n_correct / (n_gt_matches) # recall might be larger than 1, since the calculation of conf_matrix_gt # uses groundtruth depths and camera poses, but epipolar distance is used here. # matching info if alpha == "dynamic": alpha = dynamic_alpha(len(correct_mask)) color = error_colormap(epi_errs, conf_thr, alpha=alpha) text = [ f"#Matches {len(kpts0)}", f"Precision({conf_thr:.2e}) ({100 * precision:.1f}%): {n_correct}/{len(kpts0)}", f"Recall({conf_thr:.2e}) ({100 * recall:.1f}%): {n_correct}/{n_gt_matches}", ] # make the figure figure = make_matching_figure(img0, img1, kpts0, kpts1, color, text=text) return figure def _make_confidence_figure(data, b_id): # TODO: Implement confidence figure raise NotImplementedError() def make_matching_figures(data, config, mode="evaluation"): """Make matching figures for a batch. Args: data (Dict): a batch updated by PL_LoFTR. config (Dict): matcher config Returns: figures (Dict[str, List[plt.figure]] """ assert mode in ["evaluation", "confidence", "gt"] # 'confidence' figures = {mode: []} for b_id in range(data["image0"].size(0)): if mode == "evaluation": fig = _make_evaluation_figure( data, b_id, alpha=config.TRAINER.PLOT_MATCHES_ALPHA ) elif mode == "confidence": fig = _make_confidence_figure(data, b_id) else: raise ValueError(f"Unknown plot mode: {mode}") figures[mode].append(fig) return figures def dynamic_alpha( n_matches, milestones=[0, 300, 1000, 2000], alphas=[1.0, 0.8, 0.4, 0.2] ): if n_matches == 0: return 1.0 ranges = list(zip(alphas, alphas[1:] + [None])) loc = bisect.bisect_right(milestones, n_matches) - 1 _range = ranges[loc] if _range[1] is None: return _range[0] return _range[1] + (milestones[loc + 1] - n_matches) / ( milestones[loc + 1] - milestones[loc] ) * (_range[0] - _range[1]) def error_colormap(err, thr, alpha=1.0): assert alpha <= 1.0 and alpha > 0, f"Invaid alpha value: {alpha}" x = 1 - np.clip(err / (thr * 2), 0, 1) return np.clip( np.stack([2 - x * 2, x * 2, np.zeros_like(x), np.ones_like(x) * alpha], -1), 0, 1, )