处理来自 YOLOv5 TFlite 的输出数据

如何解决处理来自 YOLOv5 TFlite 的输出数据

❔问题

嗨，我已经成功训练了一个基于 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,7] 数组中获取框、类和分数。（2021 年 1 月 15 日，我将 pytorch、tensorflow 和 yolov5 更新到最新版本）

包含在 [1,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)

解决方法

感谢@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,h = boxes[...,0],boxes[...,1],2],3] #xywh
    xyxy = [x - w / 2,y - h / 2,x + w / 2,y + h / 2]  # xywh to xyxy   [4,25200]

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

获取输出数据：

"""Output data"""
output_data = interpreter.get_tensor(output_details[0]['index'])  # get tensor  x(1,7)
xyxy,scores = YOLOdetect(output_data) #boxes(x,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)
        ...