rknn_detect_yolov5.py

from rknn.api import RKNN
import cv2
import numpy as np
import cv2
import time

"""
yolov5 预测脚本 for rknn
"""


def sigmoid(x):
    return 1 / (1 + np.exp(-x))


def filter_boxes(boxes, box_confidences, box_class_probs) -> (np.ndarray, np.ndarray, np.ndarray):
    box_scores = box_confidences * box_class_probs  # 条件概率， 在该cell存在物体的概率的基础上是某个类别的概率
    box_classes = np.argmax(box_scores, axis=-1)  # 找出概率最大的类别索引
    box_class_scores = np.max(box_scores, axis=-1)  # 最大类别对应的概率值
    pos = np.where(box_class_scores >= OBJ_THRESH)  # 找出概率大于阈值的item
    # pos = box_class_scores >= OBJ_THRESH  # 找出概率大于阈值的item
    boxes = boxes[pos]
    classes = box_classes[pos]
    scores = box_class_scores[pos]
    return boxes, classes, scores


def nms_boxes(boxes, scores):
    x = boxes[:, 0]
    y = boxes[:, 1]
    w = boxes[:, 2]
    h = boxes[:, 3]

    areas = w * h
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)

        xx1 = np.maximum(x[i], x[order[1:]])
        yy1 = np.maximum(y[i], y[order[1:]])
        xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]])
        yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]])

        w1 = np.maximum(0.0, xx2 - xx1 + 0.00001)
        h1 = np.maximum(0.0, yy2 - yy1 + 0.00001)
        inter = w1 * h1

        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        inds = np.where(ovr <= NMS_THRESH)[0]
        order = order[inds + 1]
    keep = np.array(keep)
    return keep


def draw(image, boxes, scores, classes):
    """Draw the boxes on the image.

    # Argument:
        image: original image.
        boxes: ndarray, boxes of objects.
        classes: ndarray, classes of objects.
        scores: ndarray, scores of objects.
        all_classes: all classes name.
    """
    labels = []
    box_ls = []
    for box, score, cl in zip(boxes, scores, classes):
        x, y, w, h = box
        print('class: {}, score: {}'.format(CLASSES[cl], score))
        print('box coordinate left,top,right,down: [{}, {}, {}, {}]'.format(x, y, x + w, y + h))
        x *= image.shape[1]
        y *= image.shape[0]
        w *= image.shape[1]
        h *= image.shape[0]
        top = max(0, np.floor(x).astype(int))
        left = max(0, np.floor(y).astype(int))
        right = min(image.shape[1], np.floor(x + w + 0.5).astype(int))
        bottom = min(image.shape[0], np.floor(y + h + 0.5).astype(int))
        print('class: {}, score: {}'.format(CLASSES[cl], score))
        print('box coordinate left,top,right,down: [{}, {}, {}, {}]'.format(top, left, right, bottom))
        labels.append(CLASSES[cl])
        box_ls.append((top, left, right, bottom))
        cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)
        cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score),
                    (top, left - 6),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.6, (0, 0, 255), 2)
    return labels, box_ls


def load_model0(model_path, npu_id):
    rknn = RKNN()
    devs = rknn.list_devices()
    device_id_dict = {}
    for index, dev_id in enumerate(devs[-1]):
        if dev_id[:2] != 'TS':
            device_id_dict[0] = dev_id
        if dev_id[:2] == 'TS':
            device_id_dict[1] = dev_id

    print('-->loading model : ' + model_path)
    rknn.load_rknn(model_path)
    print('--> Init runtime environment on: ' + device_id_dict[npu_id])
    ret = rknn.init_runtime(device_id=device_id_dict[npu_id])
    if ret != 0:
        print('Init runtime environment failed')
        exit(ret)
    print('done')
    return rknn


def load_rknn_model(PATH):
    # Create RKNN object
    rknn = RKNN()
    # Load tensorflow model
    print('--> Loading model')
    ret = rknn.load_rknn(PATH)
    if ret != 0:
        print('load rknn model failed')
        exit(ret)
    print('done')
    ret = rknn.init_runtime(device_id='TS018083200400178', rknn2precompile=True)
    if ret != 0:
        print('Init runtime environment failed')
        exit(ret)
    print('done')
    return rknn


RKNN_MODEL_PATH = r"xxx.rknn"
SIZE = (416, 416)
CLASSES = ("classA", "classB")
OBJ_THRESH = 0.4
NMS_THRESH = 0.5
MASKS = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
ANCHORS = [[10, 13], [16, 30], [33, 23], [30, 61], [62, 45], [59, 119], [116, 90], [156, 198], [373, 326]]


def predict(img_src, rknn):
    img = cv2.resize(img_src, SIZE)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    t0 = time.time()
    print("img shape               \t:", img.shape)
    pred_onx = rknn.inference(inputs=[img])
    print("time:                    \t", time.time() - t0)
    boxes, classes, scores = [], [], []
    for t in range(3):
        input0_data = sigmoid(pred_onx[t][0])
        input0_data = np.transpose(input0_data, (1, 2, 0, 3))
        grid_h, grid_w, channel_n, predict_n = input0_data.shape
        anchors = [ANCHORS[i] for i in MASKS[t]]
        box_confidence = input0_data[..., 4]
        box_confidence = np.expand_dims(box_confidence, axis=-1)
        box_class_probs = input0_data[..., 5:]
        box_xy = input0_data[..., :2]
        box_wh = input0_data[..., 2:4]
        col = np.tile(np.arange(0, grid_w), grid_h).reshape(-1, grid_w)
        row = np.tile(np.arange(0, grid_h).reshape(-1, 1), grid_w)
        col = col.reshape((grid_h, grid_w, 1, 1)).repeat(3, axis=-2)
        row = row.reshape((grid_h, grid_w, 1, 1)).repeat(3, axis=-2)
        grid = np.concatenate((col, row), axis=-1)
        box_xy = box_xy * 2 - 0.5 + grid
        box_wh = (box_wh * 2) ** 2 * anchors
        box_xy /= (grid_w, grid_h)  # 计算原尺寸的中心
        box_wh /= SIZE  # 计算原尺寸的宽高
        box_xy -= (box_wh / 2.)  # 计算原尺寸的中心
        box = np.concatenate((box_xy, box_wh), axis=-1)
        res = filter_boxes(box, box_confidence, box_class_probs)
        boxes.append(res[0])
        classes.append(res[1])
        scores.append(res[2])
    boxes, classes, scores = np.concatenate(boxes), np.concatenate(classes), np.concatenate(scores)
    nboxes, nclasses, nscores = [], [], []
    for c in set(classes):
        inds = np.where(classes == c)
        b = boxes[inds]
        c = classes[inds]
        s = scores[inds]
        keep = nms_boxes(b, s)
        nboxes.append(b[keep])
        nclasses.append(c[keep])
        nscores.append(s[keep])
    if len(nboxes) < 1:
        return [], []
    boxes = np.concatenate(nboxes)
    classes = np.concatenate(nclasses)
    scores = np.concatenate(nscores)
    boxes, classes, scores
    label_list = []
    box_list = []
    for box, score, cl in zip(boxes, scores, classes):
        x, y, w, h = box
        x *= img_src.shape[1]
        y *= img_src.shape[0]
        w *= img_src.shape[1]
        h *= img_src.shape[0]
        top = max(0, np.floor(x).astype(int))
        left = max(0, np.floor(y).astype(int))
        right = min(img_src.shape[1], np.floor(x + w + 0.5).astype(int))
        bottom = min(img_src.shape[0], np.floor(y + h + 0.5).astype(int))
        label_list.append(CLASSES[cl])
        box_list.append((top, left, right, bottom))
    return label_list, np.array(box_list)


if __name__ == '__main__':
    RKNN_MODEL_PATH = r"xxx.rknn"
    SIZE = (416, 416)
    CLASSES = ("classA", "classB")
    OBJ_THRESH = 0.4
    NMS_THRESH = 0.5
    rknn = load_rknn_model(RKNN_MODEL_PATH)
    predict.__defaults__ = (None, rknn)
    img = cv2.imread("xxxx.png")
    res = predict(img)