import torch import torch.nn.functional as F from models.rord import RoRD from torchvision import transforms from utils.transforms import SobelTransform import numpy as np import cv2 from PIL import Image def extract_keypoints_and_descriptors(model, image): """ 从 RoRD 模型中提取关键点和描述子。 参数： model (RoRD): RoRD 模型。 image (torch.Tensor): 输入图像张量，形状为 [1, 1, H, W]。 返回： tuple: (keypoints_input, descriptors) - keypoints_input: [N, 2] float tensor，关键点在输入图像中的坐标。 - descriptors: [N, 128] float tensor，L2 归一化的描述子。 """ with torch.no_grad(): detection_map, _, desc_rord = model(image) desc = desc_rord # 使用 RoRD 描述子头 # 从检测图中提取关键点 thresh = 0.5 binary_map = (detection_map > thresh).float() coords = torch.nonzero(binary_map[0, 0] > thresh).float() # [N, 2]，每个行是 (i_d, j_d) keypoints_input = coords * 16.0 # 将特征图坐标映射到输入图像坐标（stride=16） # 从描述子图中提取描述子 # detection_map 的形状为 [1, 1, H/16, W/16]，desc 的形状为 [1, 128, H/8, W/8] # 将 detection_map 的坐标映射到 desc 的坐标：(i_d * 2, j_d * 2) keypoints_desc = (coords * 2).long() # [N, 2]，整数坐标 H_desc, W_desc = desc.shape[2], desc.shape[3] mask = (keypoints_desc[:, 0] < H_desc) & (keypoints_desc[:, 1] < W_desc) keypoints_desc = keypoints_desc[mask] keypoints_input = keypoints_input[mask] # 提取描述子 descriptors = desc[0, :, keypoints_desc[:, 0], keypoints_desc[:, 1]].T # [N, 128] # L2 归一化描述子 descriptors = F.normalize(descriptors, p=2, dim=1) return keypoints_input, descriptors def mutual_nearest_neighbor(template_descs, layout_descs): """ 使用互最近邻（MNN）找到模板和版图之间的匹配。 参数： template_descs (torch.Tensor): 模板描述子，形状为 [M, 128]。 layout_descs (torch.Tensor): 版图描述子，形状为 [N, 128]。 返回： list: [(i_template, i_layout)]，互最近邻匹配对的列表。 """ M, N = template_descs.size(0), layout_descs.size(0) if M == 0 or N == 0: return [] similarity_matrix = template_descs @ layout_descs.T # [M, N]，点积矩阵 # 找到每个模板描述子的最近邻 nn_template_to_layout = torch.argmax(similarity_matrix, dim=1) # [M] # 找到每个版图描述子的最近邻 nn_layout_to_template = torch.argmax(similarity_matrix, dim=0) # [N] # 找到互最近邻 mutual_matches = [] for i in range(M): j = nn_template_to_layout[i] if nn_layout_to_template[j] == i: mutual_matches.append((i.item(), j.item())) return mutual_matches def ransac_filter(matches, template_kps, layout_kps): """ 使用 RANSAC 对匹配进行几何验证，并返回内点。 参数： matches (list): [(i_template, i_layout)]，匹配对列表。 template_kps (torch.Tensor): 模板关键点，形状为 [M, 2]。 layout_kps (torch.Tensor): 版图关键点，形状为 [N, 2]。 返回： tuple: (inlier_matches, num_inliers) - inlier_matches: [(i_template, i_layout)]，内点匹配对。 - num_inliers: int，内点数量。 """ src_pts = np.array([template_kps[i].cpu().numpy() for i, _ in matches]) dst_pts = np.array([layout_kps[j].cpu().numpy() for _, j in matches]) if len(src_pts) < 4: return [], 0 try: H, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, ransacReprojThreshold=5.0) if H is None: return [], 0 inliers = mask.ravel() > 0 num_inliers = np.sum(inliers) inlier_matches = [matches[k] for k in range(len(matches)) if inliers[k]] return inlier_matches, num_inliers except cv2.error: return [], 0 def match_template_to_layout(model, layout_image, template_image): """ 使用 RoRD 模型执行模板匹配，迭代找到所有匹配并屏蔽已匹配区域。 参数： model (RoRD): RoRD 模型。 layout_image (torch.Tensor): 版图图像张量，形状为 [1, 1, H_layout, W_layout]。 template_image (torch.Tensor): 模板图像张量，形状为 [1, 1, H_template, W_template]。 返回： list: [{'x': x_min, 'y': y_min, 'width': w, 'height': h}]，所有检测到的边框。 """ # 提取版图和模板的关键点和描述子 layout_kps, layout_descs = extract_keypoints_and_descriptors(model, layout_image) template_kps, template_descs = extract_keypoints_and_descriptors(model, template_image) # 初始化活动版图关键点掩码 active_layout = torch.ones(len(layout_kps), dtype=bool) bboxes = [] while True: # 获取当前活动的版图关键点和描述子 current_layout_kps = layout_kps[active_layout] current_layout_descs = layout_descs[active_layout] if len(current_layout_descs) == 0: break # MNN 匹配 matches = mutual_nearest_neighbor(template_descs, current_layout_descs) if len(matches) == 0: break # 将当前版图索引映射回原始版图索引 active_indices = torch.nonzero(active_layout).squeeze(1) matches_original = [(i_template, active_indices[i_layout].item()) for i_template, i_layout in matches] # RANSAC 过滤 inlier_matches, num_inliers = ransac_filter(matches_original, template_kps, layout_kps) if num_inliers > 10: # 设置内点阈值 # 获取内点在版图中的关键点 inlier_layout_kps = [layout_kps[j].cpu().numpy() for _, j in inlier_matches] inlier_layout_kps = np.array(inlier_layout_kps) # 计算边框 x_min = int(inlier_layout_kps[:, 0].min()) y_min = int(inlier_layout_kps[:, 1].min()) x_max = int(inlier_layout_kps[:, 0].max()) y_max = int(inlier_layout_kps[:, 1].max()) bboxes.append({'x': x_min, 'y': y_min, 'width': x_max - x_min, 'height': y_max - y_min}) # 屏蔽内点 for _, j in inlier_matches: active_layout[j] = False else: break return bboxes if __name__ == "__main__": # 设置变换 transform = transforms.Compose([ SobelTransform(), transforms.ToTensor(), transforms.Normalize(mean=[0.5], std=[0.5]) ]) # 加载模型 model = RoRD().cuda() model.load_state_dict(torch.load('path/to/weights.pth')) model.eval() # 加载版图和模板图像 layout_image = Image.open('path/to/layout.png').convert('L') layout_tensor = transform(layout_image).unsqueeze(0).cuda() template_image = Image.open('path/to/template.png').convert('L') template_tensor = transform(template_image).unsqueeze(0).cuda() # 执行匹配 detected_bboxes = match_template_to_layout(model, layout_tensor, template_tensor) # 打印检测到的边框 print("检测到的边框：") for bbox in detected_bboxes: print(bbox)