import time import cv2 import skimage import os import torch from PIL import Image import matplotlib.pyplot as plt import matplotlib as mpl import numpy as np from models.line_net.line_net import linenet_resnet50_fpn, get_line_net_efficientnetv2, get_line_net_convnext_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'],strict=False) # 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_ = [] 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]): 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 # box内无线段时，选box内点组成线段最长的 两个点组成的线段返回 def box_line1(imgs, pred): # 默认置信度 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] points = pred[-1]['wires']['juncs'].cpu().numpy()[0] / 128 * 512 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_ = [] 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]): if score[j] > score_max: tmp = line[j] score_max = score[j] # # 如果 box 内无线段，则通过点坐标找最长线段 # if score_max == 0.0: # 说明 box 内无线段 # box_points = [ # [x, y] for x, y in points # if i[0] <= y <= i[2] and i[1] <= x <= i[3] # ] # # if len(box_points) >= 2: # 至少需要两个点才能组成线段 # max_distance = 0.0 # longest_segment = [[0.0, 0.0], [0.0, 0.0]] # # # 找出 box 内点组成的最长线段 # for p1 in box_points: # for p2 in box_points: # if p1 != p2: # distance = np.linalg.norm(np.array(p1) - np.array(p2)) # if distance > max_distance: # max_distance = distance # longest_segment = [p1, p2] # # tmp = longest_segment # score_max = 0.0 # 默认分数为 0.0 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 show_box(imgs, pred, t_start): col = [ '#1f77b4', '#aec7e8', '#ff7f0e', '#ffbb78', '#2ca02c', '#98df8a', '#d62728', '#ff9896', '#9467bd', '#c5b0d5', '#8c564b', '#c49c94', '#e377c2', '#f7b6d2', '#7f7f7f', '#c7c7c7', '#bcbd22', '#dbdb8d', '#17becf', '#9edae5', '#8dd3c7', '#ffffb3', '#bebada', '#fb8072', '#80b1d3', '#fdb462', '#b3de69', '#fccde5', '#bc80bd', '#ccebc5', '#ffed6f', '#8da0cb', '#e78ac3', '#e5c494', '#b3b3b3', '#fdbf6f', '#ff7f00', '#cab2d6', '#637939', '#b5cf6b', '#cedb9c', '#8c6d31', '#e7969c', '#d6616b', '#7b4173', '#ad494a', '#843c39', '#dd8452', '#f7f7f7', '#cccccc', '#969696', '#525252', '#f7fcfd', '#e5f5f9', '#ccece6', '#99d8c9', '#66c2a4', '#2ca25f', '#008d4c', '#005a32', '#f7fcf0', '#e0f3db', '#ccebc5', '#a8ddb5', '#7bccc4', '#4eb3d3', '#2b8cbe', '#08589e', '#f7fcfd', '#e0ecf4', '#bfd3e6', '#9ebcda', '#8c96c6', '#8c6bb4', '#88419d', '#810f7c', '#4d004b', '#f7f7f7', '#efefef', '#d9d9d9', '#bfbfbf', '#969696', '#737373', '#525252', '#252525', '#000000', '#ffffff', '#ffeda0', '#fed976', '#feb24c', '#fd8d3c', '#fc4e2a', '#e31a1c', '#bd0026', '#800026', '#ff6f61', '#ff9e64', '#ff6347', '#ffa07a', '#fa8072' ] # print(len(col)) im = imgs.permute(1, 2, 0) boxes = pred[0]['boxes'].cpu().numpy() box_scores = pred[0]['scores'].cpu().numpy() # 可视化预测结 fig, ax = plt.subplots(figsize=(10, 10)) ax.imshow(np.array(im)) for idx, box in enumerate(boxes): # if box_scores[idx] < 0.7: # continue x0, y0, x1, y1 = box ax.add_patch( plt.Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, edgecolor=col[idx], linewidth=1)) t_end = time.time() print(f'show_box used:{t_end - t_start}') plt.show() def show_predict(imgs, pred, t_start): col = [ '#1f77b4', '#aec7e8', '#ff7f0e', '#ffbb78', '#2ca02c', '#98df8a', '#d62728', '#ff9896', '#9467bd', '#c5b0d5', '#8c564b', '#c49c94', '#e377c2', '#f7b6d2', '#7f7f7f', '#c7c7c7', '#bcbd22', '#dbdb8d', '#17becf', '#9edae5', '#8dd3c7', '#ffffb3', '#bebada', '#fb8072', '#80b1d3', '#fdb462', '#b3de69', '#fccde5', '#bc80bd', '#ccebc5', '#ffed6f', '#8da0cb', '#e78ac3', '#e5c494', '#b3b3b3', '#fdbf6f', '#ff7f00', '#cab2d6', '#637939', '#b5cf6b', '#cedb9c', '#8c6d31', '#e7969c', '#d6616b', '#7b4173', '#ad494a', '#843c39', '#dd8452', '#f7f7f7', '#cccccc', '#969696', '#525252', '#f7fcfd', '#e5f5f9', '#ccece6', '#99d8c9', '#66c2a4', '#2ca25f', '#008d4c', '#005a32', '#f7fcf0', '#e0f3db', '#ccebc5', '#a8ddb5', '#7bccc4', '#4eb3d3', '#2b8cbe', '#08589e', '#f7fcfd', '#e0ecf4', '#bfd3e6', '#9ebcda', '#8c96c6', '#8c6bb4', '#88419d', '#810f7c', '#4d004b', '#f7f7f7', '#efefef', '#d9d9d9', '#bfbfbf', '#969696', '#737373', '#525252', '#252525', '#000000', '#ffffff', '#ffeda0', '#fed976', '#feb24c', '#fd8d3c', '#fc4e2a', '#e31a1c', '#bd0026', '#800026', '#ff6f61', '#ff9e64', '#ff6347', '#ffa07a', '#fa8072' ] im = imgs.permute(1, 2, 0) # 处理为 [512, 512, 3] boxes = pred[0]['boxes'].cpu().numpy() box_scores = pred[0]['scores'].cpu().numpy() lines = pred[0]['line'].cpu().numpy() line_scores = pred[0]['line_score'].cpu().numpy() # 可视化预测结 fig, ax = plt.subplots(figsize=(10, 10)) ax.imshow(np.array(im)) idx = 0 tmp = np.array([[0.0, 0.0], [0.0, 0.0]]) for box, line, box_score, line_score in zip(boxes, lines, box_scores, line_scores): x0, y0, x1, y1 = box # 框中无线的跳过 if np.array_equal(line, tmp): continue a, b = line if box_score >= 0 or line_score >= 0: ax.add_patch( plt.Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, edgecolor=col[idx], linewidth=1)) ax.scatter(a[1], a[0], c='#871F78', s=10) ax.scatter(b[1], b[0], c='#871F78', s=10) ax.plot([a[1], b[1]], [a[0], b[0]], c=col[idx], linewidth=1) idx = idx + 1 t_end = time.time() print(f'predict used:{t_end - t_start}') plt.show() def show_line(imgs, pred, t_start): im = imgs.permute(1, 2, 0) lines = pred[-1]['wires']['lines'][0].cpu().numpy() / 128 * 512 # print(pred[-1]['wires']['score']) scores = pred[-1]['wires']['score'].cpu().numpy()[0] t1 = time.time() print(f't1:{t1 - t_start}') for i in range(1, len(lines)): if (lines[i] == lines[0]).all(): lines = lines[:i] scores = scores[:i] break diag = (im.shape[0] ** 2 + im.shape[1] ** 2) ** 0.5 line, line_score = postprocess(lines, scores, diag * 0.01, 0, False) print(f'lines num:{len(line)}') t2 = time.time() print(f't1:{t2 - t1}') # 可视化预测结 fig, ax = plt.subplots(figsize=(10, 10)) ax.imshow(np.array(im)) for idx, (a, b) in enumerate(line): # if line_score[idx] < 0.7: # continue ax.scatter(a[1], a[0], c='#871F78', s=2) ax.scatter(b[1], b[0], c='#871F78', s=2) ax.plot([a[1], b[1]], [a[0], b[0]], c='red', linewidth=1) t_end = time.time() print(f'show_line used:{t_end - t_start}') plt.show() 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] # 将图像调整为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) im = img_tensor.permute(1, 2, 0) # [H, W, 3] if im.shape != (512, 512, 3): # im = cv2.resize(im.cpu().numpy().astype(np.float32), (512, 512), interpolation=cv2.INTER_NEAREST) im = cv2.resize(im.cpu().numpy().astype(np.float32), (512, 512), interpolation=cv2.INTER_NEAREST_EXACT) img_ = torch.tensor(im).permute(2, 0, 1) # [3, 512, 512] t_end = time.time() print(f'switch img used:{t_end - t_start}') with torch.no_grad(): predictions = model([img_.to(device)]) print(predictions) t_end1 = time.time() print(f'model test used:{t_end1 - t_end}') # show_line_optimized(img_, predictions, t_start) # 只画线 show_line(img_, predictions, t_end1) t_end2 = time.time() show_box(img_, predictions, t_end2) # 只画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_line1(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) # model = get_line_net_efficientnetv2(2, pretrained_backbone=True).to(device) model=get_line_net_convnext_fpn(num_classes=2).to(device) # pt_path = r"C:\Users\m2337\Downloads\best_lmap代替x，训练24轮结果.pth" # pt_path = r"C:\Users\m2337\Downloads\best_lmap代替x，训练75轮.pth" pt_path = r"\\192.168.50.222\share\rlq\weights\convnext25051401.pth" # pt_path = r"C:\Users\m2337\Downloads\best_e20.pth" img_path = r"\\192.168.50.222\share\zyh\513\a_dataset\images\val\2025-05-13-08-56-03_LaserData_ID019504_color_scale.jpg" predict(pt_path, model, img_path) t_end = time.time() print(f'predict used:{t_end - t_start}')