用于语义分割的 Pytorch 自定义 randomcrop

Question

我正在尝试实现自定义数据集加载器。首先，我以 (0.98, 1,1) 之间的相同比例调整图像和标签的大小，然后我随机裁剪具有相同参数的图像和标签，以便我可以将它们输入神经网络。但是，我从 PyTorch 函数中收到错误。这是我的代码：

class RandomCrop(object):

    def __init__(self, size, padding=None, pad_if_needed=True, fill=0, padding_mode='constant'):
        
        self.size = size
        self.padding = padding
        self.pad_if_needed = pad_if_needed
        self.fill = fill
        self.padding_mode = padding_mode

    @staticmethod
    def get_params(img, output_size):
        
        w, h = img.size
        th, tw = output_size
        if w == tw and h == th:
            return 0, 0, h, w

        i = random.randint(0, h - th)
        j = random.randint(0, w - tw)
        return i, j, th, tw

    def __call__(self, data):
     
        img,mask = data["image"],data["mask"]
       

        # pad the width if needed
        if self.pad_if_needed and img.size[0] < self.size[1]:
            img = F.pad(img, (self.size[1] - img.size[0], 0), self.fill, self.padding_mode)
            mask = F.pad(mask, (self.size[1] - mask.size[0], 0), self.fill, self.padding_mode)
        # pad the height if needed
        if self.pad_if_needed and img.size[1] < self.size[0]:
            img = F.pad(img, (0, self.size[0] - img.size[1]), self.fill, self.padding_mode)
            mask = F.pad(mask, (0, self.size[0] - mask.size[1]), self.fill, self.padding_mode)
       
        i, j, h, w = self.get_params(img, self.size)
        crop_image = transforms.functional.crop(img, i, j, h, w)
        crop_mask = transforms.functional.crop(mask, i, j, h, w)

        return{"image": crop_image, "mask": crop_mask }

这里是错误：

AttributeError: 'Image' object has no attribute 'dim'

Answer 1

我错误地导入了 nn.functional.pad 而不是 transforms.functional.pad。改完之后一切顺利

用于语义分割的 Pytorch 自定义 randomcrop

Pytorch custom randomcrop for semantic segmentation

pytorch

torchvision