使用 pytorch 应用简单的转换来获取二值图像

Applying a simple transformation to get a binary image using pytorch

我想在将图像传递给数据加载器之前将其二值化,我创建了一个运行良好的数据集 class。但在 __getitem__() 方法中,我想对图像进行阈值处理:

    def __getitem__(self, idx):
        # Open image, apply transforms and return with label
        img_path = os.path.join(self.dir, self.filelist[filename"])
        image = Image.open(img_path)
        label = self.x_data.iloc[idx]["label"]

        # Applying transformation to the image
        if self.transforms is not None:
           image = self.transforms(image)

        # applying threshold here:
        my_threshold = 240
        image = image.point(lambda p: p < my_threshold and 255)
        image = torch.tensor(image)

        return image, label

然后我尝试调用数据集:

    data_transformer = transforms.Compose([
        transforms.Resize((10, 10)),
        transforms.Grayscale()
        //transforms.ToTensor()
    ])

train_set = MyNewDataset(data_path, data_transformer, rows_train)

由于我在 PIL 对象上应用了阈值,我需要在之后应用到张量对象的转换,但由于某种原因它崩溃了。有人可以帮助我吗?

为什么不在从 PIL.Image 转换为 torch.Tensor 后应用二值化?

class ThresholdTransform(object):
  def __init__(self, thr_255):
    self.thr = thr_255 / 255.  # input threshold for [0..255] gray level, convert to [0..1]

  def __call__(self, x):
    return (x > self.thr).to(x.dtype)  # do not change the data type

完成此转换后,只需添加它:

data_transformer = transforms.Compose([
        transforms.Resize((10, 10)),
        transforms.Grayscale(),
        transforms.ToTensor(),
        ThresholdTransform(thr_255=240)
    ])