解码 Core ML YOLO 对象检测器

• 下载带循环的解码 - 690.3 KB

引言

本系列文章假设您熟悉 Python、Conda 和 ONNX，并且具有在 Xcode 中开发 iOS 应用程序的一些经验。欢迎您下载此项目的源代码。我们将使用 macOS 10.15+、Xcode 11.7+ 和 iOS 13+ 运行代码。

理解 YOLO v2 的输出

YOLO v2 接受固定分辨率为 416 x 416 的输入图像，这些图像被分割成一个 13 x 13 的网格。该模型的预测返回一个形状为 (1, 425, 13, 13) 的单一数组。第一维度表示批次（这对于我们的目的并不重要），最后两个维度对应于 13 x 13 网格。但是，我们在上一篇文章中看到的每个单元格中的 425 个值是什么？

这些值包含有关检测到的对象的置信度分数和相应的边界框坐标的编码信息

[x1,y1,w1,h1,s1,c011,c021,c031,…,c791,c801,x2,y2,…,x5,y5,w5,h5,s5,c015,…c795,c805],

其中

• i – 给定网格单元格内的边界框索引（值：1-5）
• x_i, y_i, w_i, h_i – 框的坐标（x、y、宽度和高度，分别）用于框 i
• s_i – 给定单元格包含一个对象的置信度分数
• c01_i - c80_i – COCO 数据集中包含的 80 个对象类中每个类的置信度分数。

快速检查：每个单元格 5 个框乘以 85 个值（四个坐标，每个单元格一个置信度分数 + 每个对象类 80 个置信度分数）正好等于 425。

准备 YOLO 输出解码

我们需要几个常量

GRID_SIZE = 13
CELL_SIZE = int(416 / GRID_SIZE)
BOXES_PER_CELL = 5

ANCHORS = [[0.57273, 0.677385], 
           [1.87446, 2.06253], 
           [3.33843, 5.47434], 
           [7.88282, 3.52778], 
           [9.77052, 9.16828]]

GRID_SIZE 反映了 YOLO 如何将图像分割成单元格，CELL_SIZE 描述了每个单元格的宽度和高度（以像素为单位），而 BOXES_PER_CELL 是模型为每个单元格考虑的预定义框的数量。

ANCHORS 数组包含用于计算每个单元格中五个框中每个框的坐标的因子。请注意，不同的 YOLO 版本使用不同的锚点，因此您始终需要检查模型训练使用了哪些值。以上值在 原始 YOLO 存储库（yolov2.cfg 文件）中找到。

我们还需要从附加的 coco_names.txt 文件中加载检测到的对象的标签

with open('./models/coco_names.txt', 'r') as f:
    COCO_CLASSES = [c.strip() for c in f.readlines()]

COLO_CLASSES 列表的前几个元素是

['person', 'bicycle', 'car', 'motorbike', 'aeroplane', 'bus', 'train', 'truck', 'boat', 'traffic light', ...]

解码 YOLO 输出

我们的 YOLO v2 模型返回“原始”神经网络输出，未通过激活函数进行归一化。为了理解它们，我们需要两个额外的函数

def sigmoid(x):
    k = np.exp(-x)
    return 1 / (1 + k)

def softmax(x):
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum()

不详细说明：sigmoid 为任何输入返回一个 0-1 范围内的值，而 softmax 为任何输入向量返回一个归一化值，其值的总和等于 1。

现在我们可以编写我们的主要解码函数

def decode_preds(raw_preds: []):
    num_classes = len(COCO_CLASSES)
    decoded_preds = []
    for cy in range(GRID_SIZE):
        for cx in range(GRID_SIZE):
            for b in range(BOXES_PER_CELL):
                box_shift = b*(num_classes + 5)
            
                tx = float(raw_preds[0, box_shift    , cy, cx])
                ty = float(raw_preds[0, box_shift + 1, cy, cx])
                tw = float(raw_preds[0, box_shift + 2, cy, cx])
                th = float(raw_preds[0, box_shift + 3, cy, cx])
                ts = float(raw_preds[0, box_shift + 4, cy, cx])

                x = (float(cx) + sigmoid(tx)) * CELL_SIZE
                y = (float(cy) + sigmoid(ty)) * CELL_SIZE
            
                w = np.exp(tw) * ANCHORS[b][0] * CELL_SIZE
                h = np.exp(th) * ANCHORS[b][1] * CELL_SIZE
            
                box_confidence = sigmoid(ts)
                classes_raw = raw_preds[0, box_shift + 5:box_shift + 5 + num_classes, cy, cx]
                classes_confidence = softmax(classes_raw)
            
                box_class_idx = np.argmax(classes_confidence)
                box_class_confidence = classes_confidence[box_class_idx]

                combined_confidence = box_confidence * box_class_confidence
            
                decoded_preds.append([box_class_idx, combined_confidence, x, y, w, h])            
    
    return sorted(decoded_preds, key=lambda p: p[1], reverse=True)

首先，该函数在每个网格单元格（cy、cx 和 b 循环）内的框上进行迭代，以解码每个边界框的后续值（假设批次中只有一个图像，因此我们使用 raw_preds[0,…]）。然后，模型返回的原始 tx、ty、tw、th 和 ts 值用于计算边界框坐标（中心 x、中心 y、宽度和高度）、box_confidence（给定框包含对象的置信度）和 class_confidence（包含 80 个 COCO 类中每个类的归一化置信度的向量）。配备了这些值，我们使用当前框（box_class_idx 及其 combined_confidence）计算检测到的最可能对象的类。

在所有计算之后，该方法返回一个按置信度分数降序排序的已解码值的列表。

让我们看看它是否适用于我们来自 Open Images 数据集 的图像

image = load_and_scale_image('https://c2.staticflickr.com/4/3393/3436245648_c4f76c0a80_o.jpg')
cml_model = ct.models.MLModel('./models/yolov2-coco-9.mlmodel')
preds = cml_model.predict(data={'input.1': image})['218']
decoded_preds = decode_preds(preds)
print([p[:3] for p in decoded_preds[:2]])

该模型似乎确信图片包含一个人和一只狗。我们应该检查这是否属实

import copy

def annotate_image(image, preds, min_score=0.5, top=10):
    annotated_image = copy.deepcopy(image)
    draw = ImageDraw.Draw(annotated_image)
    w,h = image.size
    
    colors = ['red', 'orange', 'yellow', 'green', 'blue', 'white']
    
    for class_id, label, score, xc, yc, w, h in decoded_preds[:top]:
        if score < min_score:
            continue
            
        x0 = xc - (w / 2)
        y0 = yc - (h / 2)
        color = ImageColor.colormap[colors[class_id % len(colors)]]
        draw.rectangle([(x0, y0), (x0 + w, y0 + h)], width=2, outline=color)
        draw.text((x0 + 5, y0 + 5), "{} {:0.2f}".format(label, score), fill=color)
    
    return annotated_image

annotate_image(image, decoded_preds)

不错……但重复的框有什么问题？没问题。这是 YOLO 工作方式的副作用 - 多个框（总共 425 个）可能会检测到相同的对象。我们可以通过设置最小置信度分数来稍微改进这一点。不过，目前，我们不用担心。我们将很快使用一种称为非最大抑制的算法来解决这个问题。

后续步骤

我们已经成功解码了 YOLO v2 输出，这使我们能够可视化模型的预测。如果您不喜欢代码中的循环……您是对的。这并不是我们应该对数组执行计算的方式。我们以这种方式构造代码只是为了理解解码过程。在下一篇文章中，我们将使用数组运算做同样的事情。这将使我们能够将解码逻辑直接包含在模型中。