处理来自 YOLOv5 TFlite 的输出数据

Posted 2023-03-12

【中文标题】处理来自 YOLOv5 TFlite 的输出数据

【英文标题】：Process output data from YOLOv5 TFlite

【发布时间】：2021-04-25 17:07:29

【问题描述】：

❔问题

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

我得到输出：

StatefulPartitionedCall：0 = [1,25200,7] 来自转换后的 YOLOv5 模型 Netron YOLOv5s.tflite model

但我希望输出如下：

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

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

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

0.011428807862102985, 0.006756599526852369, 0.04274776205420494, 0.034441519528627396, 0.00012877583503723145, 0.33658933639526367, 0.4722323715686798
0.023071227595210075, 0.006947836373001337, 0.046426184475421906, 0.023744791746139526, 0.0002465546131134033, 0.29862138628959656, 0.4498370885848999
0.03636947274208069, 0.006819264497607946, 0.04913407564163208, 0.025004519149661064, 0.00013208389282226562, 0.3155967593193054, 0.4081345796585083
0.04930267855525017, 0.007249316666275263, 0.04969717934727669, 0.023645592853426933, 0.0001222355494974181, 0.3123127520084381, 0.40113094449043274
...

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

【问题讨论】：

显然创建模型的过程有错误...或转换为tflite。您应该在 Netron 看到 3-4 个输出，而不仅仅是 1 个。再次检查...模型的创建

【参考方案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)
        ...