时间戳

首页 > 解决方案 > PyTorch ToTensor 将 C x H x W (5 x 600 x 900) 更改为 H x C x W (900 x 5 x 600)

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

问题描述

这是我的数据加载器。当我使用 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

标签： image-processingpytorchtorchvision

解决方案

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

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

结果image张量是 shape (C x H x W)，输入张量是 shape (H x W x C)。您可以通过打印张量形状来验证这一点。

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

推荐阅读