0

我正在尝试实现多类图像分割任务。图像的掩码属于以下类型 -

面具

我需要将图像转换为与每个类对应的标签，其中类是红色、绿色、蓝色和黑色（背景），我正在使用下面的代码 -

 def mask_to_class(self,mask):
    target = torch.from_numpy(mask)
    h,w = target.shape[0],target.shape[1]
    masks = torch.empty(h, w, dtype=torch.long)
    colors = torch.unique(target.view(-1,target.size(2)),dim=0).numpy()
    target = target.permute(2, 0, 1).contiguous()
    mapping = {tuple(c): t for c, t in zip(colors.tolist(), range(len(colors)))}
    for k in mapping:
        idx = (target==torch.tensor(k, dtype=torch.uint8).unsqueeze(1).unsqueeze(2))
        validx = (idx.sum(0) == 3) 
        masks[validx] = torch.tensor(mapping[k], dtype=torch.long)
   return masks

问题是图像中存在的不同颜色的数量应该是 4，但在这种情况下它们是 338，而且它们在不同的图像中也是不规则的。

可以使用以下代码重现问题 -

image = Image.open("Image_Path")
image = np.array(image)
target = torch.from_numpy(image)
h,w = target.shape[0],target.shape[1]
masks = torch.empty(h, w, dtype=torch.long)
colors = torch.unique(target.view(-1,target.size(2)),dim=0).numpy()
print(colors.shape)

(338, 3)

颜色的形状应该是(4,3)，但它来了338,3)。我无法找到它背后的原因。下面是另一个图像，它与上图的颜色相同，并给出了(4,3)所需的颜色形状。我在哪里做错了？ 图 2

4

0 回答 0