使用unet模型推理图像，生成可视化结果和LabelMe格式的标注文件

predict_oilDepth_num19.py

点击查看代码

import time
import cv2
import json
import shutil
import numpy as np
from PIL import Image

# 导入自定义的 U-Net 模型类
from unet import Unet

if __name__ == "__main__":
    # ================= 配置路径与参数 =================
    # 指定训练好的模型权重文件路径 (.pth)
    model_path = r"E:\PycharmWorkSpace\duankou\model_data\3000_066_meiwei_heightDifference_num19_2448_1920_0406.pth"

    # 指定待预测的原始图片文件夹路径
    dir_origin_path = r"E:\PycharmWorkSpace\duankou\img\oilDepth_num19\src"

    # 指定预测结果图片的保存路径（在原路径基础上增加后缀）
    dir_save_path = dir_origin_path + "_predict_segcls8_260420_vgg"
    # 指定生成的标注文件（JSON）的保存路径
    dir_anno_path = dir_save_path + "_anno"

    # 定义分割的类别数量（包括背景）
    num_classed = 3
    # 定义模型输入图像的尺寸 [高, 宽]
    input_shape = [2448, 1920]
    # 定义类别名称列表，索引对应类别ID
    label_classes = ["_background_", "tong", "youmo"]

    # 设置融合类型（具体含义取决于 unet.py 中的实现，通常指预测时的融合策略）
    mix_type = 1

    # ================= 初始化模型 =================
    # 实例化 U-Net 模型，加载权重和配置
    unet = Unet(model_path=model_path,
                num_classes=num_classed,
                label_classes=label_classes,
                input_shape=input_shape,
                mix_type=mix_type,
                backbone="vgg")  # 指定主干网络为 VGG，也可以改为 "resnet50"
    # backbone="resnet50")

    # ================= 批量预测处理 =================
    import os
    from tqdm import tqdm  # 导入进度条工具

    # 获取原始图片文件夹下的所有文件名
    img_names = os.listdir(dir_origin_path)

    # 使用 tqdm 显示处理进度
    for img_name in tqdm(img_names):
        # 判断文件是否为常见的图片格式
        if img_name.lower().endswith(
                ('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
            # 拼接完整的图片路径
            image_path = os.path.join(dir_origin_path, img_name)
            # 使用 PIL 打开图片
            image = Image.open(image_path)

            # 调用模型的预测函数，获取预测后的图片和掩码
            # return_mask=True 表示同时返回二值化或分类后的掩码图像
            r_image, r_mask = unet.detect_image_mask(image, return_mask=True)
            cv2.imwrite(r"D:\pic\see\Num19heightDifference\python_mask.png",r_mask);
            # 将掩码转换为 JSON 格式的对象（包含多边形点集等信息）
            obj = unet.change2json_from_mask(r_mask)

            # 将 JSON 中的 imageData 设为 None（通常用于减小文件大小，不保存 base64 编码的图像数据）
            obj["imageData"] = None
            # 在 JSON 对象中记录图片文件名
            obj["imagePath"] = img_name

            # 如果预测结果保存文件夹不存在，则创建
            if not os.path.exists(dir_save_path):
                os.makedirs(dir_save_path)

            # --- 保存拼接图 ---
            # 创建一个新的空白图片，宽度为原图的两倍，高度不变
            merged = Image.new('RGB', (r_image.width * 2, r_image.height))
            # 将原图粘贴到左侧 (0, 0)
            merged.paste(image, (0, 0))
            # 将预测结果图粘贴到右侧 (r_image.width, 0)
            merged.paste(r_image, (r_image.width, 0))
            # 保存拼接后的对比图
            merged.save(os.path.join(dir_save_path, img_name))

            # --- 保存标注文件 ---
            # 如果标注文件夹不存在，则创建
            if not os.path.exists(dir_anno_path):
                os.makedirs(dir_anno_path)

            # 定义源图片路径、目标图片路径和目标 JSON 路径
            src_img_path = os.path.join(dir_origin_path, img_name)
            dst_img_path = os.path.join(dir_anno_path, img_name)
            dst_log_path = os.path.join(dir_anno_path, img_name[:-4] + ".json")  # 将扩展名替换为 .json

            # 复制原图到标注文件夹（保持文件名一致）
            shutil.copy(src_img_path, dst_img_path)

            # 将包含分割信息的 JSON 对象写入文件
            with open(dst_log_path, 'w', encoding='utf-8') as f:
                json.dump(obj, f, ensure_ascii=False, indent=2)

unet.py

点击查看代码

# -*- coding: utf-8 -*-
import colorsys
import copy
import time
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn

# 导入自定义的 U-Net 网络结构
from nets.unet import Unet as unet
# 导入工具函数：颜色转换、预处理、resize、打印配置
#from utils.utils import cvtColor, preprocess_input, resize_image, show_config
from utils.utils import cvtColor, preprocess_input, resize_image, show_config

# --------------------------------------------#
#   Unet 类定义
#   负责模型的加载、预测和后处理
# --------------------------------------------#
class Unet(object):
    # 默认参数字典
    _defaults = {
        # -------------------------------------------------------------------#
        #   model_path: 指向 logs 文件夹下的权值文件 (.pth)
        #   训练好后会有多个权值文件，通常选择验证集损失较低的那个。
        # -------------------------------------------------------------------#
        "model_path": r"D:\work\unet-pytorch-main\runs\VOCdevkit_LayersThickness_epoch800cls8\best_epoch_weights.pth",
        # -------------------------------------------------------------------#
        #   num_classes: 需要区分的类别个数 + 1 (包含背景)
        #   必须与训练时一致，否则会出现 shape 不匹配错误
        # -------------------------------------------------------------------#
        "num_classes": 8,
        # 类别名称列表（用于生成 JSON 标注）
        "label_classes": [],
        # -------------------------------------------------------------------#
        #   backbone: 主干网络类型，可选 "vgg" 或 "resnet50"
        # -------------------------------------------------------------------#
        "backbone": "vgg",
        # -------------------------------------------------------------------#
        #   input_shape: 输入图片的大小 [高度, 宽度]
        #   必须与训练时一致
        # -------------------------------------------------------------------#
        "input_shape": [2048, 1920],
        # -------------------------------------------------------------------#
        #   mix_type: 检测结果的可视化方式
        #   0: 原图与预测图混合 (blend)
        #   1: 仅保留预测图 (纯分割结果)
        #   2: 扣去背景，仅保留原图中的目标
        # -------------------------------------------------------------------#
        "mix_type": 1,
        # -------------------------------------------------------------------#
        #   cuda: 是否使用 GPU 加速
        # -------------------------------------------------------------------#
        "cuda": False,
    }

    # ---------------------------------------------------#
    #   初始化函数
    # ---------------------------------------------------#
    def __init__(self, **kwargs):
        # 将默认参数加载到实例字典中
        self.__dict__.update(self._defaults)
        # 如果有传入新的参数，覆盖默认参数
        for name, value in kwargs.items():
            setattr(self, name, value)

        # 根据类别数量生成对应的颜色列表（用于可视化不同类别）
        if self.num_classes <= 6:
            # 类别少时使用预定义颜色
            self.colors = [(0, 0, 0), (255, 255, 255), (0, 255, 0), (0, 0, 255), (255, 0, 128), (128, 255, 128)]
        else:
            # 类别多时通过 HSV 颜色空间自动生成
            hsv_tuples = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)]
            self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
            self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors))

        # 加载模型
        self.generate()
        # 打印配置信息
        show_config(**self._defaults)

    # ---------------------------------------------------#
    #   获得模型 (加载权重)
    # ---------------------------------------------------#
    def generate(self, onnx=False):
        # 1. 实例化网络结构
        self.net = unet(num_classes=self.num_classes, backbone=self.backbone)

        # 2. 确定运行设备 (CPU 或 CUDA)
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        # 3. 加载权重文件
        self.net.load_state_dict(torch.load(self.model_path, map_location=device))

        # 4. 设置为评估模式 (关闭 Dropout 等)
        self.net = self.net.eval()
        print('{} model, and classes loaded.'.format(self.model_path))

        # 5. 如果有 GPU 且配置允许，使用 DataParallel 并移动到 GPU
        if not onnx:
            if self.cuda:
                self.net = nn.DataParallel(self.net)
                self.net = self.net.cuda()

    # ---------------------------------------------------#
    #   检测图片 (主推理函数)
    # ---------------------------------------------------#
    def detect_image(self, image, count=False, name_classes=None):
        # 1. 图像预处理：转换为 RGB
        image = cvtColor(image)
        # 备份原图用于后续混合显示
        old_img = copy.deepcopy(image)
        orininal_h = np.array(image).shape[0]
        orininal_w = np.array(image).shape[1]

        # 2. Resize 图像 (保持比例，加灰条)
        image_data, nw, nh = resize_image(image, (self.input_shape[1], self.input_shape[0]))

        # 3. 归一化、Transpose (HWC -> CHW)、增加 Batch 维度 (BCHW)
        image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, np.float32)), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(image_data)
            if self.cuda:
                images = images.cuda()

            # 4. 网络前向预测
            pr = self.net(images)[0]

            # 5. 后处理：Softmax 激活 -> 转为 NumPy
            pr = F.softmax(pr.permute(1, 2, 0), dim=-1).cpu().numpy()

            # 6. 去除灰条 (裁剪回原始比例区域)
            pr = pr[int((self.input_shape[0] - nh) // 2): int((self.input_shape[0] - nh) // 2 + nh), \
                int((self.input_shape[1] - nw) // 2): int((self.input_shape[1] - nw) // 2 + nw)]

            # 7. Resize 回原始图像大小
            pr = cv2.resize(pr, (orininal_w, orininal_h), interpolation=cv2.INTER_LINEAR)

            # 8. 获取每个像素点的类别 (取概率最大的索引)
            pr = pr.argmax(axis=-1)

        # 9. 统计各类别像素数量 (可选)
        if count:
            classes_nums = np.zeros([self.num_classes])
            total_points_num = orininal_h * orininal_w
            print('-' * 63)
            print("|%25s | %15s | %15s|" % ("Key", "Value", "Ratio"))
            print('-' * 63)
            for i in range(self.num_classes):
                num = np.sum(pr == i)
                ratio = num / total_points_num * 100
                if num > 0:
                    print("|%25s | %15s | %14.2f%%|" % (str(name_classes[i]), str(num), ratio))
                classes_nums[i] = num
            print("classes_nums:", classes_nums)

        # 10. 根据 mix_type 生成可视化结果
        if self.mix_type == 0:
            # 0: 原图与预测图混合
            seg_img = np.reshape(np.array(self.colors, np.uint8)[np.reshape(pr, [-1])], [orininal_h, orininal_w, -1])
            image = Image.fromarray(np.uint8(seg_img))
            image = Image.blend(old_img, image, 0.7)
        elif self.mix_type == 1:
            # 1: 仅显示预测图
            seg_img = np.reshape(np.array(self.colors, np.uint8)[np.reshape(pr, [-1])], [orininal_h, orininal_w, -1])
            image = Image.fromarray(np.uint8(seg_img))
        elif self.mix_type == 2:
            # 2: 扣除背景，保留原图目标
            seg_img = (np.expand_dims(pr != 0, -1) * np.array(old_img, np.float32)).astype('uint8')
            image = Image.fromarray(np.uint8(seg_img))

        return image

    # ---------------------------------------------------#
    #   检测图片并返回掩码 (用于生成 JSON)
    # ---------------------------------------------------#
    def detect_image_mask(self, image, return_mask=False):
        # 流程与 detect_image 基本一致，省略重复注释
        image = cvtColor(image)
        old_img = copy.deepcopy(image)
        orininal_h = np.array(image).shape[0]
        orininal_w = np.array(image).shape[1]

        image_data, nw, nh = resize_image(image, (self.input_shape[1], self.input_shape[0]))
        image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, np.float32)), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(image_data)
            if self.cuda:
                images = images.cuda()
            pr = self.net(images)[0]
            pr = F.softmax(pr.permute(1, 2, 0), dim=-1).cpu().numpy()
            pr = pr[int((self.input_shape[0] - nh) // 2): int((self.input_shape[0] - nh) // 2 + nh), \
                int((self.input_shape[1] - nw) // 2): int((self.input_shape[1] - nw) // 2 + nw)]
            pr = cv2.resize(pr, (orininal_w, orininal_h), interpolation=cv2.INTER_LINEAR)
            pr = pr.argmax(axis=-1)

        # 生成可视化图像
        if self.mix_type == 0:
            seg_img = np.reshape(np.array(self.colors, np.uint8)[np.reshape(pr, [-1])], [orininal_h, orininal_w, -1])
            image = Image.fromarray(np.uint8(seg_img))
            image = Image.blend(old_img, image, 0.3)
        elif self.mix_type == 1:
            seg_img = np.reshape(np.array(self.colors, np.uint8)[np.reshape(pr, [-1])], [orininal_h, orininal_w, -1])
            image = Image.fromarray(np.uint8(seg_img))
        elif self.mix_type == 2:
            seg_img = (np.expand_dims(pr != 0, -1) * np.array(old_img, np.float32)).astype('uint8')
            image = Image.fromarray(np.uint8(seg_img))

        # 如果需要返回 mask 数组，则返回 (可视化图, mask数组)
        if return_mask:
            return image, pr.astype(np.uint8)
        return image

    # ---------------------------------------------------#
    #   将预测的 Mask 转换为 LabelMe 格式的 JSON 数据
    # ---------------------------------------------------#
    def change2json_from_mask(self, mask):
        # 初始化 JSON 数据结构
        data = {'shapes': [], 'imageHeight': mask.shape[0], 'imageWidth': mask.shape[1]}
        classes_count = self.num_classes

        # 遍历每一个类别 (从 1 开始，0 通常是背景)
        for idx in range(1, classes_count):
            try:
                label_class = self.label_classes[idx]
            except:
                label_class = idx

            # 1. 提取当前类别的二值 Mask
            binary_mask = (mask == idx).astype(np.uint8) * 255

            # 2. 使用 OpenCV 查找轮廓
            # RETR_EXTERNAL: 只检测外轮廓
            # CHAIN_APPROX_SIMPLE: 压缩水平、垂直和对角线方向的元素，只保留端点
            contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            # 3. 遍历所有轮廓
            for contour in contours:
                # 过滤掉极小的噪点 (面积小于 5 的忽略)
                if cv2.contourArea(contour) < 5:
                    continue

                # 提取点集
                points = []
                for point in contour:
                    # point[0] 是 [x, y] 坐标
                    points.append(point[0].tolist())

                # 只有点数大于 2 才能构成多边形
                if len(points) > 2:
                    shape = {
                        'label': label_class,
                        'points': points,
                        'shape_type': 'polygon'
                    }
                    data['shapes'].append(shape)
        return data

    # ---------------------------------------------------#
    #   获取 FPS (用于测试推理速度)
    # ---------------------------------------------------#
    def get_FPS(self, image, test_interval):
        image = cvtColor(image)
        image_data, nw, nh = resize_image(image, (self.input_shape[1], self.input_shape[0]))
        image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, np.float32)), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(image_data)
            if self.cuda:
                images = images.cuda()

            # 预热 (Warm up)，防止第一次推理包含初始化时间导致误差
            pr = self.net(images)[0]
            pr = F.softmax(pr.permute(1, 2, 0), dim=-1).cpu().numpy().argmax(axis=-1)
            pr = pr[int((self.input_shape[0] - nh) // 2): int((self.input_shape[0] - nh) // 2 + nh), \
                int((self.input_shape[1] - nw) // 2): int((self.input_shape[1] - nw) // 2 + nw)]

            # 正式计时
            t1 = time.time()
            for _ in range(test_interval):
                pr = self.net(images)[0]
                pr = F.softmax(pr.permute(1, 2, 0), dim=-1).cpu().numpy().argmax(axis=-1)
                pr = pr[int((self.input_shape[0] - nh) // 2): int((self.input_shape[0] - nh) // 2 + nh), \
                    int((self.input_shape[1] - nw) // 2): int((self.input_shape[1] - nw) // 2 + nw)]
            t2 = time.time()

            # 计算平均单张图片耗时
            tact_time = (t2 - t1) / test_interval
            return tact_time

    # ---------------------------------------------------#
    #   转换为 ONNX 格式 (用于模型部署)
    # ---------------------------------------------------#
    def convert_to_onnx(self, simplify, model_path):
        import onnx
        self.generate(onnx=True)

        # 创建一个虚拟输入张量
        im = torch.zeros(1, 3, *self.input_shape).to('cpu')
        input_layer_names = ["images"]
        output_layer_names = ["output"]

        print(f'Starting export with onnx {onnx.__version__}.')
        # 导出 ONNX 模型
        torch.onnx.export(self.net, im, f=model_path, verbose=False, opset_version=12,
                          training=torch.onnx.TrainingMode.EVAL, do_constant_folding=True,
                          input_names=input_layer_names, output_names=output_layer_names, dynamic_axes=None)

        # 检查模型合法性
        model_onnx = onnx.load(model_path)
        onnx.checker.check_model(model_onnx)

        # 简化模型 (可选)
        if simplify:
            import onnxsim
            print(f'Simplifying with onnx-simplifier {onnxsim.__version__}.')
            model_onnx, check = onnxsim.simplify(model_onnx, dynamic_input_shape=False, input_shapes=None)
            assert check, 'assert check failed'
            onnx.save(model_onnx, model_path)

        print('Onnx model save as {}'.format(model_path))

    # ---------------------------------------------------#
    #   获取 mIoU 所需的预测图 (灰度图，像素值代表类别)
    # ---------------------------------------------------#
    def get_miou_png(self, image):
        image = cvtColor(image)
        orininal_h = np.array(image).shape[0]
        orininal_w = np.array(image).shape[1]

        image_data, nw, nh = resize_image(image, (self.input_shape[1], self.input_shape[0]))
        image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, np.float32)), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(image_data)
            if self.cuda:
                images = images.cuda()
            pr = self.net(images)[0]
            pr = F.softmax(pr.permute(1, 2, 0), dim=-1).cpu().numpy()
            pr = pr[int((self.input_shape[0] - nh) // 2): int((self.input_shape[0] - nh) // 2 + nh), \
                int((self.input_shape[1] - nw) // 2): int((self.input_shape[1] - nw) // 2 + nw)]
            pr = cv2.resize(pr, (orininal_w, orininal_h), interpolation=cv2.INTER_LINEAR)
            pr = pr.argmax(axis=-1)

        # 返回单通道的灰度图，像素值即为类别 ID
        image = Image.fromarray(np.uint8(pr))
        return image