pytorch - 将图像原始尺寸传递给变换函数

Question

我正在使用 pytorch 和 albumentations 在语义分割模型中进行图像转换。

对于某些转换，我需要传递原始图像大小作为参数，但我只是不知道如何做到这一点。

我的转换函数和数据集类位于单独的模块中，如下所示。在

medium_transform()

中的

transforms.py

函数中，我想访问我正在

original_height

中经过

original_width

的

self.transforms(**result)

和

dataset.py

。

数据集.py

from torch.utils.data import Dataset
import cv2

class SegmentationDataset(Dataset):
    def __init__(self, imagePaths, maskPaths, transforms):
        self.imagePaths = imagePaths
        self.maskPaths = maskPaths
        self.transforms = transforms

    def __getitem__(self, idx):
        image = cv2.imread(self.imagePaths[idx], cv2.COLOR_BGR2RGB)
        mask = cv2.imread(self.maskPaths[idx], 0)
        result = {"image": image, "mask": mask, "original_height": image.shape[0], "original_width": image.shape[1]}

        # check to see if we are applying any transformations
        if self.transforms is not None:
            result = self.transforms(**result)

        return result

变换.py

import albumentations as A
from . import config

def medium_transforms(**params):
    print(params) #here params is an empty dict
    return [
        A.OneOf([
          A.RandomSizedCrop(min_max_height=(50, 101), height=params["original_height"], width=params["original_width"], p=config.MEDIUM_TRANSFORMS_PROBABILITY),
          A.PadIfNeeded(min_height=params["original_height"], min_width=params["original_width"], p=config.MEDIUM_TRANSFORMS_PROBABILITY)
      ], p=1)
    ]

def compose(transforms_to_compose):
    # combine all augmentations into single pipeline
    return A.Compose([
      item for sublist in transforms_to_compose for item in sublist
    ])

train.py

import torch
from . import transforms
from .dataset import SegmentationDataset

trainImages = ["./images/test.png"]
trainMasks = ["./masks/test.png"]

train_transforms = transforms.compose([
  transforms.medium_transforms()
])

train_dataset = SegmentationDataset(imagePaths=trainImages, maskPaths=trainMasks, transforms=train_transforms)

Answer 1

一种解决方案是使用部分函数。

transforms.py：

import albumentations as A
from functools import partial
from . import config

def medium_transforms(original_height, original_width, **params):
    return [
        A.OneOf([
            A.RandomSizedCrop(min_max_height=(50, 101), height=original_height, width=original_width, p=config.MEDIUM_TRANSFORMS_PROBABILITY),
            A.PadIfNeeded(min_height=original_height, min_width=original_width, p=config.MEDIUM_TRANSFORMS_PROBABILITY)
        ], p=1)
    ]

def compose(transforms_to_compose):
    return A.Compose([item for sublist in transforms_to_compose for item in sublist])

medium_transforms_partial = partial(medium_transforms, original_height=None, original_width=None)

train.py：

import torch
from . import transforms
from .dataset import SegmentationDataset

trainImages = ["./images/test.png"]
trainMasks = ["./masks/test.png"]

train_transforms = transforms.compose([
  transforms.medium_transforms_partial
])

train_dataset = SegmentationDataset(imagePaths=trainImages, maskPaths=trainMasks, transforms=train_transforms)

在此代码中，

functools.partial

函数创建了一个部分应用函数

medium_transforms_partial

，并以original_height和original_width作为参数。这允许您在调用

medium_transforms_partial

时传递这些参数。