U-Net模型进行训练钢材表面缺陷语义分割数据集 通过钢材缺陷分割数据集的权重模型推理识别钢材分割文章目录环境搭建1. 安装CUDA驱动2. 安装Anaconda3. 创建Python虚拟环境4. 安装依赖项数据集准备使用U-Net训练模型设备配置训练循环开始训练以下文字及代码仅供参考学习使用。钢材表面缺陷语义分割数据集4432张数据jpg和mask掩码png有颜色映射关系另外有转换成coco格式json和yolo格式txt三种缺陷类型像素标签0为背景 1为夹杂物In 2为补丁Pa 3为划痕Sc共4432张数据jpg和mask掩码png有颜色映射关系另外有转换成coco格式json和yolo格式txt含三种缺陷类型像素标签0为背景 1为夹杂物In 2为补丁Pa 3为划痕Scmask标签颜色映射为了使用U-Net模型对钢材表面缺陷进行语义分割我们需要从环境搭建开始到数据集准备、模型训练和推理。以下是详细的步骤指南。仅供参考学习使用labelme查看环境搭建1. 安装CUDA驱动确保您的系统已经安装了与GPU兼容的CUDA驱动版本。可以使用以下命令检查nvidia-smi2. 安装Anaconda访问 Anaconda官网 下载并安装适合您操作系统的版本。3. 创建Python虚拟环境打开终端或Anaconda Prompt然后输入以下命令来创建并激活新的Python环境conda create--nameunet_envpython3.9conda activate unet_env4. 安装依赖项在激活的环境中运行以下命令以安装必要的库pipinstalltorch torchvision torchaudio pipinstallopencv-python pipinstallmatplotlib pipinstallscikit-image pipinstallalbumentations pipinstalltqdm pipinstalltimm pipinstallsegmentation-models-pytorch数据集准备假设同学你的数据集按照如下结构组织steel_defect_dataset/ ├── images/ │ ├── train/ │ ├── val/ │ └── test/ ├── masks/ │ ├── train/ │ ├── val/ │ └── test/ └── data.yamldata.yaml文件内容示例请根据实际情况调整路径train_images:./steel_defect_dataset/images/traintrain_masks:./steel_defect_dataset/masks/trainval_images:./steel_defect_dataset/images/valval_masks:./steel_defect_dataset/masks/valtest_images:./steel_defect_dataset/images/testtest_masks:./steel_defect_dataset/masks/testnc:3names:[In,Pa,Sc]使用U-Net训练模型使用segmentation_models.pytorch库中的U-Net模型进行训练。的Python脚本示例用于加载U-Net模型并使用提供的数据集进行训练。仅供参考学习使用。optimizer torch.optim.Adam(model.parameters(), lr1e-4)设备配置device torch.device(“cuda” if torch.cuda.is_available() else “cpu”)model.to(device)训练循环def train_model(model, train_loader, val_loader, loss_fn, optimizer, num_epochs10):for epoch in range(num_epochs):model.train()loop tqdm(train_loader)for batch_idx, (data, targets) in enumerate(loop):data data.to(devicedevice)targets targets.long().to(devicedevice)# 前向传播 predictions model(data) loss loss_fn(predictions, targets) # 反向传播和优化 optimizer.zero_grad() loss.backward() optimizer.step() # 更新进度条 loop.set_postfix(lossloss.item()) # 验证阶段 model.eval() with torch.no_grad(): num_correct 0 num_pixels 0 dice_score 0 for data, targets in val_loader: data data.to(device) targets targets.to(device).unsqueeze(1) predictions torch.softmax(model(data), dim1) preds torch.argmax(predictions, dim1).float() num_correct (preds targets).sum() num_pixels torch.numel(preds) dice_score (2 * (preds * targets).sum()) / ((preds targets).sum() 1e-8) print(fGot {num_correct}/{num_pixels} with acc {num_correct/num_pixels*100:.2f}) print(fDice score: {dice_score/len(val_loader)})开始训练train_model(model, train_loader, val_loader, loss_fn, optimizer, num_epochs10)### 推理代码 训练完成后您可以使用训练好的模型对新图片进行预测。以下是一个简单的例子 python import cv2 from torchvision import transforms # 加载训练好的模型 model Unet(encoder_nameresnet34, classes3, activationNone) model.load_state_dict(torch.load(path/to/best_model.pth)) model.eval() # 图像预处理 preprocess transforms.Compose([ transforms.ToPILImage(), transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize(mean[0.0, 0.0, 0.0], std[1.0, 1.0, 1.0]), ]) # 对单张图片进行预测 image_path path/to/new/image.jpg img cv2.imread(image_path) img_tensor preprocess(img).unsqueeze(0).to(device) with torch.no_grad(): output model(img_tensor) prediction torch.argmax(output.squeeze(), dim0).cpu().numpy() # 显示结果 def label_to_color_image(label): colormap np.array([[0, 0, 0], [255, 0, 0], [0, 255, 0], [0, 0, 255]]) return colormap[label] color_prediction label_to_color_image(prediction) cv2.imshow(Prediction, color_prediction) cv2.waitKey(0) cv2.destroyAllWindows()
