处理 YOLOv5 TFlite 的输出数据

Question

\xe2\x9d\x94问题

\n

您好，我已经成功训练了一个基于 YOLOv5s 的自定义模型，并将该模型转换为 TFlite。我觉得这个问题很傻，但是你如何使用输出数据呢？

\n

我得到的输出：

\n

\n
• StatefulPartitionedCall: 0 = [1,25200,7]\n来自转换后的 YOLOv5 模型\n Netron YOLOv5s.tflite 模型
\n

\n

但我期望输出如下：

\n

\n
• StatefulPartitionedCall:3 = [1, 10, 4] # 个框
\n
• StatefulPartitionedCall:2 = [1, 10] # 类
\n
• StatefulPartitionedCall:1 = [1, 10] #scores
\n
• StatefulPartitionedCall:0 = [1] #count\n（这个来自 TensorFlow lite mobilenet 模型（训练为提供 10 个输出数据，默认为 tflite））\n Netron mobilenet.tflite 模型
\n

\n

它也可能是其他形式的输出，但老实说，我不知道如何从 [1,25200,7] 数组中获取框、类、分数。\n（2021 年 1 月 15 日，我更新了 pytorch、tensorflow和yolov5到最新版本）

\n

[1, 25200, 7]数组中包含的数据可以在这个文件中找到：outputdata.txt

\n

0.011428807862102985, 0.006756599526852369, 0.04274776205420494, 0.034441519528627396, 0.00012877583503723145, 0.33658933639526367, 0.4722323715686798\n0.023071227595210075, 0.006947836373001337, 0.046426184475421906, 0.023744791746139526, 0.0002465546131134033, 0.29862138628959656, 0.4498370885848999\n0.03636947274208069, 0.006819264497607946, 0.04913407564163208, 0.025004519149661064, 0.00013208389282226562, 0.3155967593193054, 0.4081345796585083\n0.04930267855525017, 0.007249316666275263, 0.04969717934727669, 0.023645592853426933, 0.0001222355494974181, 0.3123127520084381, 0.40113094449043274\n...\n

\n

我应该添加非最大抑制还是其他东西，有人可以帮助我吗？（github YOLOv5 #1981）

\n

Answer 1

感谢@Glenn Jocher，我找到了解决方案。输出是[xywh, conf, class0, class1, ...]

我当前的代码现在是：

def classFilter(classdata):
    classes = []  # create a list
    for i in range(classdata.shape[0]):         # loop through all predictions
        classes.append(classdata[i].argmax())   # get the best classification location
    return classes  # return classes (int)

def YOLOdetect(output_data):  # input = interpreter, output is boxes(xyxy), classes, scores
    output_data = output_data[0]                # x(1, 25200, 7) to x(25200, 7)
    boxes = np.squeeze(output_data[..., :4])    # boxes  [25200, 4]
    scores = np.squeeze( output_data[..., 4:5]) # confidences  [25200, 1]
    classes = classFilter(output_data[..., 5:]) # get classes
    # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
    x, y, w, h = boxes[..., 0], boxes[..., 1], boxes[..., 2], boxes[..., 3] #xywh
    xyxy = [x - w / 2, y - h / 2, x + w / 2, y + h / 2]  # xywh to xyxy   [4, 25200]

    return xyxy, classes, scores  # output is boxes(x,y,x,y), classes(int), scores(float) [predictions length]

获取输出数据：

"""Output data"""
output_data = interpreter.get_tensor(output_details[0]['index'])  # get tensor  x(1, 25200, 7)
xyxy, classes, scores = YOLOdetect(output_data) #boxes(x,y,x,y), classes(int), scores(float) [25200]

对于盒子：

for i in range(len(scores)):
    if ((scores[i] > 0.1) and (scores[i] <= 1.0)):
        H = frame.shape[0]
        W = frame.shape[1]
        xmin = int(max(1,(xyxy[0][i] * W)))
        ymin = int(max(1,(xyxy[1][i] * H)))
        xmax = int(min(H,(xyxy[2][i] * W)))
        ymax = int(min(W,(xyxy[3][i] * H)))

        cv2.rectangle(frame, (xmin,ymin), (xmax,ymax), (10, 255, 0), 2)
        ...