石油泄漏检测分割 无人机航拍图 石油泄漏分割数据集 YOLO unet分割任务

使用深度学习模型图像分割模型unet处理无人机航拍图中的石油泄漏检测分割。完成对无人机航拍图石油泄漏分割数据集 发的分割任务文章目录使用深度学习模型图像分割模型unet处理无人机航拍图中的石油泄漏检测分割。完成对无人机航拍图石油泄漏分割数据集 发的分割任务✅ 1. 数据集结构要求标注说明 2. 安装依赖库 3. 数据增强与加载器构建数据增强示例训练时使用:构建Dataloader 4. 模型选择与构建使用 Segmentation Models PyTorch⚙️ 5. 损失函数与优化器️ 6. 训练循环 7. 验证与可视化可视化预测结果 8. 推理部署单张图片测试 9. 性能评估指标可选 10. 模型导出ONNX如何处理呢无人机航拍图石油泄漏分割数据集数据集——1268张RGB图像1920*1080像素分为油、水和其他包括船只、码头、建筑物等三个类别。数据采集无人机Dronematrix YACOB和DJI Mavic2配备高分辨率相机在2021年9月至2023年9月期间拍摄涵盖了不同的环境条件、光照和天气情况。使用深度学习模型图像分割模型处理无人机航拍图中的石油泄漏检测问题。语义分割任务目标是将图像中每个像素分类为油、水或其他三类。—使用U-Net 模型或其改进版本如 U-Net、Attention U-Net 等来进行多类别图像语义分割。以下是完整流程✅ 1. 数据集结构要求确保你的数据集组织如下dataset/ ├── images/ │ ├── train/ │ │ ├── img001.png │ │ └── ... │ ├── val/ │ └── test/ └── masks/ ├── train/ │ ├── img001_mask.png │ └── ... ├── val/ └── test/标注说明mask 图像是单通道灰度图uint8像素值对应类别0 → 水water1 → 油oil2 → 其他other 2. 安装依赖库pipinstalltorch torchvision scikit-learn matplotlib albumentations segmentation_models_pytorch 3. 数据增强与加载器构建importosfromtorch.utils.dataimportDataset,DataLoaderfromPILimportImageimportnumpyasnpimportalbumentationsasAfromalbumentations.pytorchimportToTensorV2# 自定义Dataset类classOilSpillDataset(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_nameself.images[idx]img_pathos.path.join(self.image_dir,img_name)mask_pathos.path.join(self.mask_dir,img_name.replace(.jpg,_mask.png))imagenp.array(Image.open(img_path).convert(RGB))masknp.array(Image.open(mask_path),dtypenp.uint8)ifself.transform:augmentationsself.transform(imageimage,maskmask)imageaugmentations[image]maskaugmentations[mask]returnimage,mask.long()数据增强示例训练时使用:transform_trainA.Compose([A.Resize(height512,width512),A.HorizontalFlip(p0.5),A.VerticalFlip(p0.5),A.RandomRotate90(p0.5),A.ColorJitter(p0.3),A.Normalize(mean(0.485,0.456,0.406),std(0.229,0.224,0.225)),ToTensorV2(),])构建Dataloadertrain_datasetOilSpillDataset(image_dirdataset/images/train,mask_dirdataset/masks/train,transformtransform_train)train_loaderDataLoader(train_dataset,batch_size8,shuffleTrue,num_workers4) 4. 模型选择与构建使用 Segmentation Models PyTorch我们使用U-Net ResNet34 编码器importsegmentation_models_pytorchassmpimporttorch# 使用预训练的UNet模型modelsmp.Unet(encoder_nameresnet34,encoder_weightsimagenet,in_channels3,classes3# 3个类别油、水、其他)devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)model.to(device)⚙️ 5. 损失函数与优化器criteriontorch.nn.CrossEntropyLoss()optimizertorch.optim.Adam(model.parameters(),lr1e-4)️ 6. 训练循环fromtqdmimporttqdm num_epochs20forepochinrange(num_epochs):model.train()running_loss0.0forimages,masksintqdm(train_loader):imagesimages.to(device)masksmasks.to(device)outputsmodel(images)losscriterion(outputs,masks)optimizer.zero_grad()loss.backward()optimizer.step()running_lossloss.item()*images.size(0)epoch_lossrunning_loss/len(train_loader.dataset)print(fEpoch [{epoch1}/{num_epochs}], Loss:{epoch_loss:.4f}) 7. 验证与可视化验证过程类似训练只需将model.eval()并关闭梯度计算。可视化预测结果importmatplotlib.pyplotaspltdefvisualize_prediction(image,mask,pred):plt.figure(figsize(10,5))plt.subplot(1,3,1)plt.title(Image)plt.imshow(image.permute(1,2,0).cpu().numpy())plt.subplot(1,3,2)plt.title(True Mask)plt.imshow(mask.cpu().numpy(),cmapgray)plt.subplot(1,3,3)plt.title(Predicted Mask)plt.imshow(pred.argmax(0).cpu().numpy(),cmapgray)plt.show() 8. 推理部署单张图片测试fromtorchvisionimporttransformsdefpredict_image(model,image_path,device):model.eval()imageImage.open(image_path).convert(RGB)transformtransforms.Compose([transforms.Resize((512,512)),transforms.ToTensor(),transforms.Normalize(mean[0.485,0.456,0.406],std[0.229,0.224,0.225]),])image_tensortransform(image).unsqueeze(0).to(device)withtorch.no_grad():outputmodel(image_tensor)pred_masktorch.softmax(output,dim1).argmax(dim1).squeeze(0).cpu().numpy()returnimage,pred_mask 9. 性能评估指标可选可以使用以下指标进行评估IoU (Intersection over Union)Dice CoefficientPixel Accuracyfromsklearn.metricsimportjaccard_score,confusion_matrixdefcalculate_metrics(y_true,y_pred):ioujaccard_score(y_true.flatten(),y_pred.flatten(),averagemacro)cmconfusion_matrix(y_true.flatten(),y_pred.flatten())returniou,cm 10. 模型导出ONNXdummy_inputtorch.randn(1,3,512,512).to(device)torch.onnx.export(model,dummy_input,oil_spill_unet.onnx,export_paramsTrue)以上代码仅供参考