编下pytorch自定义数据集遥感分割代码

以下是一个简单的遥感图像分割数据集的 PyTorch 实现示例：

import torch
from torch.utils.data import Dataset
from PIL import Image
import os

class RemoteSensingDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.image_filenames = [os.path.join(root_dir, f) for f in os.listdir(root_dir) if f.endswith('.jpg')]
        self.mask_filenames = [f.replace('.jpg', '_mask.png') for f in self.image_filenames]

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

    def __getitem__(self, idx):
        image_filename = self.image_filenames[idx]
        mask_filename = self.mask_filenames[idx]

        image = Image.open(image_filename).convert('RGB')
        mask = Image.open(mask_filename).convert('L')

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

        return image, torch.tensor(mask, dtype=torch.long)

这个遥感图像分割数据集的实现假设数据集包含一系列 JPEG 图像和对应的 PNG 标注图像（命名规则为 xxx.jpg 和 xxx_mask.png，其中 xxx 是相同的文件名前缀）。在 __init__ 方法中，我们列出了所有 JPEG 图像的文件名，然后通过将 .jpg 扩展名替换为 _mask.png，得到了所有标注图像的文件名。在 __getitem__ 方法中，我们读取了 JPEG 和 PNG 图像，将 JPEG 图像转换为 RGB 格式，然后使用 torch.tensor 将 PNG 图像转换为张量。标注图像是单通道图像，所以我们使用 dtype=torch.long 指定了张量数据类型。最后，我们将图像和标注图像作为元组返回。

使用这个自定义数据集的示例代码如下：

from torchvision.transforms import Compose, ToTensor, Normalize
from torch.utils.data import DataLoader

transform = Compose([ToTensor(), Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])
dataset = RemoteSensingDataset('path/to/dataset', transform=transform)
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)

for images, masks in dataloader:
    # 模型训练代码

这个示例使用 PyTorch 的 DataLoader 类加载了自定义数据集，并使用了一个简单的图像转换器 ToTensor() 和数据标准化器 Normalize()。然后我们可以使用 DataLoader 迭代数据集中的数据，并将图像和标注图像作为输入传递给模型训练代码。