48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
|
import faiss
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
# 导入 models.rotation_cnn 模块中的 RotationInvariantNet 类
|
|||
|
|
from models.rotation_cnn import RotationInvariantNet
|
|||
|
|
|
|||
|
|
from models.rotation_cnn import get_rotational_features
|
|||
|
|
|
|||
|
|
# 导入 data_utils 中的 layout_to_tensor 函数（假设该函数存在）
|
|||
|
|
from data_units import layout_to_tensor # 如果 data_utils.py 存在此函数
|
|||
|
|
|
|||
|
|
from data_units import tile_layout
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
# 配置参数（需根据实际调整）
|
|||
|
|
block_size = 64 # 分块尺寸
|
|||
|
|
target_module_path = "target.png"
|
|||
|
|
large_layout_path = "layout_large.png"
|
|||
|
|
|
|||
|
|
# 加载模型
|
|||
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|||
|
|
model = RotationInvariantNet().to(device)
|
|||
|
|
model.load_state_dict(torch.load("rotation_cnn.pth"))
|
|||
|
|
model.eval()
|
|||
|
|
|
|||
|
|
# 预处理目标模块与大版图
|
|||
|
|
target_tensor = layout_to_tensor(target_module_path, (block_size, block_size))
|
|||
|
|
target_feat = get_rotational_features(model, torch.tensor(target_tensor).to(device))
|
|||
|
|
|
|||
|
|
large_layout = layout_to_tensor(large_layout_path)
|
|||
|
|
tiles = tile_layout(large_layout)
|
|||
|
|
|
|||
|
|
# 构建特征索引（使用Faiss加速）
|
|||
|
|
index = faiss.IndexFlatL2(64) # 特征维度由模型决定
|
|||
|
|
features_db = []
|
|||
|
|
for (x,y,tile) in tiles:
|
|||
|
|
feat = get_rotational_features(model, torch.tensor(tile).to(device))
|
|||
|
|
features_db.append(feat)
|
|||
|
|
index.add(np.stack(features_db))
|
|||
|
|
|
|||
|
|
# 检索相似区域
|
|||
|
|
D, I = index.search(target_feat[np.newaxis,:], k=10)
|
|||
|
|
for idx in I[0]:
|
|||
|
|
x,y,_ = tiles[idx]
|
|||
|
|
print(f"匹配区域坐标: ({x}, {y}), 相似度: {D[0][idx]}")
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|