lcnn/models/line_detect/lm_0223.py

import os
import torch
import numpy as np
from PIL import Image
import skimage.io
import skimage.color
from torchvision import transforms
import shutil
import matplotlib.pyplot as plt
from models.line_detect.line_net import linenet_resnet50_fpn
from models.line_detect.postprocess import postprocess
from rtree import index
import time
import multiprocessing as mp

mp.set_start_method('spawn', force=True)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def load_model(model_path):
    model = linenet_resnet50_fpn().to(device)
    if os.path.exists(model_path):
        checkpoint = torch.load(model_path, map_location=device)
        model.load_state_dict(checkpoint['model_state_dict'])
        print(f"Loaded model from {model_path}")
    else:
        raise FileNotFoundError(f"No saved model found at {model_path}")
    model.eval()
    return model


def preprocess_image(image_path):
    img = Image.open(image_path).convert("RGB")
    transform = transforms.ToTensor()
    img_tensor = transform(img)
    resized_img = skimage.transform.resize(
        img_tensor.permute(1, 2, 0).cpu().numpy().astype(np.float32), (512, 512)
    )
    return torch.tensor(resized_img).permute(2, 0, 1)


def save_plot(output_path: str):
    plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
    print(f"Saved plot to {output_path}")
    plt.close()


def get_colors():
    """返回一组预定义的颜色列表"""
    return [
        '#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'
    ]


def process_box(box, lines, scores):
    """处理单个边界框，找到最佳匹配的线段"""
    valid_lines = []  # 存储有效的线段
    valid_scores = []  # 存储有效的分数
    # print(f'score:{len(scores)}')
    for i in box:
        best_line = None
        max_length = 0.0
        # 遍历所有线段
        for j in range(lines.shape[1]):
            line_j = lines[0, j].cpu().numpy() / 128 * 512
            # 检查线段是否完全在box内
            if (all(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])
                length = scores[j].item()
                # print(length)
                if length > max_length:
                    best_line = line_j
                    max_length = length
        if best_line is not None:
            valid_lines.append(best_line)
            valid_scores.append(max_length)

        else:
            valid_lines.append([[0.0, 0.0], [0.0, 0.0]])
            valid_scores.append(0.0)
    return valid_lines, valid_scores


def box_line_optimized_parallel(pred):
    """并行处理边界框和线段的匹配"""
    lines = pred[-1]['wires']['lines']  # 形状为[1, 2500, 2, 2]
    scores = pred[-1]['wires']['score'][0]  # 假设形状为[2500]
    boxes = [box_['boxes'].cpu().numpy() for box_ in pred[0:-1]]  # 所有box
    # num_processes = min(mp.cpu_count(), len(boxes))  # 使用可用的核心数
    # with mp.Pool(processes=num_processes) as pool:
    #     results = pool.starmap(
    #         process_box,
    #         [(box, lines, scores) for box in boxes]
    #     )
    results = process_box(boxes, lines, scores)
    # 更新预测结果
    filtered_pred = []
    for idx_box, (valid_lines, valid_scores) in enumerate(results):
        if valid_lines:
            pred[idx_box]['line'] = torch.tensor(valid_lines)
            pred[idx_box]['line_score'] = torch.tensor(valid_scores)
            filtered_pred.append(pred[idx_box])
    return filtered_pred


def predict(image_path):
    start_time = time.time()

    model_path = r'D:\python\PycharmProjects\20250214\weight\linenet_wts\r50fpn_wts_e350\best.pth'
    model = load_model(model_path)

    img_tensor = preprocess_image(image_path)
    im = img_tensor.permute(1, 2, 0).cpu().numpy()

    with torch.no_grad():
        predictions = model([img_tensor.to(device)])

    t_start = time.time()
    filtered_pred = box_line_optimized_parallel(predictions)
    t_end = time.time()
    print(f'Matched boxes and lines used: {t_end - t_start:.2f} seconds')

    output_path_box = show_box(im, predictions, t_start)
    output_path_line = show_line(im, predictions, t_start)
    show_predict(im, filtered_pred, t_start)

    # combined_image_path = "combined_result.png"
    # combine_images(
    #     [output_path_boxandline, output_path_box, output_path_line],
    #     titles=["Box and Line", "Box", "Line"],
    #     output_path=combined_image_path
    # )

    end_time = time.time()
    print(f'Total time: {end_time - start_time:.2f} seconds')


def combine_images(image_paths: list, titles: list, output_path: str):
    """将多个图像合并为一张图片"""
    fig, axes = plt.subplots(1, len(image_paths), figsize=(15, 5))
    for ax, img_path, title in zip(axes, image_paths, titles):
        ax.imshow(plt.imread(img_path))
        ax.set_title(title)
        ax.axis("off")
    plt.savefig(output_path, bbox_inches="tight", pad_inches=0)
    plt.close()


def show_box(im, predictions, t_start):
    """绘制边界框并保存结果"""
    boxes = predictions[0]['boxes'].cpu().numpy()
    box_scores = predictions[0]['scores'].cpu().numpy()
    colors = get_colors()
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(im)
    for idx, (box, score) in enumerate(zip(boxes, box_scores)):
        if score < 0.7:
            continue
        x0, y0, x1, y1 = box
        ax.add_patch(
            plt.Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, edgecolor=colors[idx % len(colors)], linewidth=1))
    t_end = time.time()
    print(f'show_box used: {t_end - t_start:.2f} seconds')
    plt.show()
    output_path = "temp_result_box.png"
    save_plot(output_path)
    return output_path


def show_line(im, predictions, t_start):
    """绘制线段并保存结果"""
    lines = predictions[-1]['wires']['lines'][0].cpu().numpy() / 128 * 512
    line_scores = predictions[-1]['wires']['score'].cpu().numpy()[0]
    diag = (im.shape[0] ** 2 + im.shape[1] ** 2) ** 0.5
    nlines, nscores = postprocess(lines, line_scores, diag * 0.01, 0, False)
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(im)
    for (a, b), s in zip(nlines, nscores):
        if s < 0.9:
            continue
        ax.scatter([a[1], b[1]], [a[0], 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:.2f} seconds')
    plt.show()
    output_path = "temp_result_line.png"
    save_plot(output_path)
    return output_path


# def show_predict(im, filtered_pred, t_start):
#     colors = get_colors()
#     fig, ax = plt.subplots(figsize=(10, 10))
#     ax.imshow(im)
#     for idx, pred in enumerate(filtered_pred):
#         boxes = pred['boxes'].cpu().numpy()
#         box_scores = pred['scores'].cpu().numpy()
#         lines = pred['line'].cpu().numpy()
#         line_scores = pred['line_score'].cpu().numpy()
#         diag = (im.shape[0] ** 2 + im.shape[1] ** 2) ** 0.5
#         nlines, nscores = postprocess(lines, line_scores, diag * 0.01, 0, False)
#         for box_idx, (box, line, box_score, line_score) in enumerate(zip(boxes, nlines, box_scores, nscores)):
#             if box_score < 0.7 or line_score < 0.9:
#                 continue
#
#             if line is None or len(line) == 0:
#                 continue
#
#             x0, y0, x1, y1 = box
#             a, b = line
#             color = colors[(idx + box_idx) % len(colors)]  # 每个边界框分配一个唯一颜色
#             ax.add_patch(plt.Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, edgecolor=color, linewidth=1))
#             ax.scatter(a[1], a[0], c=color, s=10)
#             ax.scatter(b[1], b[0], c=color, s=10)
#             ax.plot([a[1], b[1]], [a[0], b[0]], c=color, linewidth=1)
#     t_end = time.time()
#     print(f'show_predict used: {t_end - t_start:.2f} seconds')
#     plt.show()
#     output_path = "temp_result.png"
#     save_plot(output_path)
#     return output_path
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'
    ]
    print(imgs.shape)
    # 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(imgs))
    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.7 or line_score >= 0.9:
            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()


if __name__ == "__main__":
    predict(r'C:\Users\m2337\Desktop\p\22.png')