Jetson设备Yolov7 Tiny+Darknet+Python3.6列表索引错误求助

环境信息 #

• Jetson设备，Python 3.6.9
• YOLOv7 Tiny + Darknet框架
• 本地2021款M1 Mac运行正常，Jetson上触发报错

错误详情 #

错误栈：

Traceback (most recent call last):
  ...
  File "main.py", line 267, in start_detection
    detections = model.process_frame(frame_model)
  File "/home/jetson/plate-jetson/plate-jetson/yolo.py", line 87, in process_frame
    box = boxes[i]
TypeError: only integer scalar arrays can be converted to a scalar index

打印indices结果：

[[  508]
 [ 8739]
 [ 8417]
 ...
 [ 2405]
 [ 9898]
 [10138]]

相关代码 #

import cv2
import numpy as np
        
class YoloDetection():
    def __init__(self, model_path: str, config: str, classes: str, width: int, height: int,
                 scale=.00392, thr=0.1, nms=0.4,#00392#.4
                 backend=0,
                 framework=3,
                 target=0,
                 mean=[0, 0, 0]):
        
        super(YoloDetection,self).__init__()
        choices = ['caffe', 'tensorflow', 'torch', 'darknet']
        backends = (
            cv2.dnn.DNN_BACKEND_DEFAULT, cv2.dnn.DNN_BACKEND_HALIDE, cv2.dnn.DNN_BACKEND_INFERENCE_ENGINE,
            cv2.dnn.DNN_BACKEND_OPENCV)
        targets = (
            cv2.dnn.DNN_TARGET_CPU, cv2.dnn.DNN_TARGET_OPENCL, cv2.dnn.DNN_TARGET_OPENCL_FP16, cv2.dnn.DNN_TARGET_MYRIAD)

        self.__confThreshold = thr
        self.__nmsThreshold = nms
        self.__mean = mean
        self.__scale = scale
        self.__width = width
        self.__height = height

        # Load a network
        self.__net = cv2.dnn.readNet(model_path, config, choices[framework])
        self.__net.setPreferableBackend(backends[backend])
        self.__net.setPreferableTarget(targets[target])
        self.__classes = None

        if classes:
            with open(classes, 'rt') as f:
                self.__classes = f.read().rstrip('\n').split('\n')


    def get_output_layers_name(self, net):
        return net.getUnconnectedOutLayersNames()

    def post_process_output(self, frame, outs):
        frame_height = frame.shape[0]
        frame_width = frame.shape[1]

        class_ids = []
        confidences = []
        boxes = []

        class_ids = []
        confidences = []
        boxes = []
        for out in outs:
            for detection in out:
                scores = detection[5:]
                class_id = np.argmax(scores)
                confidence = scores[class_id]
                if confidence > self.__confThreshold:
                    center_x = int(detection[0] * frame_width)
                    center_y = int(detection[1] * frame_height)
                    width = int(detection[2] * frame_width)
                    height = int(detection[3] * frame_height)
                    left = center_x - width / 2
                    top = center_y - height / 2
                    class_ids.append(class_id)
                    confidences.append(float(confidence))
                    boxes.append([left, top, width, height])

        
        indices = cv2.dnn.NMSBoxes(boxes, confidences, self.__confThreshold, self.__nmsThreshold)
        return (indices, boxes, confidences, class_ids)

    def process_frame(self, frame: np.ndarray):
        frame_height = frame.shape[0]
        frame_width = frame.shape[1]

        blob = cv2.dnn.blobFromImage(frame, self.__scale, (self.__width, self.__height), self.__mean, True, crop=False)

        # Run a model
        self.__net.setInput(blob)
        outs = self.__net.forward(self.get_output_layers_name(self.__net))
        (indices, boxes, confidences, class_ids) = self.post_process_output(frame, outs)
        detected_objects = []

        for i in indices:
            
            box = boxes[i]
            left = box[0]
            top = box[1]
            width = box[2]
            height = box[3]
            x = int(left)
            y = int(top)
            nw = int(width)
            nh = int(height)
            if x < 0:
                x = 0
            if y < 0:
                y = 0
            if x + nw > frame_width:
                nw = frame_width - x
            if y + nh > frame_height:
                nh = frame_height - y
            detected_objects.append([self.__classes[class_ids[i]], x, y, nw, nh, confidences[i]])
        return detected_objects

解决方案 #

问题根源是不同设备上的OpenCV版本对cv2.dnn.NMSBoxes的返回值格式处理不一致：Mac上返回一维整数数组，而Jetson上返回二维数组（每个元素是单元素数组），直接用数组去索引列表boxes就会触发类型错误。

只需要统一indices的格式为一维整数数组即可，有两种修改方式：

方式1：在process_frame中处理格式 #

修改循环前的代码，将二维数组扁平化：

def process_frame(self, frame: np.ndarray):
    frame_height = frame.shape[0]
    frame_width = frame.shape[1]

    blob = cv2.dnn.blobFromImage(frame, self.__scale, (self.__width, self.__height), self.__mean, True, crop=False)

    # Run a model
    self.__net.setInput(blob)
    outs = self.__net.forward(self.get_output_layers_name(self.__net))
    (indices, boxes, confidences, class_ids) = self.post_process_output(frame, outs)
    detected_objects = []

    # 统一indices格式为一维数组
    indices = indices.flatten() if len(indices.shape) > 1 else indices

    for i in indices:
        # 确保i是整数标量
        i = int(i)
        box = boxes[i]
        left = box[0]
        top = box[1]
        width = box[2]
        height = box[3]
        x = int(left)
        y = int(top)
        nw = int(width)
        nh = int(height)
        if x < 0:
            x = 0
        if y < 0:
            y = 0
        if x + nw > frame_width:
            nw = frame_width - x
        if y + nh > frame_height:
            nh = frame_height - y
        detected_objects.append([self.__classes[class_ids[i]], x, y, nw, nh, confidences[i]])
    return detected_objects

方式2：在post_process_output中处理格式 #

直接在返回前修改indices：

def post_process_output(self, frame, outs):
    # ... 原有代码省略 ...
        
    indices = cv2.dnn.NMSBoxes(boxes, confidences, self.__confThreshold, self.__nmsThreshold)
    # 统一indices格式为一维数组
    if len(indices.shape) > 1:
        indices = indices.flatten()
    return (indices, boxes, confidences, class_ids)

两种方式都能让indices成为一维整数数组，循环中的i会变成整数标量，即可正常索引列表元素，解决报错问题。

内容的提问来源于stack exchange，提问作者Matthias