如何在 AWS Sagemaker 中检索分段掩码中使用的标签

Question

从分段掩码中，我试图检索掩码中表示的标签。

这是我在 AWS Sagemaker 中通过语义分割模型运行的图像。

用于进行预测和显示掩码的代码。

from sagemaker.predictor import json_serializer, json_deserializer, RealTimePredictor
from sagemaker.content_types import CONTENT_TYPE_CSV, CONTENT_TYPE_JSON

%%time
ss_predict = sagemaker.RealTimePredictor(endpoint=ss_model.endpoint_name, 
                                     sagemaker_session=sess,
                                    content_type = 'image/jpeg',
                                    accept = 'image/png')

return_img = ss_predict.predict(img)

from PIL import Image
import numpy as np
import io

num_labels = 21
mask = np.array(Image.open(io.BytesIO(return_img)))
plt.imshow(mask, vmin=0, vmax=num_labels-1, cmap='jet')
plt.show()

该图像是创建的分割蒙版，它代表摩托车，其他一切都是背景。

[

从代码中可以看出，有 21 个可能的标签，其中 2 个用于蒙版，一个用于摩托车，另一个用于背景。我现在想弄清楚的是，如何在 21 个可能的选项中打印出这个掩码中实际使用的标签？

如果您需要任何进一步的信息，请告诉我，我们非常感谢您的帮助。

Answer 1

在某个地方你应该有从标签整数到标签类的映射，例如

label_map = {0: 'background', 1: 'motorbike', 2: 'train', ...}

如果您使用 Pascal VOC 数据集，则为 (1=飞机、2=自行车、3=鸟、4=船、5=瓶子、6=公共汽车、7=汽车、8=猫、9=椅子、10 =牛、11=餐桌、12=狗、13=马、14=摩托车、15=人、16=盆栽、17=羊、18=沙发、19=火车、20=电视/显示器） - 参见此处：http://host.robots.ox.ac.uk/pascal/VOC/voc2012/segexamples/index.html

然后你可以简单地使用该地图：

used_classes = np.unique(mask)
for cls in used_classes:
    print("Found class: {}".format(label_map[cls]))