如何在PyTorch DataLoader中为输入图像拼接单波段图像得到4波段输入

核心实现逻辑为：在数据集初始化阶段提前加载固定的待拼接单波段图像，在每次读取样本时，将3波段输入图像与预加载的单波段图像在通道维度拼接，最终输出4波段输入张量。

修改后完整代码 #

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

class lakeDataSet(Dataset):
    def __init__(self, root, transform, x_path):
        super().__init__()
        self.root = root
        self.img_dir = os.path.join(root,'image-c3/c3-crop')   #9UAV
        self.mask_dir = os.path.join(root,'label-c3/c3-crop')
        self.files = [fname for fname in os.listdir(self.img_dir) if fname.endswith('.jpg')]
        self.transform = transform
        # 提前加载待拼接的x.jpg，转单通道灰度图后应用transform处理
        self.x_band = self.transform(Image.open(x_path).convert('L'))

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

    def __getitem__(self, i):
        fname = self.files[i]
        img_path = os.path.join(self.img_dir, fname)
        mask_path = os.path.join(self.mask_dir, fname)

        # 读取原3波段图像，transform后shape为 (3, 高, 宽)
        img = self.transform(Image.open(img_path))
        mask = self.transform(Image.open(mask_path))
        # 通道维度拼接，得到4波段输入张量
        img_4band = torch.cat([img, self.x_band], dim=0)
        return img_4band, mask

实例化数据集时传入x.jpg的本地路径即可：

dataset = lakeDataSet(root="你的数据集根目录", transform=你的预处理逻辑, x_path="x.jpg的本地路径")

注意事项 #

• 若你使用的transform包含随机空间增强（如随机裁剪、随机翻转、随机旋转），不能提前预计算self.x_band，需要在__getitem__方法中对输入图像和x.jpg原图应用完全相同的增强参数后再拼接，避免空间位置不匹配
• 需提前确认x.jpg的分辨率与数据集中所有输入图像的分辨率完全一致，否则会触发拼接维度不匹配的报错
• 若x.jpg本身就是单波段图像，可删除convert('L')操作

内容的提问来源于stack exchange，提问作者bibinwilson