1、源码下载并配置环境GitHub - milesial/Pytorch-UNet: PyTorch implementation of the U-Net for image semantic segmentation with high quality images下载源码后并解压导D:\DINO\Pytorch-UNet-master路径下打开requirements.txt通过输入指令 conda list查看哪个模块没安装完成安装即可。2、准备自定义数据集2.1数据标注与目录结构标注工具可以使用Labelme工具对裂缝区域进行标注它会生成JSON文件。格式转换你需要将JSON文件转换为二值掩码图片例如裂缝区域像素值为255背景为0。目录组织将处理好的图片按以下结构存放这是最通用的方式。dataset/├── train/│ ├── images/ # 存放训练原图 (如 img_001.jpg)│ └── masks/ # 存放对应的标签图 (如 img_001.png建议为png)└── val/├── images/ # 存放验证原图└── masks/ # 存放对应的标签图如何将json文件转换为二值掩码图像1首先安装依赖 pip install labelme # labelme用于解析JSON2编写转换脚本json2mask.pyimport json import numpy as np import cv2 import os from labelme import utils def json_to_mask(json_path, output_mask_path): # 读取JSON文件 with open(json_path, r) as f: data json.load(f) # 获取图像尺寸 image_height data[imageHeight] image_width data[imageWidth] # 创建一个全黑的掩码单通道背景为0 mask np.zeros((image_height, image_width), dtypenp.uint8) # 遍历所有的标注形状假设你的裂缝标注是多边形 for shape in data[shapes]: label shape[label] # 标注名称例如 crack points shape[points] # 多边形顶点列表 # 将点坐标转换为整数OpenCV需要 points np.array(points, dtypenp.int32) # 在掩码上填充多边形裂缝区域设为255白色 cv2.fillPoly(mask, [points], color255) # 保存掩码为PNG格式无损压缩 cv2.imwrite(output_mask_path, mask) print(f掩码已保存至{output_mask_path}) # 批量处理示例 if __name__ __main__: json_folder path/to/your/jsons # 存放所有JSON的文件夹 output_mask_folder path/to/masks # 输出掩码的文件夹 os.makedirs(output_mask_folder, exist_okTrue) for json_file in os.listdir(json_folder): if json_file.endswith(.json): json_path os.path.join(json_folder, json_file) # 掩码文件名与原JSON同名但扩展名为.png mask_name json_file.replace(.json, .png) mask_path os.path.join(output_mask_folder, mask_name) json_to_mask(json_path, mask_path)3运行脚本python json2mask.py3.编写数据加载代码 (dataset.py)创建一个dataset.py文件定义如何将我们的数据读入模型。这里可以参考一些成熟项目的做法# dataset.py import os import torch from torch.utils.data import Dataset from PIL import Image import numpy as np class CrackDataset(Dataset): def __init__(self, image_dir, mask_dir, transformNone): self.image_dir image_dir self.mask_dir mask_dir self.transform transform self.images [f for f in os.listdir(image_dir) if f.endswith((.jpg, .png))] def __len__(self): return len(self.images) def __getitem__(self, idx): # 读取图片 img_name self.images[idx] img_path os.path.join(self.image_dir, img_name) image Image.open(img_path).convert(RGB) # 读取对应的掩码 (假设掩码是png格式且文件名与图片相同) mask_name os.path.splitext(img_name)[0] .png mask_path os.path.join(self.mask_dir, mask_name) mask Image.open(mask_path).convert(L) # 以灰度图读取 # 将掩码二值化裂缝区域为1背景为0 mask np.array(mask) mask (mask 127).astype(np.float32) # 如果有数据增强可以在这里应用 if self.transform: image self.transform(image) return image, torch.from_numpy(mask).unsqueeze(0) # 添加通道维度形状为 (1, H, W)4.模型训练将标注好的数据文件放置在D:\DINO\Pytorch-UNet-master\data文件夹内其中D:\DINO\Pytorch-UNet-master\data\train_imgs文件夹内放置训练的图像。D:\DINO\Pytorch-UNet-master\data\train_masks文件夹内放置的为掩码图像。如下图所示用VScode点开train.py将dir_img改为原图的绝对路径将dir_mask改为掩膜的绝对路径。将修改为虚拟环境激活及训练指令如下图所示训练过程时间有点长但值得等待出现以下情况选择3离线训练即可。 进入离线模式所有日志保存在本地wandb/文件夹不会上传。无需账户立即继续运行。5、预测模型训练好之后会将参数保存在checkpoints文件夹下因为当初设置的epochs等于5所以有5个结果。输入指令如下python .\predict.py -m checkpoints\checkpoint_epoch5.pth -i C:\Users\Administrator\Pictures\car1.jpeg测试结果如下6、模型转换为onnx新建文件pth2onnx并编写代码如下import torch import onnx import onnxruntime import numpy as np from unet import UNet # 你的模型定义 def export(): # 1. 加载模型 #unet UNet(in_channels3, num_classes1) unet UNet(n_channels3, n_classes2, bilinearFalse) state_dict torch.load(checkpoints/checkpoint_epoch5.pth, map_locationcuda) mask_values state_dict.pop(mask_values, [0, 1]) unet.load_state_dict(state_dict) unet.eval() # 2. 虚拟输入 dummy_input torch.randn(1, 3, 512, 512) # 3. 导出 ONNX torch.onnx.export( unet, dummy_input, unet.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}}, opset_version11, do_constant_foldingTrue ) print(导出完成) # 4. 验证 onnx_model onnx.load(unet.onnx) onnx.checker.check_model(onnx_model) print(ONNX 检查通过) # 5. 测试推理 ort_session onnxruntime.InferenceSession(unet.onnx) test_input np.random.randn(1, 3, 512, 512).astype(np.float32) outputs ort_session.run(None, {input: test_input}) print(ONNX Runtime 推理成功输出形状, outputs[0].shape) if __name__ __main__: export()运行指令转换后的结果如下可见onnx的输出格式为[batch_size, num_classes, height, width]所以batch_size1一次处理一张图像。num_classes2模型将像素分为2类这里指前景和背景两个类别。height512, width512输出特征图的空间尺寸与输入图像尺寸相同通常UNet保持尺寸不变通过上采样恢复。其输入输出用netron工具查看输入输出如下图所示通过C推理可完成其关注的mask提取。