处理YOLOv5 TFlite的输出数据

Question

❔问题

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

我得到的输出：

• StatefulPartitionedCall:0=[125200,7]来自转换后的YOLOv5模型Netron YOLOv5s.tflite型号

但我希望输出如下：

• StatefulPartitionedCall:3=[1，10，4]#个框
• StatefulPartitionedCall:2=[1，10]#类
• StatefulPartitionedCall:1=[1，10]#分数
• StatefulPartitionedCall:0=[1]#count(这一个来自tensorflow lite移动网络模型(训练为提供10个输出数据，默认为tflite))Netron mobilenet.tflite型号

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

[125200，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)
...