From 419a7db54334a381a5f7d1ebe65b22fff312147d Mon Sep 17 00:00:00 2001 From: Jiao77 Date: Thu, 25 Sep 2025 22:05:39 +0800 Subject: [PATCH] finish Inference and Matching Part. --- README.md | 45 +++++++++++++ configs/base_config.yaml | 12 ++++ docs/NextStep.md | 133 +++++++++++++++++++++++++++++++++++++++ docs/feature_work.md | 4 +- match.py | 80 +++++++++++++++++++++-- models/rord.py | 90 ++++++++++++++++++++++---- 6 files changed, 346 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index b19d798..c31ac0a 100644 --- a/README.md +++ b/README.md @@ -274,6 +274,51 @@ RoRD 模型基于 D2-Net 架构，使用 VGG-16 作为骨干网络，**专门针 - **二值化特征距离**: 强化几何边界特征，弱化灰度变化 - **几何感知困难负样本**: 基于结构相似性而非像素相似性选择负样本 +## 🔎 推理与匹配（FPN 路径与 NMS） + +项目已支持通过 FPN 单次推理产生多尺度特征，并在匹配阶段引入半径 NMS 去重以减少冗余关键点： + +在 `configs/base_config.yaml` 中启用 FPN 与 NMS： + +```yaml +model: + fpn: + enabled: true + out_channels: 256 + levels: [2, 3, 4] + +matching: + use_fpn: true + nms: + enabled: true + radius: 4 + score_threshold: 0.5 +``` + +运行匹配并将过程写入 TensorBoard： + +```bash +uv run python match.py \ + --config configs/base_config.yaml \ + --layout /path/to/layout.png \ + --template /path/to/template.png \ + --tb_log_matches +``` + +如需回退旧“图像金字塔”路径，将 `matching.use_fpn` 设为 `false` 即可。 + +也可使用 CLI 快捷开关临时覆盖： + +```bash +# 关闭 FPN（等同 matching.use_fpn=false） +uv run python match.py --config configs/base_config.yaml --fpn_off \ + --layout /path/to/layout.png --template /path/to/template.png + +# 关闭关键点去重（NMS） +uv run python match.py --config configs/base_config.yaml --no_nms \ + --layout /path/to/layout.png --template /path/to/template.png +``` + ### 训练策略 - 几何结构学习 模型通过**几何结构学习**策略进行训练： - **曼哈顿变换生成训练对**: 利用90度旋转等曼哈顿变换 diff --git a/configs/base_config.yaml b/configs/base_config.yaml index ee03075..80db1ce 100644 --- a/configs/base_config.yaml +++ b/configs/base_config.yaml @@ -5,6 +5,13 @@ training: patch_size: 256 scale_jitter_range: [0.8, 1.2] +model: + fpn: + enabled: true + out_channels: 256 + levels: [2, 3, 4] + norm: "bn" + matching: keypoint_threshold: 0.5 ransac_reproj_threshold: 5.0 @@ -12,6 +19,11 @@ matching: pyramid_scales: [0.75, 1.0, 1.5] inference_window_size: 1024 inference_stride: 768 + use_fpn: true + nms: + enabled: true + radius: 4 + score_threshold: 0.5 evaluation: iou_threshold: 0.5 diff --git a/docs/NextStep.md b/docs/NextStep.md index 5a5ddbb..b6b57ee 100644 --- a/docs/NextStep.md +++ b/docs/NextStep.md @@ -122,3 +122,136 @@ - [ ] 修改配置和脚本，接入 SummaryWriter。 - [ ] 准备示例 Notebook/文档，展示 TensorBoard 面板截图。 - [ ] 后续评估是否需要接入 W&B、MLflow 等更高级平台。 + +--- + +# 推理与匹配改造计划（FPN + NMS） + +日期：2025-09-25 + +## 目标 +- 将当前的“图像金字塔 + 多次推理”的匹配流程，升级为“单次推理 + 特征金字塔 (FPN)”以显著提速。 +- 在滑动窗口提取关键点后增加去重（NMS/半径抑制），降低冗余点与后续 RANSAC 的计算量。 +- 保持与现有 YAML 配置、TensorBoard 记录和命令行接口的一致性；以 uv 为包管理器管理依赖和运行。 + +## 设计概览 +- FPN：在 `models/rord.py` 中，从骨干网络多层提取特征（例如 VGG 的 relu2_2/relu3_3/relu4_3），通过横向 1x1 卷积与自顶向下上采样构建 P2/P3/P4 金字塔特征；为每个尺度共享或独立地接上检测头与描述子头，导出同维度描述子。 +- 匹配路径：`match.py` 新增 FPN 路径，单次前向获得多尺度特征，逐层与模板进行匹配与几何验证；保留旧路径（图像金字塔）作为回退，通过配置开关切换。 +- 去重策略：在滑窗聚合关键点后，基于“分数优先 + 半径抑制 (radius NMS)”进行去重；半径和分数阈值配置化。 + +## 配置变更（YAML） +在 `configs/base_config.yaml` 中新增/扩展： + +```yaml +model: + fpn: + enabled: true # 开启 FPN 推理 + out_channels: 256 # 金字塔特征通道数 + levels: [2, 3, 4] # 输出层级，对应 P2/P3/P4 + norm: "bn" # 归一化类型：bn/gn/none + +matching: + use_fpn: true # 使用 FPN 路径；false 则沿用图像金字塔 + nms: + enabled: true + radius: 4 # 半径抑制像素半径 + score_threshold: 0.5 # 关键点保留的最低分数 + # 其余已有参数保留，如 ransac_reproj_threshold/min_inliers/inference_window_size... +``` + +注意：所有相对路径依旧使用 `utils.config_loader.to_absolute_path` 以配置文件所在目录为基准解析。 + +## 实施步骤 + +1) 基线分支与依赖 +- 新开分支保存改造： + ```bash + git checkout -b feature/fpn-matching + uv sync + ``` +- 目前不引入新三方库，继续使用现有 `torch/opencv/numpy`。 + +2) 模型侧改造（`models/rord.py`） +- 提取多层特征：在骨干网络中暴露中间层输出（如 C2/C3/C4）。 +- 构建 FPN： + - 使用 1x1 conv 降维对齐通道； + - 自顶向下上采样并逐级相加； + - 3x3 conv 平滑，得到 P2/P3/P4； + - 可选归一化（BN/GN）。 +- 头部适配：复用或复制现有检测头/描述子头到每个 P 层，输出： + - det_scores[L]：B×1×H_L×W_L + - descriptors[L]：B×D×H_L×W_L（D 与现有描述子维度一致） +- 前向接口： + - 训练模式：维持现有输出以兼容训练； + - 匹配/评估模式：支持 `return_pyramid=True` 返回 {P2,P3,P4} 的 det/desc。 + +3) 匹配侧改造（`match.py`） +- 配置读取：根据 `matching.use_fpn` 决定走 FPN 或旧图像金字塔路径。 +- FPN 路径： + - 对 layout 与 template 各做一次前向，获得 {det, desc}×L； + - 对每个层级 L： + - 从 det_scores[L] 以 `score_threshold` 抽取关键点坐标与分数； + - 半径 NMS 去重（见步骤 4）； + - 使用 desc[L] 在对应层做特征最近邻匹配（可选比值测试）并估计单应性 H_L（RANSAC）； + - 融合多个层级的候选：选取内点数最多或综合打分最佳的实例； + - 将层级坐标映射回原图坐标；输出 bbox 与 H。 +- 旧路径保留：若 `use_fpn=false`，继续使用当前图像金字塔多次推理策略，便于回退与对照实验。 + +4) 关键点去重（NMS/半径抑制） +- 输入：关键点集合 K = {(x, y, score)}。 +- 算法：按 score 降序遍历，若与已保留点的欧氏距离 < radius 则丢弃，否则保留。 +- 复杂度：O(N log N) 排序 + O(N·k) 检查（k 为邻域个数，可通过网格划分加速）。 +- 参数：`matching.nms.radius`、`matching.nms.score_threshold`。 + +5) TensorBoard 记录（扩展） +- Scalars： + - `match_fpn/level_L/keypoints_before_nms`、`keypoints_after_nms` + - `match_fpn/level_L/inliers`、`best_instance_inliers` + - `match_fpn/instances_found`、`runtime_ms` +- Text/Image： + - 关键点可视化（可选），最佳实例覆盖图； + - 记录使用的层级与最终选中尺度信息。 + +6) 兼容性与回退 +- 通过 YAML `matching.use_fpn` 开关控制路径； +- 保持 CLI 不变，新增可选 `--fpn-off`（等同 use_fpn=false）仅作为临时调试； +- 若新路径异常可快速回退旧路径，保证生产可用性。 + +## 开发里程碑与工时预估 +- M1（0.5 天）：配置与分支、占位接口、日志钩子。 +- M2（1.5 天）：FPN 实现与前向接口；单图 smoke 测试。 +- M3（1 天）：`match.py` FPN 路径、尺度回映射与候选融合。 +- M4（0.5 天）：NMS 实现与参数打通； +- M5（0.5 天）：TensorBoard 指标与可视化； +- M6（0.5 天）：对照基线的性能与速度评估，整理报告。 + +## 质量门禁与验收标准 +- 构建：`uv sync` 无错误；`python -m compileall` 通过； +- 功能：在 2–3 张样例上，FPN 路径输出的实例数量与旧路径相当或更优； +- 速度：相同输入，FPN 路径总耗时较旧路径下降 ≥ 30%； +- 稳定性：无异常崩溃；在找不到匹配时能优雅返回空结果； +- 指标：TensorBoard 中关键点数量、NMS 前后对比、内点数、总实例数与运行时均可见。 + +## 快速试用（命令） +```bash +# 同步环境 +uv sync + +# 基于 YAML 启用 FPN 匹配（推荐） +uv run python match.py \ + --config configs/base_config.yaml \ + --layout /path/to/layout.png \ + --template /path/to/template.png \ + --tb_log_matches + +# 临时关闭 FPN（对照实验） +# 可通过把 configs 中 matching.use_fpn 设为 false，或后续提供 --fpn-off 开关 + +# 打开 TensorBoard 查看匹配指标 +uv run tensorboard --logdir runs +``` + +## 风险与回滚 +- FPN 输出与原检测/描述子头的维度/分布不一致，需在实现时对齐通道与归一化； +- 多层融合策略（如何选取最佳实例）可能影响稳定性，可先以“内点数最大”作为启发式； +- 如出现精度下降或不稳定，立即回退 `matching.use_fpn=false`，保留旧流程并开启数据记录比对差异。 diff --git a/docs/feature_work.md b/docs/feature_work.md index ed7b36b..7e1ef83 100644 --- a/docs/feature_work.md +++ b/docs/feature_work.md @@ -62,13 +62,13 @@ > *目标：大幅提升大尺寸版图的匹配速度和多尺度检测能力。* -- [ ] **将模型改造为特征金字塔网络 (FPN) 架构** +- [x] **将模型改造为特征金字塔网络 (FPN) 架构** - **✔️ 价值**: 当前的多尺度匹配需要多次缩放图像并推理，速度慢。FPN 只需一次推理即可获得所有尺度的特征，极大加速匹配过程。 - **📝 执行方案**: 1. 修改 `models/rord.py`，从骨干网络的不同层级（如 VGG 的 `relu2_2`, `relu3_3`, `relu4_3`）提取特征图。 2. 添加上采样和横向连接层来融合这些特征图，构建出特征金字塔。 3. 修改 `match.py`，使其能够直接从 FPN 的不同层级获取特征，替代原有的图像金字塔循环。 -- [ ] **在滑动窗口匹配后增加关键点去重** +- [x] **在滑动窗口匹配后增加关键点去重** - **✔️ 价值**: `match.py` 中的滑动窗口在重叠区域会产生大量重复的关键点，增加后续匹配的计算量并可能影响精度。 - **📝 执行方案**: 1. 在 `match.py` 的 `extract_features_sliding_window` 函数返回前。 diff --git a/match.py b/match.py index 1368410..71b00ed 100644 --- a/match.py +++ b/match.py @@ -44,6 +44,24 @@ def extract_keypoints_and_descriptors(model, image_tensor, kp_thresh): return keypoints, descriptors + +# --- (新增) 简单半径 NMS 去重 --- +def radius_nms(kps: torch.Tensor, scores: torch.Tensor, radius: float) -> torch.Tensor: + if kps.numel() == 0: + return torch.empty((0,), dtype=torch.long, device=kps.device) + idx = torch.argsort(scores, descending=True) + keep = [] + taken = torch.zeros(len(kps), dtype=torch.bool, device=kps.device) + for i in idx: + if taken[i]: + continue + keep.append(i.item()) + di = kps - kps[i] + dist2 = (di[:, 0]**2 + di[:, 1]**2) + taken |= dist2 <= (radius * radius) + taken[i] = True + return torch.tensor(keep, dtype=torch.long, device=kps.device) + # --- (新增) 滑动窗口特征提取函数 --- def extract_features_sliding_window(model, large_image, transform, matching_cfg): """ @@ -88,6 +106,40 @@ def extract_features_sliding_window(model, large_image, transform, matching_cfg) return torch.cat(all_kps, dim=0), torch.cat(all_descs, dim=0) +# --- (新增) FPN 路径的关键点与描述子抽取 --- +def extract_from_pyramid(model, image_tensor, kp_thresh, nms_cfg): + with torch.no_grad(): + pyramid = model(image_tensor, return_pyramid=True) + all_kps = [] + all_desc = [] + for level_name, (det, desc, stride) in pyramid.items(): + binary = (det > kp_thresh).squeeze(0).squeeze(0) + coords = torch.nonzero(binary).float() # y,x + if len(coords) == 0: + continue + scores = det.squeeze()[binary] + # 采样描述子 + coords_for_grid = coords.flip(1).view(1, -1, 1, 2) + coords_for_grid = coords_for_grid / torch.tensor([(desc.shape[3]-1)/2, (desc.shape[2]-1)/2], device=desc.device) - 1 + descs = F.grid_sample(desc, coords_for_grid, align_corners=True).squeeze().T + descs = F.normalize(descs, p=2, dim=1) + + # 映射回原图坐标 + kps = coords.flip(1) * float(stride) + + # NMS + if nms_cfg and nms_cfg.get('enabled', False): + keep = radius_nms(kps, scores, float(nms_cfg.get('radius', 4))) + if len(keep) > 0: + kps = kps[keep] + descs = descs[keep] + all_kps.append(kps) + all_desc.append(descs) + if not all_kps: + return torch.tensor([], device=image_tensor.device), torch.tensor([], device=image_tensor.device) + return torch.cat(all_kps, dim=0), torch.cat(all_desc, dim=0) + + # --- 互近邻匹配 (无变动) --- def mutual_nearest_neighbor(descs1, descs2): if len(descs1) == 0 or len(descs2) == 0: @@ -113,8 +165,13 @@ def match_template_multiscale( """ 在不同尺度下搜索模板，并检测多个实例 """ - # 1. 对大版图使用滑动窗口提取全部特征 - layout_kps, layout_descs = extract_features_sliding_window(model, layout_image, transform, matching_cfg) + # 1. 版图特征提取：根据配置选择 FPN 或滑窗 + device = next(model.parameters()).device + if getattr(matching_cfg, 'use_fpn', False): + layout_tensor = transform(layout_image).unsqueeze(0).to(device) + layout_kps, layout_descs = extract_from_pyramid(model, layout_tensor, float(matching_cfg.keypoint_threshold), getattr(matching_cfg, 'nms', {})) + else: + layout_kps, layout_descs = extract_features_sliding_window(model, layout_image, transform, matching_cfg) if log_writer: log_writer.add_scalar("match/layout_keypoints", len(layout_kps), log_step) @@ -154,8 +211,11 @@ def match_template_multiscale( scaled_template = template_image.resize((new_W, new_H), Image.LANCZOS) template_tensor = transform(scaled_template).unsqueeze(0).to(layout_kps.device) - # 提取缩放后模板的特征 - template_kps, template_descs = extract_keypoints_and_descriptors(model, template_tensor, keypoint_threshold) + # 提取缩放后模板的特征：FPN 或单尺度 + if getattr(matching_cfg, 'use_fpn', False): + template_kps, template_descs = extract_from_pyramid(model, template_tensor, keypoint_threshold, getattr(matching_cfg, 'nms', {})) + else: + template_kps, template_descs = extract_keypoints_and_descriptors(model, template_tensor, keypoint_threshold) if len(template_kps) < 4: continue @@ -227,6 +287,8 @@ if __name__ == "__main__": parser.add_argument('--experiment_name', type=str, default=None, help="TensorBoard 实验名称，覆盖配置文件设置") parser.add_argument('--tb_log_matches', action='store_true', help="启用模板匹配过程的 TensorBoard 记录") parser.add_argument('--disable_tensorboard', action='store_true', help="禁用 TensorBoard 记录") + parser.add_argument('--fpn_off', action='store_true', help="关闭 FPN 匹配路径（等同于 matching.use_fpn=false）") + parser.add_argument('--no_nms', action='store_true', help="关闭关键点去重（NMS）") parser.add_argument('--layout', type=str, required=True) parser.add_argument('--template', type=str, required=True) parser.add_argument('--output', type=str) @@ -262,6 +324,16 @@ if __name__ == "__main__": tb_path.parent.mkdir(parents=True, exist_ok=True) writer = SummaryWriter(tb_path.as_posix()) + # CLI 快捷开关覆盖 YAML 配置 + try: + if args.fpn_off: + matching_cfg.use_fpn = False + if args.no_nms and hasattr(matching_cfg, 'nms'): + matching_cfg.nms.enabled = False + except Exception: + # 若 OmegaConf 结构不可写，忽略并在后续逻辑中以 getattr 的方式读取 + pass + transform = get_transform() model = RoRD().cuda() model.load_state_dict(torch.load(model_path)) diff --git a/models/rord.py b/models/rord.py index 2d97705..81ed0a6 100644 --- a/models/rord.py +++ b/models/rord.py @@ -2,20 +2,26 @@ import torch import torch.nn as nn +import torch.nn.functional as F from torchvision import models class RoRD(nn.Module): - def __init__(self): + def __init__(self, fpn_out_channels: int = 256, fpn_levels=(2, 3, 4)): """ 修复后的 RoRD 模型。 - 实现了共享骨干网络，以提高计算效率和减少内存占用。 - 确保检测头和描述子头使用相同尺寸的特征图。 + - 新增（可选）FPN 推理路径，提供多尺度特征用于高效匹配。 """ super(RoRD, self).__init__() vgg16_features = models.vgg16(pretrained=False).features - - # 共享骨干网络 - 只使用到 relu4_3，确保特征图尺寸一致 + + # VGG16 特征各阶段索引（conv & relu 层序列） + # relu2_2 索引 8，relu3_3 索引 15，relu4_3 索引 22 + self.features = vgg16_features + + # 共享骨干（向后兼容单尺度路径，使用到 relu4_3） self.backbone = nn.Sequential(*list(vgg16_features.children())[:23]) # 检测头 @@ -38,12 +44,72 @@ class RoRD(nn.Module): nn.InstanceNorm2d(128) ) - def forward(self, x): - # 共享特征提取 - features = self.backbone(x) - - # 检测器和描述子使用相同的特征图 - detection_map = self.detection_head(features) - descriptors = self.descriptor_head(features) - - return detection_map, descriptors \ No newline at end of file + # --- FPN 组件（用于可选多尺度推理） --- + self.fpn_out_channels = fpn_out_channels + self.fpn_levels = tuple(sorted(set(fpn_levels))) # e.g., (2,3,4) + + # 横向连接 1x1 将 C2(128)/C3(256)/C4(512) 对齐到相同通道数 + self.lateral_c2 = nn.Conv2d(128, fpn_out_channels, kernel_size=1) + self.lateral_c3 = nn.Conv2d(256, fpn_out_channels, kernel_size=1) + self.lateral_c4 = nn.Conv2d(512, fpn_out_channels, kernel_size=1) + + # 平滑 3x3 conv + self.smooth_p2 = nn.Conv2d(fpn_out_channels, fpn_out_channels, kernel_size=3, padding=1) + self.smooth_p3 = nn.Conv2d(fpn_out_channels, fpn_out_channels, kernel_size=3, padding=1) + self.smooth_p4 = nn.Conv2d(fpn_out_channels, fpn_out_channels, kernel_size=3, padding=1) + + # 共享的 FPN 检测/描述子头（输入通道为 fpn_out_channels） + self.det_head_fpn = nn.Sequential( + nn.Conv2d(fpn_out_channels, 128, kernel_size=3, padding=1), + nn.ReLU(inplace=True), + nn.Conv2d(128, 1, kernel_size=1), + nn.Sigmoid(), + ) + self.desc_head_fpn = nn.Sequential( + nn.Conv2d(fpn_out_channels, 128, kernel_size=3, padding=1), + nn.ReLU(inplace=True), + nn.Conv2d(128, 128, kernel_size=1), + nn.InstanceNorm2d(128), + ) + + def forward(self, x: torch.Tensor, return_pyramid: bool = False): + if not return_pyramid: + # 向后兼容的单尺度路径（relu4_3） + features = self.backbone(x) + detection_map = self.detection_head(features) + descriptors = self.descriptor_head(features) + return detection_map, descriptors + + # --- FPN 路径：提取 C2/C3/C4 --- + c2, c3, c4 = self._extract_c234(x) + p4 = self.lateral_c4(c4) + p3 = self.lateral_c3(c3) + F.interpolate(p4, size=c3.shape[-2:], mode="nearest") + p2 = self.lateral_c2(c2) + F.interpolate(p3, size=c2.shape[-2:], mode="nearest") + + p4 = self.smooth_p4(p4) + p3 = self.smooth_p3(p3) + p2 = self.smooth_p2(p2) + + pyramid = {} + if 4 in self.fpn_levels: + pyramid["P4"] = (self.det_head_fpn(p4), self.desc_head_fpn(p4), 8) + if 3 in self.fpn_levels: + pyramid["P3"] = (self.det_head_fpn(p3), self.desc_head_fpn(p3), 4) + if 2 in self.fpn_levels: + pyramid["P2"] = (self.det_head_fpn(p2), self.desc_head_fpn(p2), 2) + return pyramid + + def _extract_c234(self, x: torch.Tensor): + """提取 VGG 中间层特征：C2(relU2_2), C3(relu3_3), C4(relu4_3).""" + c2 = c3 = c4 = None + for i, layer in enumerate(self.features): + x = layer(x) + if i == 8: # relu2_2 + c2 = x + elif i == 15: # relu3_3 + c3 = x + elif i == 22: # relu4_3 + c4 = x + break + assert c2 is not None and c3 is not None and c4 is not None + return c2, c3, c4 \ No newline at end of file -- 2.49.1