烟叶病害分割数据集
分为两个数据文件夹.

文件夹1：1500张图片，以及相应的mask
文件夹2：2000张图片，以及相应的mask

1

1

1

使用U-Net模型训练烟叶病害分割数据集是一个典型的图像分割任务。

环境设置

首先，确保你的环境中安装了必要的库：

1. Anaconda：如果你还没有安装Anaconda，请先下载并安装它。
2. 创建虚拟环境：

   conda create --name unet_env python=3.8
   conda activate unet_env
   

3. 安装依赖项：
   • PyTorch（根据你的CUDA版本选择合适的命令）:

     pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
     

   • 其他必要库:

     pip install numpy matplotlib scikit-image opencv-python
     

数据集组织

假设你的数据集组织如下：

dataset/
├── folder1/
│   ├── images/
│   │   ├── img_00001.jpg
│   │   └── ...
│   ├── masks/
│   │   ├── mask_00001.png
│   │   └── ...
├── folder2/
│   ├── images/
│   │   ├── img_00001.jpg
│   │   └── ...
│   ├── masks/
│   │   ├── mask_00001.png
│   │   └── ...

数据加载器

编写一个自定义的数据加载器来处理这个数据集：

import os
from torch.utils.data import Dataset, DataLoader
from skimage.io import imread
from torchvision.transforms import ToTensor

class LeafDiseaseDataset(Dataset):
    def __init__(self, image_folder, mask_folder, transform=None):
        self.image_paths = [os.path.join(image_folder, filename) for filename in os.listdir(image_folder)]
        self.mask_paths = [os.path.join(mask_folder, filename.replace('.jpg', '.png')) for filename in os.listdir(image_folder)]
        self.transform = transform

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        image_path = self.image_paths[idx]
        mask_path = self.mask_paths[idx]

        image = imread(image_path)
        mask = imread(mask_path, as_gray=True)

        if self.transform:
            image = self.transform(image)
            mask = ToTensor()(mask)

        return image, mask

# 使用示例
image_folder1 = 'path/to/folder1/images/'
mask_folder1 = 'path/to/folder1/masks/'

dataset1 = LeafDiseaseDataset(image_folder1, mask_folder1, transform=ToTensor())
dataloader1 = DataLoader(dataset1, batch_size=4, shuffle=True)

对于第二个文件夹，重复上述步骤创建dataset2和dataloader2。

U-Net模型定义

你可以直接使用或微调已有的U-Net实现。这里提供一个简单的U-Net架构：

import torch.nn as nn
import torch

class UNet(nn.Module):
    # 这里简化了U-Net的定义，实际应用中可能需要更复杂的网络结构
    def __init__(self):
        super(UNet, self).__init__()
        # 定义编码器和解码器部分
        # ...

    def forward(self, x):
        # 前向传播逻辑
        # ...
        return x

model = UNet()
if torch.cuda.is_available():
    model.cuda()

损失函数与优化器

选择适合的损失函数（如交叉熵损失）和优化器（如Adam）：

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

训练循环

最后，编写训练循环：

num_epochs = 20
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, masks in dataloader1:  # 对于第二个数据集，同样地遍历dataloader2
        if torch.cuda.is_available():
            images = images.cuda()
            masks = masks.cuda()

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, masks.long().squeeze(1))  # 注意调整维度以匹配损失函数要求
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * images.size(0)
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(dataloader1.dataset)}")

这只是一个基本框架，具体实现时可能需要根据实际情况调整。例如，处理数据增强、学习率调度等。此外，确保在开始训练之前将所有路径替换为你的实际数据集路径。