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

使用深度学习模型（图像分割模型）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. 安装依赖库

pip install torch torchvision scikit-learn matplotlib albumentations segmentation_models_pytorch

---

🧱 3. 数据增强与加载器构建

import os
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import numpy as np
import albumentations as A
from albumentations.pytorch import ToTensorV2

# 自定义Dataset类
class OilSpillDataset(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):
        return len(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)

        if self.transform:
            augmentations = self.transform(image=image, mask=mask)
            image = augmentations["image"]
            mask = augmentations["mask"]

        return image, 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 编码器

import segmentation_models_pytorch as smp
import torch

# 使用预训练的UNet模型
model = smp.Unet(
    encoder_name="resnet34",
    encoder_weights="imagenet",
    in_channels=3,
    classes=3  # 3个类别：油、水、其他
)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

---

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

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

---

🏋️ 6. 训练循环

from tqdm import tqdm

num_epochs = 20

for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0

    for images, masks in tqdm(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() 并关闭梯度计算。

可视化预测结果：

import matplotlib.pyplot as plt

def visualize_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. 推理部署（单张图片测试）

from torchvision import transforms

def predict_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)

    with torch.no_grad():
        output = model(image_tensor)
        pred_mask = torch.softmax(output, dim=1).argmax(dim=1).squeeze(0).cpu().numpy()

    return image, pred_mask

---

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

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

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

from sklearn.metrics import jaccard_score, confusion_matrix

def calculate_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())
    return iou, 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)

---

以上代码仅供参考，