image-processing - PyTorch ToTensor 将 C x H x W (5 x 600 x 900) 更改为 H x C x W (900 x 5 x 600)

Question

这是我的数据加载器。当我使用 ToTensor 时，它会将图像的尺寸更改为 H x C x W。置换是否可以解决这个问题，或者这可能会改变一些方向？

class DPWHDataset(Dataset):
  def __init__(self, mean=None, std=None, phase=None, dataset=None):
    self.data = dataset
    self.mean = mean
    self.std = std
    self.phase = phase
    self.transforms = get_transforms(phase, mean, std)

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

  def __getitem__(self, idx):
    image_name = self.data[idx]

    image_path = image_prefix + image_name + ".jpg"
    mask_path = binary_mask_prefix + image_name + "_mask.png"

    mask = cv2.imread(mask_path, 0)
    print(image_path)

    # image = np.array(Image.open(image_path))
    # mask = np.array(Image.open(mask_path))
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    mask = create_channel_mask(mask)
    # augmented = self.transforms(image=image, mask=mask)
    # image = augmented['image']
    # mask = augmented['mask']

    image = torchvision.transforms.ToTensor()(image)
    image = torchvision.transforms.Normalize(mean=self.mean, std=self.std)(image)
    mask = torchvision.transforms.ToTensor()(mask)
    return image, mask

Answer 1

根据文档，torchvision.transforms.ToTensor将 PIL Image 或 numpy.ndarray (H x W x C) 转换为形状为 (C x H x W) 的 torch.FloatTensor。因此，在以下行中:

image = torchvision.transforms.ToTensor()(image)

生成的图像张量的形状为(C x H x W)，输入张量的形状为(H x W x C)。您可以通过打印张量形状来验证这一点。

是的，您可以使用 torch.permute 调整形状，不会造成任何问题。