pytorch/torchvision处理32位/16位灰度图的坑

Posted 2023-04-27

参考技术A 之前利用PIL把dicom的slice保存为了16位灰度图, 用 torchvision.transform 做图像增强时发现会报错.

Dataset 的 __getitem__ 函数如下

output:

查询了一下 torchvision.transform.ToTensor() 函数, 发现对输入值域要求为 [0-255] . 估计是我 [0-2048] 的范围出发了某种判断, 使得该函数以为输入图片为某种其他格式.

将函数改为如下, 解决了问题

注意, 若不显式注明 dtype='float32' , 会自动转换为 float64 的 tensor , 不确定对训练结果和速度有何影响 ( pytorch 的默认数据类型为 float32 ).

总之, 下次直接把图片保存为 numpy 格式会更方便些.

如何在 Pytorch 中使用 torchvision.transforms 对分割任务进行数据增强？

【中文标题】如何在 Pytorch 中使用 torchvision.transforms 对分割任务进行数据增强？

【英文标题】：How to use torchvision.transforms for data augmentation of segmentation task in Pytorch?

【发布时间】：2020-02-01 12:46:33

【问题描述】：

我对 PyTorch 中执行的数据增强有点困惑。

因为我们在处理分割任务，我们需要数据和掩码来进行相同的数据增强，但其中一些是随机的，例如随机旋转。

Keras 提供了random seed 保证 data 和 mask 做同样的操作，如下代码所示：

    data_gen_args = dict(featurewise_center=True,
                         featurewise_std_normalization=True,
                         rotation_range=25,
                         horizontal_flip=True,
                         vertical_flip=True)


    image_datagen = ImageDataGenerator(**data_gen_args)
    mask_datagen = ImageDataGenerator(**data_gen_args)

    seed = 1
    image_generator = image_datagen.flow(train_data, seed=seed, batch_size=1)
    mask_generator = mask_datagen.flow(train_label, seed=seed, batch_size=1)

    train_generator = zip(image_generator, mask_generator)

我在Pytorch官方文档中没有找到类似的描述，所以不知道如何保证数据和掩码可以同步处理。

Pytorch 确实提供了这样的功能，但我想将其应用于自定义 Dataloader。

例如：

def __getitem__(self, index):
    img = np.zeros((self.im_ht, self.im_wd, channel_size))
    mask = np.zeros((self.im_ht, self.im_wd, channel_size))

    temp_img = np.load(Image_path + ':0>4'.format(self.patient_index[index]) + '.npy')
    temp_label = np.load(Label_path + ':0>4'.format(self.patient_index[index]) + '.npy')

    for i in range(channel_size):
        img[:,:,i] = temp_img[self.count[index] + i]
        mask[:,:,i] = temp_label[self.count[index] + i]

    if self.transforms:
        img = np.uint8(img)
        mask = np.uint8(mask)
        img = self.transforms(img)
        mask = self.transforms(mask)

    return img, mask

这种情况下img和mask会分开变换，因为一些随机旋转等操作是随机的，所以mask和image的对应关系可能会改变。换句话说，图像可能已经旋转，但蒙版没有这样做。

编辑 1

我用augmentations.py中的方法，但是报错：

Traceback (most recent call last):
  File "test_transform.py", line 87, in <module>
    for batch_idx, image, mask in enumerate(train_loader):
  File "/home/dirk/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 314, in __next__
    batch = self.collate_fn([self.dataset[i] for i in indices])
  File "/home/dirk/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 314, in <listcomp>
    batch = self.collate_fn([self.dataset[i] for i in indices])
  File "/home/dirk/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/utils/data/dataset.py", line 103, in __getitem__
    return self.dataset[self.indices[idx]]
  File "/home/dirk/home/data/dirk/segmentation_unet_pytorch/data.py", line 164, in __getitem__
    img, mask = self.transforms(img, mask)
  File "/home/dirk/home/data/dirk/segmentation_unet_pytorch/augmentations.py", line 17, in __call__
    img, mask = a(img, mask)
TypeError: __call__() takes 2 positional arguments but 3 were given

这是我__getitem__()的代码：

data_transforms = 
    'train': Compose([
        RandomHorizontallyFlip(),
        RandomRotate(degree=25),
        transforms.ToTensor()
    ]),


train_set = DatasetUnetForTestTransform(fold=args.fold, random_index=args.random_index,transforms=data_transforms['train'])

# __getitem__ in class DatasetUnetForTestTransform
def __getitem__(self, index):
    img = np.zeros((self.im_ht, self.im_wd, channel_size))
    mask = np.zeros((self.im_ht, self.im_wd, channel_size))
    temp_img = np.load(Label_path + ':0>4'.format(self.patient_index[index]) + '.npy')
    temp_label = np.load(Label_path + ':0>4'.format(self.patient_index[index]) + '.npy')
    temp_img, temp_label = crop_data_label_from_0(temp_img, temp_label)
    for i in range(channel_size):
        img[:,:,i] = temp_img[self.count[index] + i]
        mask[:,:,i] = temp_label[self.count[index] + i]

    if self.transforms:
        img = T.ToPILImage()(np.uint8(img))
        mask = T.ToPILImage()(np.uint8(mask))
        img, mask = self.transforms(img, mask)

    img = T.ToTensor()(img).copy()
    mask = T.ToTensor()(mask).copy()
    return img, mask

编辑 2

我发现ToTensor之后，相同标签之间的骰子变成255而不是1，如何解决？

# Dice computation
def DSC_computation(label, pred):
    pred_sum = pred.sum()
    label_sum = label.sum()
    inter_sum = np.logical_and(pred, label).sum()
    return 2 * float(inter_sum) / (pred_sum + label_sum)

随时询问是否需要更多代码来解释问题。

【参考方案1】：

需要像RandomCrop 这样的输入参数的转换有一个get_param 方法，它将返回该特定转换的参数。然后可以使用转换的功能接口将其应用于图像和蒙版：

from torchvision import transforms
import torchvision.transforms.functional as F

i, j, h, w = transforms.RandomCrop.get_params(input, (100, 100))
input = F.crop(input, i, j, h, w)
target = F.crop(target, i, j, h, w)

此处提供示例： https://github.com/pytorch/vision/releases/tag/v0.2.0

此处提供 VOC 和 COCO 的完整示例： https://github.com/pytorch/vision/blob/master/references/segmentation/transforms.py https://github.com/pytorch/vision/blob/master/references/segmentation/train.py

关于错误，

ToTensor() 未被覆盖以处理额外的掩码参数，因此它不能在data_transforms 中。此外，__getitem__ 在返回之前对img 和mask 执行ToTensor。

data_transforms = 
    'train': Compose([
        RandomHorizontallyFlip(),
        RandomRotate(degree=25),
        #transforms.ToTensor()  => remove this line
    ]),

【讨论】：

感谢您的回答，我使用该方法实现了transform，但是在问题描述中添加了一个错误。

用修复更新了答案。

【参考方案2】：

torchvision也提供了类似的功能[document]。

这是一个简单的例子，

import torchvision
from torchvision import transforms

trans = transforms.Compose([transforms.CenterCrop((178, 178)),
                                    transforms.Resize(128),
                                    transforms.RandomRotation(20),
                                    transforms.ToTensor(),
                                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
dset = torchvision.datasets.MNIST(data_root, transforms=trans)

---

编辑

自定义您自己的 CelebA 数据集时的一个简短示例。请注意，要应用转换，您需要在__getitem__ 中调用transform 列表。

class CelebADataset(Dataset):
    def __init__(self, root, transforms=None, num=None):
        super(CelebADataset, self).__init__()

        self.img_root = os.path.join(root, 'img_align_celeba')
        self.attr_root = os.path.join(root, 'Anno/list_attr_celeba.txt')
        self.transforms = transforms

        df = pd.read_csv(self.attr_root, sep='\s+', header=1, index_col=0)
        #print(df.columns.tolist())
        if num is None:
            self.labels = df.values
            self.img_name = df.index.values
        else:
            self.labels = df.values[:num]
            self.img_name = df.index.values[:num]

    def __getitem__(self, index):
        img = Image.open(os.path.join(self.img_root, self.img_name[index]))
        # only use blond_hair, eyeglass, male, smile
        indices = [9, 15, 20, 31]
        label = np.take(self.labels[index], indices)
        label[label==-1] = 0

        if self.transforms is not None:
            img = self.transforms(img)

        return np.asarray(img), label

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

---

编辑 2

乍一看，我可能会错过一些东西。您问题的重点是如何将“相同”的数据预处理应用于 img 和标签。据我了解，没有可用的 Pytorch 内置函数。所以，我之前所做的就是自己实现增强。

class RandomRotate(object):
    def __init__(self, degree):
        self.degree = degree

    def __call__(self, img, mask):
        rotate_degree = random.random() * 2 * self.degree - self.degree
        return img.rotate(rotate_degree, Image.BILINEAR), 
                           mask.rotate(rotate_degree, Image.NEAREST)

请注意，输入应为 PIL 格式。请参阅this 了解更多信息。

【讨论】：

我重新编辑了我的问题。您的示例用于分类任务。如果用于分割任务如何修改它？具体来说，我们需要同时旋转图片和标签，那么这里应该如何保证它们的对应呢？如果使用img, label = self.transforms(img), self.transforms(label)，变换中的随机旋转会破坏图像和标签之间的对应关系。换句话说，图像可能已经旋转，但蒙版没有这样做。

@Dirk Li，感谢您的评论。我已经添加了我之前的处理方式。

我使用了augmentations.py中的方法，但是出现了错误。并且我在问题描述中添加了错误描述。

我的错。请在augmentation.py 中使用Compose。我已经相应地更新了我的答案。

感谢您的回答，我现在知道问题的原因了。但是还有一个地方让我很困惑。我使用labe1, label2 = self.transforms(label,label) 来测试transform 是否改变了数据。我用label1 和label2 计算骰子，结果是255（以前是1）。我知道ToTensor会把数据除以255，但是我还是不明白Dice为什么变成255，怎么解决。顺便说一句，增强中的 random_rotation 会将 Dice 减少到 244。这是否意味着它改变了对应关系？将骰子计算添加到问题描述中。

【参考方案3】：

另一个想法是沿着通道维度堆叠图像和蒙版，然后将它们一起转换。显然，这只适用于几何类型的转换，您需要对两者使用相同的 dtype。我使用这样的东西：

# Apply these to image and mask
affine_transforms = transforms.Compose([
    transforms.RandomAffine(degrees=180),
    ...
])

# Apply these to image only
image_transforms = transforms.Compose([
    transforms.GaussianBlur(),
    ...
])

# Loader...
def __getitem__(self, index: int):
    # Get the image and mask, here shape=(HxW) for both
    image = self.images[index]
    mask = self.masks[index]

    # Stack the image and mask together so they get the same geometric transformations
    stacked = torch.cat([image, mask], dim=0)  # shape=(2xHxW)
    stacked = self.affine_transforms(stacked)

    # Split them back up again
    image, mask = torch.chunk(stacked, chunks=2, dim=0)

    # Image transforms are only applied to the image
    image = self.image_transforms(image)

    return image, mask