PyTorch 实现 深度学习语义分割模型unet训练 道路裂缝分割数据集实现深度学习道路裂缝分割检测任务 语义分割。文章目录 数据集概述 下载 U-Net 示例代码PyTorch 实现1. 自定义 Dataset 类2. 构建 U-Net 模型简化版3. 训练脚本4. 推理 可视化 模型评估指标语义分割 模型导出ONNX / TorchScript道路的裂缝数据集适合做深度学习语义分割。标签掩膜图像。250张1你提供的是一个用于语义分割的小型道路裂缝数据集具有以下特点 数据集概述图像数量: 250张图像类型: RGB 图像如.jpg标签格式: PNG 格式的二值 mask 掩膜单通道表示裂缝区域像素值为1和非裂缝区域像素值为0任务类型:语义分割用途: 道路裂缝检测、自动驾驶辅助系统、路面维护监控等— 下载 U-Net 示例代码PyTorch 实现1. 自定义 Dataset 类importosfromPILimportImageimportnumpyasnpimporttorchfromtorch.utils.dataimportDatasetclassCrackDataset(Dataset):def__init__(self,image_dir,mask_dir,transformNone):self.image_dirimage_dir self.mask_dirmask_dir self.transformtransform self.imagesos.listdir(image_dir)def__len__(self):returnlen(self.images)def__getitem__(self,idx):img_pathos.path.join(self.image_dir,self.images[idx])mask_pathos.path.join(self.mask_dir,self.images[idx])imagenp.array(Image.open(img_path).convert(RGB))masknp.array(Image.open(mask_path).convert(L),dtypenp.float32)mask[mask255.0]1.0# 二值mask转为0和1ifself.transform:augmentationsself.transform(imageimage,maskmask)imageaugmentations[image]maskaugmentations[mask]returnimage,mask.unsqueeze(0)# shape: [C, H, W]2. 构建 U-Net 模型简化版importtorchimporttorch.nnasnndefdouble_conv(in_channels,out_channels):returnnn.Sequential(nn.Conv2d(in_channels,out_channels,3,padding1),nn.BatchNorm2d(out_channels),nn.ReLU(inplaceTrue),nn.Conv2d(out_channels,out_channels,3,padding1),nn.BatchNorm2d(out_channels),nn.ReLU(inplaceTrue))classUNet(nn.Module):def__init__(self,in_channels3):super(UNet,self).__init__()self.down1double_conv(in_channels,64)self.down2double_conv(64,128)self.down3double_conv(128,256)self.down4double_conv(256,512)self.maxpoolnn.MaxPool2d(2)self.up3nn.ConvTranspose2d(512,256,2,stride2)self.up_conv3double_conv(512,256)self.up2nn.ConvTranspose2d(256,128,2,stride2)self.up_conv2double_conv(256,128)self.up1nn.ConvTranspose2d(128,64,2,stride2)self.up_conv1double_conv(128,64)self.final_convnn.Conv2d(64,1,kernel_size1)defforward(self,x):# 下采样路径conv1self.down1(x)xself.maxpool(conv1)conv2self.down2(x)xself.maxpool(conv2)conv3self.down3(x)xself.maxpool(conv3)xself.down4(x)# 上采样路径xself.up3(x)xtorch.cat([x,conv3],dim1)xself.up_conv3(x)xself.up2(x)xtorch.cat([x,conv2],dim1)xself.up_conv2(x)xself.up1(x)xtorch.cat([x,conv1],dim1)xself.up_conv1(x)returntorch.sigmoid(self.final_conv(x))3. 训练脚本fromtorch.utils.dataimportDataLoaderfromtorchvisionimporttransformsimportalbumentationsasAfromalbumentations.pytorchimportToTensorV2# 数据增强与加载transformA.Compose([A.Resize(height256,width256),A.ToFloat(max_value255),A.Normalize(mean(0.485,0.456,0.406),std(0.229,0.224,0.225)),ToTensorV2()])train_datasetCrackDataset(images/train,masks/train,transformtransform)val_datasetCrackDataset(images/val,masks/val,transformtransform)train_loaderDataLoader(train_dataset,batch_size4,shuffleTrue)val_loaderDataLoader(val_dataset,batch_size4)# 初始化模型、损失函数、优化器devicecudaiftorch.cuda.is_available()elsecpumodelUNet().to(device)criterionnn.BCELoss()optimizertorch.optim.Adam(model.parameters(),lr1e-4)# 训练循环forepochinrange(50):# 训练50轮model.train()forimages,masksintrain_loader:imagesimages.to(device)masksmasks.to(device)outputsmodel(images)losscriterion(outputs,masks)optimizer.zero_grad()loss.backward()optimizer.step()print(fEpoch [{epoch1}/50], Loss:{loss.item():.4f})# 验证阶段略4. 推理 可视化importmatplotlib.pyplotasplt model.eval()withtorch.no_grad():images,masksnext(iter(val_loader))outputsmodel(images.to(device)).cpu().numpy()foriinrange(len(outputs)):plt.figure(figsize(10,5))plt.subplot(1,3,1)plt.title(Image)plt.imshow(images[i].permute(1,2,0).cpu().numpy())plt.subplot(1,3,2)plt.title(Mask)plt.imshow(masks[i][0].cpu().numpy(),cmapgray)plt.subplot(1,3,3)plt.title(Prediction)plt.imshow(outputs[i][0]0.5,cmapgray)plt.show() 模型评估指标语义分割fromsklearn.metricsimportjaccard_score,f1_scoredefcalculate_metrics(preds,targets):predspreds.flatten()targetstargets.flatten()ioujaccard_score(targets,preds)dicef1_score(targets,preds)return{IoU:iou,Dice:dice} 模型导出ONNX / TorchScriptdummy_inputtorch.randn(1,3,256,256).to(device)torch.onnx.export(model,dummy_input,unet_crack.onnx,export_paramsTrue)以上文字及代码仅供参考学习。