无法获得TFLITE Yolov5型号的盒子

Question

我训练了一个模型，通过Yolov5可以检测图像上的'+'字符。我想在TFLITE中使用这个模型。但是，当我在模型中推断图像时，我在解释输出时遇到了麻烦。

在我的模型中我是这样推断的:

interpreter = tf.lite.Interpreter("/Users/maximereder/Desktop/best-fp16.tflite")
interpreter.allocate_tensors()
IMAGE_PATH = "/Users/maximereder/Documents/ML/dataset/plus- 
1000x750/train/IMG_1492B2.jpg"
img = cv2.resize(cv2.imread(IMAGE_PATH), (640, 640))
features = img.copy()
np_features = np.array(features, dtype=np.float32)
np_features = np.expand_dims(np_features, axis=0)
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
interpreter.set_tensor(input_details[0]['index'], np_features)
interpreter.invoke()

out_details:

[{'name': 'Identity',
'index': 422,
'shape': array([    1, 25200,    85], dtype=int32),
'shape_signature': array([    1, 25200,    85], dtype=int32),
'dtype': numpy.float32,
'quantization': (0.0, 0),
'quantization_parameters': {'scales': array([], dtype=float32),
'zero_points': array([], dtype=int32),
'quantized_dimension': 0},
'sparsity_parameters': {}}]

当我想解释模型的输出时:

detection_boxes = interpreter.get_tensor(output_details[0]['index'])
detection_classes = interpreter.get_tensor(output_details[1]['index'])
detection_scores = interpreter.get_tensor(output_details[2]['index'])
num_boxes = interpreter.get_tensor(output_details[3]['index'])

解释器输出。get_tensor (output_details[0]("指数"):

array([[[-5.6185187e-03,  1.3539949e-02,  5.7405889e-02,  4.2354122e-02,
1.3114554e-04,  9.9999905e-01],
[-4.4679684e-03,  1.7201375e-02,  6.8576269e-02,  2.5891241e-02,
3.4220223e-04,  9.9999964e-01],
[-4.9383980e-03,  1.5453462e-02,  4.2031817e-02,  2.6558569e-02,
2.5974249e-03,  9.9999815e-01],
...,
[ 9.2678040e-01,  9.2856336e-01,  7.1995622e-01,  5.6132025e-01,
1.4253161e-14,  9.9999440e-01],
[ 9.2535079e-01,  9.2647862e-01,  9.4501650e-01,  1.1292257e+00,
4.9554409e-09,  9.9999994e-01],
[ 1.0224271e+00,  9.7982901e-01,  2.2890522e+00,  1.1467136e-02,
1.4553191e-07,  9.9999893e-01]]], dtype=float32)

我得到一个错误:IndexError: list index out of range

我理解错误的含义，但为什么只有一个元素在我的输出?我该怎么解释呢?我想去拿箱子。谢谢你。

Answer 1

请查看https://github.com/ultralytics/yolov5/issues/1981

'shape': array([1,252,85]， dtype=int32)

[x,y,w,h, conf, class1，…]]共85列

col 0-3为box, col 4为conf，其余80为class

获得真正的盒子，你需要做一些处理，如重新缩放xywh, NMS(非最大抑制)

您可以查看Yolov5源代码中的detect.py和utils/general.py