pokouqiege/models/line_detect/predict2.py

# import time
#
# from models.line_detect.postprocess import show_predict
# 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 rtree import index
# # from models.wirenet.postprocess import postprocess
# from models.wirenet.postprocess import postprocess
#
# 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_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, (x_min, y_min, x_max, y_max))
#
#     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 box_line_(pred):
# #     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]
# #         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
#
#
# 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)
#
#     with torch.no_grad():
#         t_start = time.time()
#         predictions = model([img_tensor.to(device)])
#         t_end=time.time()
#         print(f'predict used:{t_end-t_start}')
#         # print(f'predictions:{predictions}')
#         boxes=predictions[0]['boxes'].shape
#         lines=predictions[-1]['wires']['lines'].shape
#         lines_scores=predictions[-1]['wires']['score'].shape
#         print(f'predictions boxes:{boxes},lines:{lines},lines_scores:{lines_scores}')
#     t_start=time.time()
#     pred = box_line_optimized(predictions)
#     t_end=time.time()
#     print(f'matched boxes and lines used:{t_end - t_start}')
#     # print(f'pred:{pred[0]}')
#     show_predict(img_tensor, 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"F:\BaiduNetdiskDownload\resnet50_best_e8.pth"
#     img_path = r"I:\datasets\wirenet_1000\images\val\00035148_0.png"
#     predict(pt_path, model, img_path)
#     t_end = time.time()
#     # print(f'predict used:{t_end - t_start}')


import time

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_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_ = []
        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, 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]

    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)
    line = pred[-1]['wires']['lines'][0].cpu().numpy() / 128 * 512
    # print(pred[-1]['wires']['score'])
    line_score = pred[-1]['wires']['score'].cpu().numpy()[0]

    t1 = time.time()
    print(f't1:{t1 - t_start}')
    diag = (im.shape[0] ** 2 + im.shape[1] ** 2) ** 0.5
    line, line_score = postprocess(line, line_score, 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)
    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_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"C:\Users\m2337\Downloads\best_lmap代替x，训练24轮结果.pth"
    # pt_path = r"C:\Users\m2337\Downloads\best_lmap代替x，训练75轮.pth"
    pt_path = r"C:\Users\m2337\Downloads\best_e150.pth"
    # pt_path = r"C:\Users\m2337\Downloads\best_e20.pth"
    img_path = r"C:\Users\m2337\Desktop\p\新建文件夹\2025-03-25-16-10-00_SaveLeftImage.png"
    predict(pt_path, model, img_path)
    t_end = time.time()
    print(f'predict used:{t_end - t_start}')