如何使用分割模型输出张量？

Question

我正在尝试运行 iOS 上的分割模型，我有几个关于如何正确使用输出张量的问题。

这里是我使用的模型上的 link： https://www.tensorflow.org/lite/models/segmentation/overview

当我运行这个模型时，我得到了维度的输出张量： 1 x 257 x 257 x 21。为什么我得到 21 作为最后一个维度？看起来每个像素我们都得到 class 分数。我们需要在这里找到 argmax 以获得正确的 class 值吗？

但为什么只有 21 classes？我在想它应该包含更多。我在哪里可以找到某个值对应于某个 class 的信息。在 ImageClassification 示例中，我们有一个 label.txt 和 1001 classes.

基于 ImageClassification 示例，我尝试解析张量：首先将其转换为大小为 1 387 029 (21 x 257 x 257) 的 Float 数组，然后使用以下代码创建图像像素像素：

    // size = 257
    // depth = 21
    // array - float array of size 1 387 029
    for i in 0..<size {
        for j in 0..<size {
            var scores: [Float] = []
            for k in 0..<depth {
                let index = i * size * depth + j * depth + k
                let score = array[index]
                scores.append(score)
            }
            if let maxScore = scores.max(),
                let maxClass = scores.firstIndex(of: maxScore) {
                let index = i * size + j

                if maxClass == 0 {
                    pixelBuffer[index] = .blue
                } else if maxClass == 12 {
                    pixelBuffer[index] = .black
                } else {
                    pixelBuffer[index] = .green
                }
            }
        }
    }

这是我得到的结果：

可以看出质量不是很好。我错过了什么？

CoreML 的分割模型（https://developer.apple.com/machine-learning/models/）在同一个例子中效果更好：

Answer 1

您的模型似乎是在 PASCAL VOC 数据上训练的，该数据有 21 类用于分割。
您可以找到类 here:

的列表

background
aeroplane
bicycle
bird
boat
bottle
bus
car
cat
chair
cow
diningtable
dog
horse
motorbike
person
pottedplant
sheep
sofa
train
tvmonitor

Answer 2

除了 Shai 的回答之外，您还可以使用 Netron 之类的工具来可视化您的网络并更深入地了解输入和输出，例如您的输入将是尺寸为 257x257x3 的图像：

并且您已经知道输出大小，对于分割模型，您得到的是 21，因为这是类您的模型所支持的数量，正如 Shai 提到的那样，然后为所有类这应该会给你一个更不错的输出，不需要调整任何东西的大小，尝试类似的东西（在伪代码中）：

output = [rows][cols]
for i in rows:
  for j in cols:
    argmax = -1
    for c in classes:
      if tensor_out[i][j][c] > argmax:
        argmax = tensor_out[i][j][c]
    output[i][j] = c

然后输出就是你的分割图像。

如何使用分割模型输出张量？

How to use segmentation model output tensor?

image-segmentation

ios

tensorflow

semantic-segmentation

tensorflow-lite