Pytorch数据变换(Transform)

实例化数据库的时候,有一个可选的参数可以对数据进行转换,满足大多神经网络的要求输入固定尺寸的图片,因此要对原图进行Rescale或者Crop操作,然后返回的数据需要转换成Tensor如:

import FaceLandmarksDataset
face_dataset = FaceLandmarksDataset(csv_file='data/faces/face_landmarks.csv',
                                    root_dir='data/faces/',
                                    transform=transforms.Compose([ Rescale(256), RandomCrop(224), ToTensor()]) )

数据转换(Transfrom)发生在数据库中的__getitem__操作中。以上代码中,transforms.Compose(transform_list),Compose即组合的意思,其参数是一个转换操作的列表。如上是[ Rescale(256), RandomCrop(224), ToTensor()],以下是实现这三个转换类。我们将把它们写成可调用的类,而不是简单的函数,这样在每次调用转换时就不需要传递它的参数。为此,我们只需要实现__call__方法,如果需要,还需要实现__init__方法。然后我们可以使用这样的变换:

 

#创建一个转换可调用类的实例
tsfm = Transform(params)
#使用转换操作实例对样本sample进行转换
transformed_sample = tsfm(sample)

 

下面观察这些转换是如何应用于图像和标注的。(注:每一个操作对应一个类

class Rescale(object):
    """Rescale the image in a sample to a given size.

    Args:
        output_size (tuple or int): Desired output size. If tuple, output is
            matched to output_size. If int, smaller of image edges is matched
            to output_size keeping aspect ratio the same.
    """

    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))
        self.output_size = output_size

    def __call__(self, sample):
        image, landmarks = sample['image'], sample['landmarks']

        h, w = image.shape[:2]
        if isinstance(self.output_size, int):
            if h > w:
                new_h, new_w = self.output_size * h / w, self.output_size
            else:
                new_h, new_w = self.output_size, self.output_size * w / h
        else:
            new_h, new_w = self.output_size

        new_h, new_w = int(new_h), int(new_w)

        img = transform.resize(image, (new_h, new_w))

        # h and w are swapped for landmarks because for images,
        # x and y axes are axis 1 and 0 respectively
        landmarks = landmarks * [new_w / w, new_h / h]

        return {'image': img, 'landmarks': landmarks}


class RandomCrop(object):
    """Crop randomly the image in a sample.

    Args:
        output_size (tuple or int): Desired output size. If int, square crop
            is made.
    """

    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))
        if isinstance(output_size, int):
            self.output_size = (output_size, output_size)
        else:
            assert len(output_size) == 2
            self.output_size = output_size

    def __call__(self, sample):
        image, landmarks = sample['image'], sample['landmarks']

        h, w = image.shape[:2]
        new_h, new_w = self.output_size

        top = np.random.randint(0, h - new_h)
        left = np.random.randint(0, w - new_w)

        image = image[top: top + new_h,
                      left: left + new_w]

        landmarks = landmarks - [left, top]

        return {'image': image, 'landmarks': landmarks}


class ToTensor(object):
    """Convert ndarrays in sample to Tensors."""

    def __call__(self, sample):
        image, landmarks = sample['image'], sample['landmarks']

        # swap color axis because
        # numpy image: H x W x C
        # torch image: C X H X W
        image = image.transpose((2, 0, 1))
        return {'image': torch.from_numpy(image),
                'landmarks': torch.from_numpy(landmarks)}

以下来介绍转换的用法。

#获取一条数据
sample = face_dataset[index]
#单独进行操作
scale = Rescale(256)
crope= RandomCrop(224)
scale(sample)
crope(sample)
#使用Compose组合操作
compose = transforms.Compose([Rescale(256),RandomCrop(224)])
compose(sample)

上述转换后数据仍然是PIL类型,如果要求返回是一个tensor,那么还得在Compose的最后一个元素进行Totensor操作。

posted @ 2019-02-20 14:35  HOU_JUN  阅读(13584)  评论(0编辑  收藏  举报