处理来自 YOLOv5 TFlite 的输出数据

Question

❔问题

您好，我已经成功地训练了一个基于 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)

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)
        ...

处理来自 YOLOv5 TFlite 的输出数据

Process output data from YOLOv5 TFlite

tensorflow

pytorch

tensorflow-lite

yolov5

❔问题