devbox
酷酷是懒虫
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

ONNX版本YOLOV5-DeepSort (rknn版本已经Ready)_rknn deepsort android

作者：酷酷是懒虫 | 2024-06-19 21:28:57

赞

踩

rknn deepsort android

目录

1. 前言

2. 储备知识

3. 准备工作

4. 代码修改的地方

5.结果展示

1. 前言

        之前一直在忙着写文档，之前一直做分类，检测和分割，现在看到跟踪算法，花了几天时间找代码调试，看了看，展示效果比单纯的检测要更加的炸裂一点。

2. 储备知识

        DeepSORT（Deep Learning to Track Multi-Object in SORT）是一种基于深度学习的多目标跟踪算法，它结合了深度学习的目标检测和传统的轨迹跟踪方法，旨在实现在复杂场景中准确和稳定地跟踪多个移动目标。以下是关于DeepSORT的检测思想、特点和应用方面的介绍：

        检测思想： DeepSORT的核心思想是结合深度学习目标检测和轨迹跟踪方法，以实现多目标跟踪。首先，利用深度学习目标检测模型（如YOLO、Faster R-CNN等）检测出每一帧图像中的所有目标物体，并提取其特征。然后，通过应用传统的轨迹跟踪算法（如卡尔曼滤波器和轨迹关联等），将目标在连续帧之间进行关联，从而生成每个目标的运动轨迹。

        特点：

  1. 多目标跟踪： DeepSORT专注于同时跟踪多个目标，适用于需要同时监测和追踪多个物体的场景，如交通监控、人群管理等。
  2. 深度特征： 通过使用深度学习模型提取目标的特征，DeepSORT可以更准确地表示目标，从而提高跟踪的精度和鲁棒性。
  3. 轨迹关联： DeepSORT使用传统的轨迹关联技术来连接不同帧之间的目标，确保在物体出现、消失、重叠等情况下仍能准确跟踪。
  4. 实时性能： DeepSORT设计用于实时应用，可以在视频流中高效地进行目标跟踪，适用于要求实时性能的应用场景。

需要了解的算法内容：详细介绍

  1. 目前主流的目标跟踪算法都是基于Tracking-by-Detecton策略，即基于目标检测的结果来进行目标跟踪。DeepSORT运用的就是这个策略，上面的视频是DeepSORT对人群进行跟踪的结果，每个bbox左上角的数字是用来标识某个人的唯一ID号。

  2. 这里就有个问题，视频中不同时刻的同一个人，位置发生了变化，那么是如何关联上的呢？答案就是匈牙利算法和卡尔曼滤波。

    匈牙利算法可以告诉我们当前帧的某个目标，是否与前一帧的某个目标相同。卡尔曼滤波可以基于目标前一时刻的位置，来预测当前时刻的位置，并且可以比传感器（在目标跟踪中即目标检测器，比如Yolo等）更准确的估计目标的位置。

3. 准备工作

        基础代码:黄老师的github,参考的是这位博主的，我做了相应的修改

4. 代码修改的地方

具体需要修改的有两个py文件

(1) main.py文件，里面的检测器yolo用onnx做推理，onnx模型参考我的博文yolov5转rknn（聪明的你应该会的）

  1. import cv2
  2. import torch
  3. import numpy as np
  4. import onnxruntime as rt
  5.  
  6.  
  7. def sigmoid(x):
  8.     return 1 / (1 + np.exp(-x))
  9.  
  10.  
  11. def nms_boxes(boxes, scores):
  12.     """Suppress non-maximal boxes.
  14.     # Arguments
  15.         boxes: ndarray, boxes of objects.
  16.         scores: ndarray, scores of objects.
  18.     # Returns
  19.         keep: ndarray, index of effective boxes.
  20.     """
  21.     x = boxes[:, 0]
  22.     y = boxes[:, 1]
  23.     w = boxes[:, 2] - boxes[:, 0]
  24.     h = boxes[:, 3] - boxes[:, 1]
  25.  
  26.     areas = w * h
  27.     order = scores.argsort()[::-1]
  28.  
  29.     keep = []
  30.     while order.size > 0:
  31.         i = order[0]
  32.         keep.append(i)
  33.  
  34.         xx1 = np.maximum(x[i], x[order[1:]])
  35.         yy1 = np.maximum(y[i], y[order[1:]])
  36.         xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]])
  37.         yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]])
  38.  
  39.         w1 = np.maximum(0.0, xx2 - xx1 + 0.00001)
  40.         h1 = np.maximum(0.0, yy2 - yy1 + 0.00001)
  41.         inter = w1 * h1
  42.  
  43.         ovr = inter / (areas[i] + areas[order[1:]] - inter)
  44.         inds = np.where(ovr <= 0.45)[0]
  45.         order = order[inds + 1]
  46.     keep = np.array(keep)
  47.     return keep
  48.  
  49.  
  50. def process(input, mask, anchors):
  51.  
  52.     anchors = [anchors[i] for i in mask]
  53.     grid_h, grid_w = map(int, input.shape[0:2])
  54.  
  55.     box_confidence = sigmoid(input[..., 4])
  56.     box_confidence = np.expand_dims(box_confidence, axis=-1)
  57.  
  58.     box_class_probs = sigmoid(input[..., 5:])
  59.  
  60.     box_xy = sigmoid(input[..., :2])*2 - 0.5
  61.  
  62.     col = np.tile(np.arange(0, grid_w), grid_w).reshape(-1, grid_w)
  63.     row = np.tile(np.arange(0, grid_h).reshape(-1, 1), grid_h)
  64.     col = col.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
  65.     row = row.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
  66.     grid = np.concatenate((col, row), axis=-1)
  67.     box_xy += grid
  68.     box_xy *= int(img_size/grid_h)
  69.  
  70.     box_wh = pow(sigmoid(input[..., 2:4])*2, 2)
  71.     box_wh = box_wh * anchors
  72.  
  73.     box = np.concatenate((box_xy, box_wh), axis=-1)
  74.  
  75.     return box, box_confidence, box_class_probs
  76.  
  77.  
  78. def filter_boxes(boxes, box_confidences, box_class_probs):
  79.     """Filter boxes with box threshold. It's a bit different with origin yolov5 post process!
  81.     # Arguments
  82.         boxes: ndarray, boxes of objects.
  83.         box_confidences: ndarray, confidences of objects.
  84.         box_class_probs: ndarray, class_probs of objects.
  86.     # Returns
  87.         boxes: ndarray, filtered boxes.
  88.         classes: ndarray, classes for boxes.
  89.         scores: ndarray, scores for boxes.
  90.     """
  91.     box_classes = np.argmax(box_class_probs, axis=-1)
  92.     box_class_scores = np.max(box_class_probs, axis=-1)
  93.     pos = np.where(box_confidences[..., 0] >= 0.5)
  94.  
  95.  
  96.     boxes = boxes[pos]
  97.     classes = box_classes[pos]
  98.     scores = box_class_scores[pos]
  99.  
  100.     return boxes, classes, scores
  101.  
  102.  
  103. def yolov5_post_process(input_data):
  104.     masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
  105.     anchors = [[10, 13], [16, 30], [33, 23], [30, 61], [62, 45],
  106.               [59, 119], [116, 90], [156, 198], [373, 326]]
  107.  
  108.     boxes, classes, scores = [], [], []
  109.     for input,mask in zip(input_data, masks):
  110.         b, c, s = process(input, mask, anchors)
  111.         b, c, s = filter_boxes(b, c, s)
  112.         boxes.append(b)
  113.         classes.append(c)
  114.         scores.append(s)
  115.  
  116.     boxes = np.concatenate(boxes)
  117.     boxes = xywh2xyxy(boxes)
  118.     classes = np.concatenate(classes)
  119.     scores = np.concatenate(scores)
  120.  
  121.     nboxes, nclasses, nscores = [], [], []
  122.     for c in set(classes):
  123.         inds = np.where(classes == c)
  124.         b = boxes[inds]
  125.         c = classes[inds]
  126.         s = scores[inds]
  127.  
  128.         keep = nms_boxes(b, s)
  129.  
  130.         nboxes.append(b[keep])
  131.         nclasses.append(c[keep])
  132.         nscores.append(s[keep])
  133.  
  134.     if not nclasses and not nscores:
  135.         return None, None, None
  136.  
  137.     boxes = np.concatenate(nboxes)
  138.     classes = np.concatenate(nclasses)
  139.     scores = np.concatenate(nscores)
  140.  
  141.     return boxes, classes, scores
  142.  
  143.  
  144. def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
  145.     # Resize and pad image while meeting stride-multiple constraints
  146.     shape = img.shape[:2]  # current shape [height, width]
  147.     if isinstance(new_shape, int):
  148.         new_shape = (new_shape, new_shape)
  149.  
  150.     # Scale ratio (new / old)
  151.     r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
  152.     if not scaleup:  # only scale down, do not scale up (for better test mAP)
  153.         r = min(r, 1.0)
  154.  
  155.     # Compute padding
  156.     ratio = r, r  # width, height ratios
  157.     new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
  158.     dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
  159.     if auto:  # minimum rectangle
  160.         dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
  161.     elif scaleFill:  # stretch
  162.         dw, dh = 0.0, 0.0
  163.         new_unpad = (new_shape[1], new_shape[0])
  164.         ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios
  165.  
  166.     dw /= 2  # divide padding into 2 sides
  167.     dh /= 2
  168.  
  169.     if shape[::-1] != new_unpad:  # resize
  170.         img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
  171.     top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
  172.     left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
  173.     img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
  174.     return img, ratio, (dw, dh)
  175.  
  176.  
  177. def clip_coords(boxes, img_shape):
  178.     # Clip bounding xyxy bounding boxes to image shape (height, width)
  179.     boxes[:, 0].clamp_(0, img_shape[1])  # x1
  180.     boxes[:, 1].clamp_(0, img_shape[0])  # y1
  181.     boxes[:, 2].clamp_(0, img_shape[1])  # x2
  182.     boxes[:, 3].clamp_(0, img_shape[0])  # y2
  183.  
  184.  
  185. def xywh2xyxy(x):
  186.     # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
  187.     y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
  188.     y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
  189.     y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
  190.     y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
  191.     y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
  192.     
  193.     return y
  194.  
  195.  
  196. CLASSES  = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
  197.         'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
  198.         'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
  199.         'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
  200.         'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
  201.         'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
  202.         'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink',
  203.         'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']
  204.  
  205.  
  206. def preprocess(img, img_size):
  207.     img0 = img.copy()
  208.     img = letterbox(img, new_shape=img_size)[0]
  209.     img = img[:, :, ::-1].transpose(2, 0, 1)
  210.     img = np.ascontiguousarray(img).astype(np.float32)
  211.     img = torch.from_numpy(img)
  212.     img /= 255.0
  213.     if img.ndimension() == 3:
  214.         img = img.unsqueeze(0)
  215.  
  216.     return img0, img
  217.  
  218.  
  219. def draw(image, boxes, scores, classes):
  220.     """Draw the boxes on the image.
  222.     # Argument:
  223.         image: original image.
  224.         boxes: ndarray, boxes of objects.
  225.         classes: ndarray, classes of objects.
  226.         scores: ndarray, scores of objects.
  227.         all_classes: all classes name.
  228.     """
  229.     for box, score, cl in zip(boxes, scores, classes):
  230.         top, left, right, bottom = box
  231.         # print('class: {}, score: {}'.format(CLASSES[cl], score))
  232.         # print('box coordinate left,top,right,down: [{}, {}, {}, {}]'.format(top, left, right, bottom))
  233.         top = int(top)
  234.         left = int(left)
  235.         right = int(right)
  236.         bottom = int(bottom)
  237.  
  238.         cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)
  239.         cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score),
  240.                     (top, left - 6),
  241.                     cv2.FONT_HERSHEY_SIMPLEX,
  242.                     0.6, (0, 0, 255), 2)
  243.  
  244.  
  245. def detect(im, img_size, sess, input_name, outputs_name):
  246.     im0, img = preprocess(im, img_size)
  247.     input_data = onnx_inference(img.numpy(), sess, input_name, outputs_name)
  248.     boxes, classes, scores = yolov5_post_process(input_data)
  249.     
  250.     if boxes is not None:
  251.         draw(im, boxes, scores, classes)
  252.     cv2.imshow('demo', im)
  253.     cv2.waitKey(1)
  254.  
  255.  
  256. def onnx_inference(img, sess, input_name, outputs_name):
  257.     
  258.     # 模型推理:模型输出节点名，模型输入节点名，输入数据(注意节点名的格式！！！！！)
  259.     outputs = sess.run(outputs_name, {input_name: img})
  260.     input0_data = outputs[0]
  261.     input1_data = outputs[1]
  262.     input2_data = outputs[2]
  263.  
  264.     input0_data = input0_data.reshape([3, 80, 80, 85])
  265.     input1_data = input1_data.reshape([3, 40, 40, 85])
  266.     input2_data = input2_data.reshape([3, 20, 20, 85])
  267.  
  268.     input_data = list()
  269.     input_data.append(np.transpose(input0_data, (1, 2, 0, 3)))
  270.     input_data.append(np.transpose(input1_data, (1, 2, 0, 3)))
  271.     input_data.append(np.transpose(input2_data, (1, 2, 0, 3)))
  272.  
  273.     return input_data
  274.  
  275.  
  276. def load_onnx_model():
  277.     # onnx模型前向推理
  278.     sess = rt.InferenceSession('./weights/modified_yolov5s.onnx')
  279.     # 模型的输入和输出节点名，可以通过netron查看
  280.     input_name = 'images'
  281.     outputs_name = ['396', '440', '484']
  282.     
  283.     return sess, input_name, outputs_name
  284.  
  285.  
  286. if __name__ == '__main__':
  287.     # create onnx_model
  288.     sess, input_name, outputs_name = load_onnx_model()
  289.     # input_model_size
  290.     img_size = 640
  291.     # read video
  292.     video = cv2.VideoCapture('./video/cut3.avi')
  293.     print("Loaded video ...")
  294.     frame_interval = 2  # 间隔帧数，例如每隔10帧获取一次
  295.     frame_count = 0
  296.     while True:
  297.         # 读取每帧图片
  298.         _, im = video.read()
  299.         if frame_count % frame_interval == 0:
  300.             if im is None:
  301.                 break
  302.             # 缩小尺寸，1920x1080->960x540
  303.             im = cv2.resize(im, (640, 640))
  304.             list_bboxs = []
  305.             # det_object
  306.             detect(im, img_size, sess, input_name, outputs_name)
  307.         frame_count += 1
  308.     video.release()
  309.     cv2.destroyAllWindows()
  310.  
  311.  
  312.

(2) feature_extractor.py的修改:

这里有4种推理情况：ckpt.t7是ReID( Re-identification利用算法),在图像库中找到要搜索的目标的技术，所以它是属于图像检索的一个子问题。

        (1) 动态的batch_size推理:由于检测到的目标是多个object，在本项目的代码REID推理中，会将目标通过torch.cat连接起来，变成(n, 64, 128)的形状，所以需要用动态的onnx模型

        (2)那我就想要静态的怎么办，安排！！！，思路就是将cat的拆分开就行了，shape变成(1, 64 , 128)，单个推理后将结果cat起来就行了，easy的。

重要！！！！ckpt文件转onnx的代码

  1. import os
  2. import cv2
  3. import time
  4. import argparse
  5. import torch
  6. import numpy as np
  7. from deep_sort import build_tracker
  8. from utils.draw import draw_boxes
  9. from utils.parser import get_config
  10. from tqdm import tqdm
  11.  
  12.  
  13. if __name__ == '__main__':
  14.     parser = argparse.ArgumentParser()
  15.     parser.add_argument("--config_deepsort", type=str, default="./configs/deep_sort.yaml", help='Configure tracker')
  16.     parser.add_argument("--cpu", dest="use_cuda", action="store_false", default=True, help='Run in CPU')
  17.     args = parser.parse_args()
  18.  
  19.     cfg = get_config()
  20.     cfg.merge_from_file(args.config_deepsort)
  21.     use_cuda = args.use_cuda and torch.cuda.is_available()
  22.     torch.set_grad_enabled(False)
  23.     model = build_tracker(cfg, use_cuda=False)
  24.  
  25.     model.reid = True
  26.     model.extractor.net.eval()
  27.  
  28.     device = 'cpu'
  29.     output_onnx = 'deepsort.onnx'
  30.     # ------------------------ export -----------------------------
  31.     print("==> Exporting model to ONNX format at '{}'".format(output_onnx))
  32.     input_names = ['input']
  33.     output_names = ['output']
  34.  
  35.     input_tensor = torch.randn(1, 3, 128, 64, device=device)
  36.  
  37.     torch.onnx.export(model.extractor.net, input_tensor, output_onnx, export_params=True, verbose=False,
  38.                       input_names=input_names, output_names=output_names, opset_version=13,
  39.                       do_constant_folding=True)

        (3)但是要转rknn怎么办，ckpt.t7转onnx后，有一个ReduceL2,不支持量化，我就转的fp16(在RK3588上是可以的，rk1808不知道行不行)，不过我尝试了将最后两个节点删除，对结果好像没有什么影响(用的是cut后的onnx推理)，有懂的朋友可以解释一下！！！

        (4) 就是rknn的推理，这里就不展示了，需要的私聊我吧

  1. import torch
  2. import torchvision.transforms as transforms
  3. import numpy as np
  4. import cv2
  5. # import onnxruntime as rt
  6. # from rknnlite.api import RKNNLite
  7.  
  8.  
  9. class Extractor(object):
  10.     def __init__(self, model_path):
  11.         self.model_path = model_path
  12.         self.device = "cpu"
  13.         self.size = (64, 128)
  14.         self.norm = transforms.Compose([
  15.             transforms.ToTensor(),
  16.             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
  17.         ])
  18.     
  19.     def _preprocess(self, im_crops):
  20.         """
  21.         TODO:
  22.             1. to float with scale from 0 to 1
  23.             2. resize to (64, 128) as Market1501 dataset did
  24.             3. concatenate to a numpy array
  25.             3. to torch Tensor
  26.             4. normalize
  27.         """
  28.         
  29.         def _resize(im, size):
  30.             return cv2.resize(im.astype(np.float32) / 255., size)
  31.         
  32.         im_batch = torch.cat([self.norm(_resize(im, self.size)).unsqueeze(0) for im in im_crops], dim=0).float()
  33.         
  34.         return im_batch
  35.     
  36.     def __call__(self, im_crops):
  37.         im_batch = self._preprocess(im_crops)
  38.  
  39.         # sess = rt.InferenceSession(self.model_path)
  40.         # 模型的输入和输出节点名，可以通过netron查看
  41.         # input_name = 'input'
  42.         # outputs_name = ['output']
  43.         
  44.         # (1)动态输出
  45.         # features = sess.run(outputs_name, {input_name: im_batch.numpy()})
  46.         # print('features:', np.array(features)[0, :, :].shape)
  47.         # return np.array(features)[0, :, :]
  48.  
  49.         
  50.         # (2)静态态输出
  51.         # sort_results = []
  52.         # n = im_batch.numpy().shape[0]
  53.         # for i in range(n):
  54.         #     img = im_batch.numpy()[i, :, :].reshape(1, 3, 128, 64)
  55.         #     feature = sess.run(outputs_name, {input_name: img})
  56.         #     feature = np.array(feature)
  57.         #     sort_results.append(feature)
  58.         # features = np.concatenate(sort_results, axis=1)[0, :, :]
  59.         # print(features.shape)
  60.         # return np.array(features)
  61.  
  62.  
  63.         # (3)去掉onnx的最后两个节点的静态模型输出
  64.         # input_name = 'input'
  65.         # outputs_name = ['204']
  66.         # sort_results = []
  67.         # n = im_batch.numpy().shape[0]
  68.         # for i in range(n):
  69.         #     img = im_batch.numpy()[i, :, :].reshape(1, 3, 128, 64)
  70.         #     feature = sess.run(outputs_name, {input_name: img})
  71.         #     feature = np.array(feature)
  72.         #     sort_results.append(feature)
  73.         # features = np.concatenate(sort_results, axis=1)[0, :, :]
  74.         # print(features.shape)
  75.         # return np.array(features)
  76.         
  77.         # (4 )rk模型修改
  78.         # rknn_lite = RKNNLite()
  79.         # rknn_lite.load_rknn('./weights/ckpt_fp16.rknn')
  80.         # ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0_1_2)
  81.         # if ret != 0:
  82.         #     print('Init runtime environment failed')
  83.         #     exit(ret)
  84.         # print('done')
  85.         
  86.         # sort_results = []
  87.         # n = im_batch.numpy().shape[0]
  88.         # for i in range(n):
  89.         #    img = im_batch.numpy()[i, :, :].reshape(1, 3, 128, 64)
  90.         #    feature = self.model_path.inference(inputs=[img])
  91.         #   feature = np.array(feature)
  92.         #    sort_results.append(feature)
  93.         # features = np.concatenate(sort_results, axis=1)[0, :, :]
  94.         # print(features.shape)
  95.         # return np.array(features)
  96.  
  97.

5.结果展示

        onnx的转换结果(测试视频地址)

 

检测结果

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/酷酷是懒虫/article/detail/737723
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号