199 lines
7.2 KiB
Python
199 lines
7.2 KiB
Python
|
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)
|