有条件地更改 RGB 图像中的像素值，然后删除通道

Question

我只想比较图像的像素，如果该像素是粉红色的（R值= 0.502，G值= 0.0，B值= 0.502），则将其更改为黑色，否则将其更改为白色。之后我想删除通道，只得到一个 (512,512,) 形状张量。

这个 getitem 在我的数据集类中。

代码：

    def __getitem__(self, index):
        img = Image.open(self.images[index]).convert("RGB") # this is one image
        mask = Image.open(self.masks[index]).convert("RGB") # this is the corresponding mask
        mask = self.transform_tensor(mask) # tensor is (3,512,512) shape
        mask = torch.where((mask[0, :, :] == 0.502) & (mask[1, :, :] == 0.0) & (mask[2, :, :] == 0.502), torch.tensor([0.0, 0.0, 0.0]), torch.tensor([1.0, 1.0, 1.0])) # error
        mask = self.transform_to_image(mask)
        mask.show()
        return self.transform_image(img), self.transform_mask(mask)

运行时错误：张量 a (512) 的大小必须与非单一维度 1 处的张量 b (3) 的大小匹配

错误：意外类型：（布尔、张量、张量）

Answer 1

看起来你真的只需要“匹配粉红色像素”的布尔数组，而不是更改像素然后删除通道，即

arr = np.all(img == [0.502, 0, 0.502], axis=-1)

您可以按原样使用此值，但如果您需要像素值，请将其乘以 1 或 255。