# Example code for running inference on a pre-trained model import os import yaml import numpy as np import cv2 import torch from torchvision import transforms from PIL import Image import matplotlib.pyplot as plt from model.net.super_retina import SuperRetina from utils.image_process import * from utils.common_util import * device = torch.device('cuda' if torch.cuda.is_available() else "cpu") def sigmoid(arr): return 1. / (1 + np.exp(-arr)) class Inference(object): def __init__(self, model_path): self.model_path = model_path config_path = os.path.join(model_path, 'config.yaml') if os.path.exists(config_path): with open(config_path) as f: predict_config = yaml.safe_load(f) self.config = predict_config['PREDICT'] else: raise FileNotFoundError("Config File doesn't Exist") self.nms_size = self.config['nms_size'] self.nms_thresh = self.config['nms_thresh'] self.scale = 8 self.knn_thresh = self.config['knn_thresh'] self.model_image_width = self.config['model_image_width'] self.model_image_height = self.config['model_image_height'] self.transformer = transforms.Compose([ transforms.Resize((self.model_image_height, self.model_image_width)), transforms.ToTensor(), ]) self.model = self.load_model() self.device = device self.knn_matcher = cv2.BFMatcher(cv2.NORM_L2) def load_model(self): model_save_path = os.path.join(self.model_path, 'SuperRetina_released.pth') print('Loading model from {}'.format(model_save_path)) model = SuperRetina(config=self.config, encoder_device=device, detector_device=device, descriptor_device=device) checkpoint = torch.load(model_save_path, map_location=device) # 处理权重文件，适配模型结构 processed_weights = {} for k, v in checkpoint['net'].items(): # 处理detector head中的DoubleConv模块权重 if any(layer in k for layer in ['dconv_up3', 'dconv_up2', 'dconv_up1']): # 对于DoubleConv模块，需要添加.conv.部分 if '.0.' in k or '.2.' in k: # 0和2是Sequential中的卷积层索引 # 将"dconv_upX.Y.weight"转换为"dconv_upX.conv.Y.weight" parts = k.split('.') module_name = parts[0] # dconv_up3, dconv_up2, 或 dconv_up1 layer_idx = parts[1] # 0 或 2 weight_type = parts[2] # weight 或 bias new_key = f"{module_name}.conv.{layer_idx}.{weight_type}" processed_weights[new_key] = v else: processed_weights[k] = v else: processed_weights[k] = v # 检查权重文件的结构 if all(k.startswith('encoder.') or k.startswith('desc_head.') or k.startswith('det_head.') for k in processed_weights.keys()): # 权重文件已经有正确的模块前缀 encoder_state = {k: v for k, v in processed_weights.items() if k.startswith('encoder.')} desc_head_state = {k: v for k, v in processed_weights.items() if k.startswith('desc_head.')} det_head_state = {k: v for k, v in processed_weights.items() if k.startswith('det_head.')} # 移除前缀以匹配模型结构 encoder_state = {k.replace('encoder.', ''): v for k, v in encoder_state.items()} desc_head_state = {k.replace('desc_head.', ''): v for k, v in desc_head_state.items()} det_head_state = {k.replace('det_head.', ''): v for k, v in det_head_state.items()} else: # 权重文件没有模块前缀，根据层名称分类 encoder_state = {k: v for k, v in processed_weights.items() if any(layer in k for layer in ['conv1a', 'conv1b', 'conv2a', 'conv2b', 'conv3a', 'conv3b', 'conv4a', 'conv4b'])} desc_head_state = {k: v for k, v in processed_weights.items() if any(layer in k for layer in ['convDa', 'convDb', 'convDc', 'trans_conv'])} det_head_state = {k: v for k, v in processed_weights.items() if any(layer in k for layer in ['dconv_up3', 'dconv_up2', 'dconv_up1', 'conv_last'])} # 分别加载到对应设备上的模块 model.encoder.load_state_dict(encoder_state) model.desc_head.load_state_dict(desc_head_state) model.det_head.load_state_dict(det_head_state) model.to(device) model.eval() return model def load_image(self, image): # Load the image and preprocess it if isinstance(image, str): print('Loading image from {}'.format(image)) image = cv2.imread(image, cv2.IMREAD_COLOR) # --- 新增代码 Start --- # 强制将图像缩放到 config 中定义的模型输入尺寸 # 注意 cv2.resize 参数顺序是 (width, height) if hasattr(self, 'model_image_width') and hasattr(self, 'model_image_height'): image = cv2.resize(image, (self.model_image_width, self.model_image_height)) # --- 新增代码 End --- # 如果是灰度图则不额外处理，如果是RGB图像则取其中的绿色通道 if len(image.shape) == 3 and image.shape[2] == 3: image = image[:, :, 1] image = pre_processing(image) image = (image * 255).astype(np.uint8) # 打印调试信息确认尺寸 print(f"Processed image shape: {image.shape}") return image def model_run_pair(self, query_image, refer_image): query_tensor = self.transformer(Image.fromarray(query_image)) refer_tensor = self.transformer(Image.fromarray(refer_image)) inputs = torch.cat((query_tensor.unsqueeze(0), refer_tensor.unsqueeze(0))) inputs = inputs.to(self.device) with torch.no_grad(): detector_pred, descriptor_pred = self.model(inputs) scores = simple_nms(detector_pred, self.nms_size) b, _, h, w = detector_pred.shape scores = scores.reshape(-1, h, w) keypoints = [ torch.nonzero(s > self.nms_thresh) for s in scores] scores = [s[tuple(k.t())] for s, k in zip(scores, keypoints)] # Discard keypoints near the image borders keypoints, scores = list(zip(*[ remove_borders(k, s, 4, h, w) for k, s in zip(keypoints, scores)])) keypoints = [torch.flip(k, [1]).float().data for k in keypoints] descriptors = [sample_keypoint_desc(k[None], d[None], 8)[0].cpu() for k, d in zip(keypoints, descriptor_pred)] keypoints = [k.cpu() for k in keypoints] # 将PyTorch张量转换为NumPy数组 query_tensor = query_tensor.cpu().numpy() refer_tensor = refer_tensor.cpu().numpy() return keypoints, descriptors, query_tensor, refer_tensor def match(self, keypoints, descriptors): query_keypoints, refer_keypoints = keypoints[0], keypoints[1] query_desc, refer_desc = descriptors[0].permute(1, 0).numpy(), descriptors[1].permute(1, 0).numpy() # mapping keypoints to scaled keypoints cv_kpts_query = [cv2.KeyPoint(int(i[0]), int(i[1]), 30) for i in query_keypoints] cv_kpts_refer = [cv2.KeyPoint(int(i[0]), int(i[1]), 30) for i in refer_keypoints] goodMatch = [] status = [] matches = [] try: matches = self.knn_matcher.knnMatch(query_desc, refer_desc, k=2) for m, n in matches: if m.distance < self.knn_thresh * n.distance: goodMatch.append(m) status.append(True) else: status.append(False) except Exception: pass print(len(goodMatch)) print(len(cv_kpts_query)) print(len(cv_kpts_refer)) return goodMatch, cv_kpts_query, cv_kpts_refer, status def compute_homography(self, goodMatch, cv_kpts_query, cv_kpts_refer): H_m = None inliers_num_rate = 0 mask = None if len(goodMatch) >= 4: src_pts = [cv_kpts_query[m.queryIdx].pt for m in goodMatch] src_pts = np.float32(src_pts).reshape(-1, 1, 2) dst_pts = [cv_kpts_refer[m.trainIdx].pt for m in goodMatch] dst_pts = np.float32(dst_pts).reshape(-1, 1, 2) H_m, mask = cv2.findHomography(src_pts, dst_pts, cv2.LMEDS) goodMatch = np.array(goodMatch)[mask.ravel() == 1] inliers_num_rate = mask.sum() / len(mask.ravel()) print(H_m, inliers_num_rate) return H_m, inliers_num_rate, goodMatch, mask def align_image_pair(self, H_m, inliers_num_rate, query_image, refer_image): if H_m is not None: h, w = self.model_image_height, self.model_image_width # 确保图像格式正确 - 从 (channels, height, width) 转换为 (height, width, channels) if len(query_image.shape) == 3 and query_image.shape[0] in [1, 3]: # 如果是 (channels, height, width) 格式，转换为 (height, width, channels) query_image = query_image.transpose(1, 2, 0) # 如果是单通道，去掉通道维度 if query_image.shape[2] == 1: query_image = query_image.squeeze(2) if len(refer_image.shape) == 3 and refer_image.shape[0] in [1, 3]: # 如果是 (channels, height, width) 格式，转换为 (height, width, channels) refer_image = refer_image.transpose(1, 2, 0) # 如果是单通道，去掉通道维度 if refer_image.shape[2] == 1: refer_image = refer_image.squeeze(2) # 确保图像是 uint8 类型 if query_image.dtype != np.uint8: query_image = (query_image * 255).astype(np.uint8) if refer_image.dtype != np.uint8: refer_image = (refer_image * 255).astype(np.uint8) print(f"query_image shape after conversion: {query_image.shape}") print(f"refer_image shape after conversion: {refer_image.shape}") try: query_align = cv2.warpPerspective(query_image, H_m, (w, h), borderMode=cv2.BORDER_CONSTANT, borderValue=(0)) merged = np.zeros((h, w, 3), dtype=np.uint8) # 处理查询图像 if len(query_align.shape) == 3: query_align_gray = cv2.cvtColor(query_align, cv2.COLOR_BGR2GRAY) else: query_align_gray = query_align # 处理参考图像 if len(refer_image.shape) == 3: refer_gray = cv2.cvtColor(refer_image, cv2.COLOR_BGR2GRAY) else: refer_gray = refer_image merged[:, :, 0] = query_align_gray merged[:, :, 1] = refer_gray return merged except Exception as e: print(f"图像对齐过程中出错: {e}") return None print("Matched Failed!") return None def draw_result(self, query_image, refer_image, cv_kpts_query, cv_kpts_refer, matches, status): def drawMatches(imageA, imageB, kpsA, kpsB, matches, status): # initialize the output visualization image (hA, wA) = imageA.shape[:2] (hB, wB) = imageB.shape[:2] print("query_image shape:", imageA.shape) print("refer_image shape:", imageB.shape) vis = np.zeros((max(hA, hB), wA + wB, 3), dtype="uint8") if len(imageA.shape) == 2: imageA = cv2.cvtColor(imageA, cv2.COLOR_GRAY2RGB) imageB = cv2.cvtColor(imageB, cv2.COLOR_GRAY2RGB) vis[0:hA, 0:wA] = imageA vis[0:hB, wA:] = imageB # loop over the matches for match, s in zip(matches, status): trainIdx, queryIdx = match.trainIdx, match.queryIdx # only process the match if the keypoint was successfully # matched if s == 1 or s == True: # draw the match ptA = (int(kpsA[queryIdx].pt[0]), int(kpsA[queryIdx].pt[1])) ptB = (int(kpsB[trainIdx].pt[0]) + wA, int(kpsB[trainIdx].pt[1])) cv2.line(vis, ptA, ptB, (0, 255, 0), 2) # return the visualization return vis query_np = np.array([kp.pt for kp in cv_kpts_query]) refer_np = np.array([kp.pt for kp in cv_kpts_refer]) refer_np[:, 0] += query_image.shape[1] matched_image = drawMatches(query_image, refer_image, cv_kpts_query, cv_kpts_refer, matches, status) plt.figure(dpi=300) plt.scatter(query_np[:, 0], query_np[:, 1], s=1, c='r') plt.scatter(refer_np[:, 0], refer_np[:, 1], s=1, c='r') plt.axis('off') plt.title('Match Result, #goodMatch: {}'.format(np.sum(status))) plt.imshow(cv2.cvtColor(matched_image, cv2.COLOR_BGR2RGB)) # Instead of showing the plot, we return the image as a numpy array # Convert the plot to a numpy array plt.tight_layout() fig = plt.gcf() fig.canvas.draw() # Get the RGBA buffer from the figure w, h = fig.canvas.get_width_height() buf = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8) buf.shape = (h, w, 3) plt.close() return buf def inference(self, query_image, refer_image): try: query_processed = self.load_image(query_image) refer_processed = self.load_image(refer_image) keypoints, descriptors, query_tensor, refer_tensor = self.model_run_pair(query_processed, refer_processed) goodMatch, cv_kpts_query, cv_kpts_refer, status = self.match(keypoints, descriptors) H_m, inliers_num_rate, filtered_goodMatch, mask = self.compute_homography(goodMatch, cv_kpts_query, cv_kpts_refer) # 创建过滤后的状态列表，只包含通过单应性矩阵验证的匹配点 if mask is not None: filtered_status = [s for s, m in zip(status, mask.ravel()) if m == 1] else: filtered_status = [] merged = self.align_image_pair(H_m, inliers_num_rate, query_tensor, refer_tensor) match_show = self.draw_result(query_processed, refer_processed, cv_kpts_query, cv_kpts_refer, filtered_goodMatch, filtered_status) return merged, match_show except Exception as e: print(f"推理过程中出现错误: {e}") import traceback traceback.print_exc() # 返回错误信息或默认图像 h, w = self.model_image_height, self.model_image_height error_image = np.zeros((h, w, 3), dtype=np.uint8) cv2.putText(error_image, f"Error: {str(e)}", (10, h//2), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) return error_image, error_image