import time import skimage from models.line_detect.postprocess import show_predict, show_box, show_box_or_line, show_box_and_line, \ show_line_optimized, show_line import os import torch from PIL import Image import matplotlib.pyplot as plt import matplotlib as mpl import numpy as np from models.line_detect.line_net import linenet_resnet50_fpn from torchvision import transforms # from models.wirenet.postprocess import postprocess from models.wirenet.postprocess import postprocess from rtree import index device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') def load_best_model(model, save_path, device): if os.path.exists(save_path): checkpoint = torch.load(save_path, map_location=device) model.load_state_dict(checkpoint['model_state_dict']) # if optimizer is not None: # optimizer.load_state_dict(checkpoint['optimizer_state_dict']) # epoch = checkpoint['epoch'] # loss = checkpoint['loss'] # print(f"Loaded best model from {save_path} at epoch {epoch} with loss {loss:.4f}") else: print(f"No saved model found at {save_path}") return model def box_line_(imgs, pred, length=False): # 默认置信度 im = imgs.permute(1, 2, 0).cpu().numpy() line_data = pred[-1]['wires']['lines'][0].cpu().numpy() / 128 * 512 line_scores = pred[-1]['wires']['score'].cpu().numpy()[0] diag = (im.shape[0] ** 2 + im.shape[1] ** 2) ** 0.5 line, score = postprocess(line_data, line_scores, diag * 0.01, 0, False) for idx, box_ in enumerate(pred[0:-1]): box = box_['boxes'] # 是一个tensor # line = pred[-1]['wires']['lines'][idx].cpu().numpy() / 128 * 512 # score = pred[-1]['wires']['score'][idx] # diag = (512 ** 2 + 512 ** 2) ** 0.5 # line, score = postprocess(line, score, diag * 0.01, 0, False) line_ = [] score_ = [] for i in box: score_max = 0.0 tmp = [[0.0, 0.0], [0.0, 0.0]] for j in range(len(line)): if (line[j][0][1] >= i[0] and line[j][1][1] >= i[0] and line[j][0][1] <= i[2] and line[j][1][1] <= i[2] and line[j][0][0] >= i[1] and line[j][1][0] >= i[1] and line[j][0][0] <= i[3] and line[j][1][0] <= i[3]): # # 计算线段长度 # length = np.linalg.norm(line[j][0] - line[j][1]) # if length > score_max: # tmp = line[j] # score_max = score[j] if score[j] > score_max: tmp = line[j] score_max = score[j] line_.append(tmp) score_.append(score_max) processed_list = torch.tensor(line_) pred[idx]['line'] = processed_list processed_s_list = torch.tensor(score_) pred[idx]['line_score'] = processed_s_list return pred def box_line_optimized(pred): # 创建R-tree索引 idx = index.Index() # 将所有线段添加到R-tree中 lines = pred[-1]['wires']['lines'] # 形状为[1, 2500, 2, 2] scores = pred[-1]['wires']['score'][0] # 假设形状为[2500] # 提取并处理所有线段 for idx_line in range(lines.shape[1]): # 遍历2500条线段 line_tensor = lines[0, idx_line].cpu().numpy() / 128 * 512 # 转换为numpy数组并调整比例 x_min = float(min(line_tensor[0][0], line_tensor[1][0])) y_min = float(min(line_tensor[0][1], line_tensor[1][1])) x_max = float(max(line_tensor[0][0], line_tensor[1][0])) y_max = float(max(line_tensor[0][1], line_tensor[1][1])) idx.insert(idx_line, (max(0, x_min - 256), max(0, y_min - 256), min(512, x_max + 256), min(512, y_max + 256))) for idx_box, box_ in enumerate(pred[0:-1]): box = box_['boxes'].cpu().numpy() # 确保将张量转换为numpy数组 line_ = [] score_ = [] for i in box: score_max = 0.0 tmp = [[0.0, 0.0], [0.0, 0.0]] # 获取与当前box可能相交的所有线段 possible_matches = list(idx.intersection((i[0], i[1], i[2], i[3]))) for j in possible_matches: line_j = lines[0, j].cpu().numpy() / 128 * 512 if (line_j[0][1] >= i[0] and line_j[1][1] >= i[0] and # 注意这里交换了x和y line_j[0][1] <= i[2] and line_j[1][1] <= i[2] and line_j[0][0] >= i[1] and line_j[1][0] >= i[1] and line_j[0][0] <= i[3] and line_j[1][0] <= i[3]): if scores[j] > score_max: tmp = line_j score_max = scores[j] line_.append(tmp) score_.append(score_max) processed_list = torch.tensor(line_) pred[idx_box]['line'] = processed_list processed_s_list = torch.tensor(score_) pred[idx_box]['line_score'] = processed_s_list return pred def predict(pt_path, model, img): model = load_best_model(model, pt_path, device) model.eval() if isinstance(img, str): img = Image.open(img).convert("RGB") transform = transforms.ToTensor() img_tensor = transform(img) # [3, 512, 512] # img_ = img_tensor # 将图像调整为512x512大小 t_start = time.time() im = img_tensor.permute(1, 2, 0) # [512, 512, 3] im_resized = skimage.transform.resize(im.cpu().numpy().astype(np.float32), (512, 512)) # (512, 512, 3) img_ = torch.tensor(im_resized).permute(2, 0, 1) t_end = time.time() print(f'switch img used:{t_end - t_start}') with torch.no_grad(): predictions = model([img_.to(device)]) # print(predictions) # show_line_optimized(img_, predictions, t_start) # 只画线 show_line(img_, predictions, t_start) # show_box(img_, predictions, t_start) # 只画kuang # show_box_or_line(img_, predictions, show_line=True, show_box=True) # 参数确定画什么 # show_box_and_line(img_, predictions, show_line=True, show_box=True) # 一起画 1x2 2张图 t_start = time.time() # pred = box_line_optimized(predictions) pred = box_line_(img_, predictions) t_end = time.time() print(f'Matched boxes and lines used: {t_end - t_start:.4f} seconds') show_predict(img_, pred, t_start) if __name__ == '__main__': t_start = time.time() print(f'start to predict:{t_start}') model = linenet_resnet50_fpn().to(device) # pt_path = r'D:\python\PycharmProjects\20250214\weight\resnet50_best_e100.pth' # pt_path = r'D:\python\PycharmProjects\linenet_wts\r50fpn_wts_e350\best.pth' pt_path = r'D:\python\PycharmProjects\20250214\weight\merged_model_weights.pth' # img_path = f'D:\python\PycharmProjects\data2\images/train/2024-11-27-15-43-13_SaveImage.png' # 工件图 # img_path = f'D:\python\PycharmProjects\data\images/train/00558656_3.png' # wireframe图 img_path = r'C:\Users\m2337\Desktop\p\112941.jpg' predict(pt_path, model, img_path) t_end = time.time() print(f'predict used:{t_end - t_start}')