企业官网建设流程全解析

使用深度学习模型（图像分割模型）unet处理无人机航拍图中的石油泄漏检测分割。完成对无人机航拍图石油泄漏分割数据集 发的分割任务

如何处理呢

无人机航拍图石油泄漏分割数据集
数据集——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 → 水（water）
  • 1 → 油（oil）
  • 2 → 其他（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,transform=None):self.image_dir=image_dir self.mask_dir=mask_dir self.transform=transform self.images=os.listdir(image_dir)def__len__(self):returnlen(self.images)def__getitem__(self,idx):img_name=self.images[idx]img_path=os.path.join(self.image_dir,img_name)mask_path=os.path.join(self.mask_dir,img_name.replace(".jpg","_mask.png"))image=np.array(Image.open(img_path).convert("RGB"))mask=np.array(Image.open(mask_path),dtype=np.uint8)ifself.transform:augmentations=self.transform(image=image,mask=mask)image=augmentations["image"]mask=augmentations["mask"]returnimage,mask.long()

数据增强示例（训练时使用）:

transform_train=A.Compose([A.Resize(height=512,width=512),A.HorizontalFlip(p=0.5),A.VerticalFlip(p=0.5),A.RandomRotate90(p=0.5),A.ColorJitter(p=0.3),A.Normalize(mean=(0.485,0.456,0.406),std=(0.229,0.224,0.225)),ToTensorV2(),])

构建Dataloader：

train_dataset=OilSpillDataset(image_dir='dataset/images/train',mask_dir='dataset/masks/train',transform=transform_train)train_loader=DataLoader(train_dataset,batch_size=8,shuffle=True,num_workers=4)

---

🤖 4. 模型选择与构建（使用 Segmentation Models PyTorch）

我们使用U-Net + ResNet34 编码器

importsegmentation_models_pytorchassmpimporttorch# 使用预训练的UNet模型model=smp.Unet(encoder_name="resnet34",encoder_weights="imagenet",in_channels=3,classes=3# 3个类别：油、水、其他)device=torch.device('cuda'iftorch.cuda.is_available()else'cpu')model.to(device)

---

⚙️ 5. 损失函数与优化器

criterion=torch.nn.CrossEntropyLoss()optimizer=torch.optim.Adam(model.parameters(),lr=1e-4)

---

🏋️ 6. 训练循环

fromtqdmimporttqdm num_epochs=20forepochinrange(num_epochs):model.train()running_loss=0.0forimages,masksintqdm(train_loader):images=images.to(device)masks=masks.to(device)outputs=model(images)loss=criterion(outputs,masks)optimizer.zero_grad()loss.backward()optimizer.step()running_loss+=loss.item()*images.size(0)epoch_loss=running_loss/len(train_loader.dataset)print(f"Epoch [{epoch+1}/{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(),cmap='gray')plt.subplot(1,3,3)plt.title("Predicted Mask")plt.imshow(pred.argmax(0).cpu().numpy(),cmap='gray')plt.show()

---

📁 8. 推理部署（单张图片测试）

fromtorchvisionimporttransformsdefpredict_image(model,image_path,device):model.eval()image=Image.open(image_path).convert("RGB")transform=transforms.Compose([transforms.Resize((512,512)),transforms.ToTensor(),transforms.Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225]),])image_tensor=transform(image).unsqueeze(0).to(device)withtorch.no_grad():output=model(image_tensor)pred_mask=torch.softmax(output,dim=1).argmax(dim=1).squeeze(0).cpu().numpy()returnimage,pred_mask

---

📈 9. 性能评估指标（可选）

可以使用以下指标进行评估：

• IoU (Intersection over Union)
• Dice Coefficient
• Pixel Accuracy

fromsklearn.metricsimportjaccard_score,confusion_matrixdefcalculate_metrics(y_true,y_pred):iou=jaccard_score(y_true.flatten(),y_pred.flatten(),average='macro')cm=confusion_matrix(y_true.flatten(),y_pred.flatten())returniou,cm

---

📦 10. 模型导出（ONNX）

dummy_input=torch.randn(1,3,512,512).to(device)torch.onnx.export(model,dummy_input,"oil_spill_unet.onnx",export_params=True)

---

以上代码仅供参考，