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

[AI智能摄像头]RV1126部署yolov5并加速

作者：Gausst松鼠会 | 2024-05-17 19:19:21

赞

踩

[AI智能摄像头]RV1126部署yolov5并加速

导出onnx模型

yolov5官方地址 git clone https://github.com/ultralytics/yolov5

利用官方命令导出python export.py --weights yolov5n.pt --include onnx

利用代码导出

  1. import os
  2. import sys
  3. os.chdir(sys.path[0])
  4. import onnx
  5. import torch
  6. sys.path.append('..')
  7. from models.common import DetectMultiBackend
  8. from models.experimental import attempt_load
  9. DEVICE='cuda' if torch.cuda.is_available else 'cpu'
  10. def main():
  11.     """create model """
  12.     input = torch.randn(1, 3, 640, 640, requires_grad=False).float().to(torch.device(DEVICE))
  13.     model = attempt_load('./model/yolov5n.pt', device=DEVICE, inplace=True, fuse=True)  # load FP32 model
  14.     #model = DetectMultiBackend('./model/yolov5n.pt', data=input)
  15.     model.to(DEVICE)
  16.  
  17.     torch.onnx.export(model,
  18.             input,
  19.             'yolov5n_self.onnx', # name of the exported onnx model
  20.             export_params=True,
  21.             opset_version=12,
  22.             do_constant_folding=False, 
  23.             input_names=["images"])
  24. if __name__=="__main__":
  25.     main()

onnx模型测试

  1. import os
  2. import sys
  3. os.chdir(sys.path[0])
  4. import onnxruntime
  5. import torch
  6. import torchvision
  7. import numpy as np
  8. import time
  9. import cv2
  10. sys.path.append('..')
  11. from ultralytics.utils.plotting import Annotator, colors
  12.  
  13. ONNX_MODEL="./yolov5n.onnx"
  14. DEVICE='cuda' if torch.cuda.is_available() else 'cpu'
  15.  
  16. def xywh2xyxy(x):
  17.     """Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right."""
  18.     y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
  19.     y[..., 0] = x[..., 0] - x[..., 2] / 2  # top left x
  20.     y[..., 1] = x[..., 1] - x[..., 3] / 2  # top left y
  21.     y[..., 2] = x[..., 0] + x[..., 2] / 2  # bottom right x
  22.     y[..., 3] = x[..., 1] + x[..., 3] / 2  # bottom right y
  23.     return y
  24.  
  25. def box_iou(box1, box2, eps=1e-7):
  26.     # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
  27.     """
  28.     Return intersection-over-union (Jaccard index) of boxes.
  30.     Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
  31.     Arguments:
  32.         box1 (Tensor[N, 4])
  33.         box2 (Tensor[M, 4])
  34.     Returns:
  35.         iou (Tensor[N, M]): the NxM matrix containing the pairwise
  36.             IoU values for every element in boxes1 and boxes2
  37.     """
  38.  
  39.     # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
  40.     (a1, a2), (b1, b2) = box1.unsqueeze(1).chunk(2, 2), box2.unsqueeze(0).chunk(2, 2)
  41.     inter = (torch.min(a2, b2) - torch.max(a1, b1)).clamp(0).prod(2)
  42.  
  43.     # IoU = inter / (area1 + area2 - inter)
  44.     return inter / ((a2 - a1).prod(2) + (b2 - b1).prod(2) - inter + eps)
  45.  
  46. def non_max_suppression(
  47.     prediction,
  48.     conf_thres=0.25,
  49.     iou_thres=0.45,
  50.     classes=None,
  51.     agnostic=False,
  52.     multi_label=False,
  53.     labels=(),
  54.     max_det=300,
  55.     nm=0,  # number of masks
  56. ):
  57.     """
  58.     Non-Maximum Suppression (NMS) on inference results to reject overlapping detections.
  60.     Returns:
  61.          list of detections, on (n,6) tensor per image [xyxy, conf, cls]
  62.     """
  63.  
  64.     # Checks
  65.     assert 0 <= conf_thres <= 1, f"Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0"
  66.     assert 0 <= iou_thres <= 1, f"Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0"
  67.  
  68.     device = prediction.device
  69.     mps = "mps" in device.type  # Apple MPS
  70.     if mps:  # MPS not fully supported yet, convert tensors to CPU before NMS
  71.         prediction = prediction.cpu()
  72.     bs = prediction.shape[0]  # batch size
  73.     nc = prediction.shape[2] - nm - 5  # number of classes
  74.     xc = prediction[..., 4] > conf_thres  # candidates
  75.  
  76.     # Settings
  77.     # min_wh = 2  # (pixels) minimum box width and height
  78.     max_wh = 7680  # (pixels) maximum box width and height
  79.     max_nms = 30000  # maximum number of boxes into torchvision.ops.nms()
  80.     time_limit = 0.5 + 0.05 * bs  # seconds to quit after
  81.     redundant = True  # require redundant detections
  82.     multi_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)
  83.     merge = False  # use merge-NMS
  84.  
  85.     t = time.time()
  86.     mi = 5 + nc  # mask start index
  87.     output = [torch.zeros((0, 6 + nm), device=prediction.device)] * bs
  88.     for xi, x in enumerate(prediction):  # image index, image inference
  89.         # Apply constraints
  90.         # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0  # width-height
  91.         x = x[xc[xi]]  # confidence
  92.  
  93.         # Cat apriori labels if autolabelling
  94.         if labels and len(labels[xi]):
  95.             lb = labels[xi]
  96.             v = torch.zeros((len(lb), nc + nm + 5), device=x.device)
  97.             v[:, :4] = lb[:, 1:5]  # box
  98.             v[:, 4] = 1.0  # conf
  99.             v[range(len(lb)), lb[:, 0].long() + 5] = 1.0  # cls
  100.             x = torch.cat((x, v), 0)
  101.  
  102.         # If none remain process next image
  103.         if not x.shape[0]:
  104.             continue
  105.  
  106.         # Compute conf
  107.         x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf
  108.  
  109.         # Box/Mask
  110.         box = xywh2xyxy(x[:, :4])  # center_x, center_y, width, height) to (x1, y1, x2, y2)
  111.         mask = x[:, mi:]  # zero columns if no masks
  112.  
  113.         # Detections matrix nx6 (xyxy, conf, cls)
  114.         if multi_label:
  115.             i, j = (x[:, 5:mi] > conf_thres).nonzero(as_tuple=False).T
  116.             x = torch.cat((box[i], x[i, 5 + j, None], j[:, None].float(), mask[i]), 1)
  117.         else:  # best class only
  118.             conf, j = x[:, 5:mi].max(1, keepdim=True)
  119.             x = torch.cat((box, conf, j.float(), mask), 1)[conf.view(-1) > conf_thres]
  120.  
  121.         # Filter by class
  122.         if classes is not None:
  123.             x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
  124.  
  125.         # Apply finite constraint
  126.         # if not torch.isfinite(x).all():
  127.         #     x = x[torch.isfinite(x).all(1)]
  128.  
  129.         # Check shape
  130.         n = x.shape[0]  # number of boxes
  131.         if not n:  # no boxes
  132.             continue
  133.         x = x[x[:, 4].argsort(descending=True)[:max_nms]]  # sort by confidence and remove excess boxes
  134.  
  135.         # Batched NMS
  136.         c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
  137.         boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
  138.         i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
  139.         i = i[:max_det]  # limit detections
  140.         if merge and (1 < n < 3e3):  # Merge NMS (boxes merged using weighted mean)
  141.             # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
  142.             iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
  143.             weights = iou * scores[None]  # box weights
  144.             x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
  145.             if redundant:
  146.                 i = i[iou.sum(1) > 1]  # require redundancy
  147.  
  148.         output[xi] = x[i]
  149.         if mps:
  150.             output[xi] = output[xi].to(device)
  151.         if (time.time() - t) > time_limit:
  152.             break  # time limit exceeded
  153.  
  154.     return output
  155.  
  156. def draw_bbox(image, result, color=(0, 0, 255), thickness=2):
  157.     # img_path = cv2.cvtColor(img_path, cv2.COLOR_BGR2RGB)
  158.     image = image.copy()
  159.     for point in result:
  160.         x1,y1,x2,y2=point
  161.         cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness)
  162.     return image
  163.  
  164. def main():
  165.     input=torch.load("input.pt").to('cpu')
  166.     input_array=np.array(input)
  167.     onnx_model = onnxruntime.InferenceSession(ONNX_MODEL)
  168.     input_name = onnx_model.get_inputs()[0].name
  169.     out = onnx_model.run(None, {input_name:input_array})
  170.     out_tensor = torch.tensor(out).to(DEVICE)
  171.     pred = non_max_suppression(out_tensor,0.25,0.45,classes=None,agnostic=False,max_det=1000)
  172.     # Process predictions
  173.     for i, det in enumerate(pred):  # per image
  174.         im0_=cv2.imread('../data/images/bus.jpg')
  175.         im0=im0_.reshape(1,3,640,640)
  176.         names=torch.load('name.pt')
  177.         annotator = Annotator(im0, line_width=3, example=str(names))
  178.         coord=[]
  179.         image=im0.reshape(640,640,3)
  180.         if len(det):
  181.             # Rescale boxes from img_size to im0 size
  182.             #det[:, :4] = scale_boxes(im0.shape[2:], det[:, :4], im0.shape).round()
  183.             # Write results
  184.             for *xyxy, conf, cls in reversed(det):
  185.                 # Add bbox to image
  186.                 c = int(cls)  # integer class
  187.                 label = f"{names[c]} {conf:.2f}"
  188.                 # 创建两个顶点坐标子数组，并将它们组合成一个列表``
  189.                 coord.append([int(xyxy[0].item()), int(xyxy[1].item()),int(xyxy[2].item()), int(xyxy[3].item())])
  190.         image=draw_bbox(image,coord)
  191.         # Stream results
  192.         save_success =cv2.imwrite('result.jpg', image)
  193.         print(f"save image end {save_success}")
  194.                
  195. if __name__=="__main__":
  196.     main()

测试结果

板端部署

环境准备

搭建好rknntoolkit以及rknpu环境

大致流程

模型转换

新建export_rknn.py用于将onnx模型转化为rknn模型

  1. import os
  2. import sys
  3. os.chdir(sys.path[0])
  4. import numpy as np
  5. import cv2
  6. from rknn.api import RKNN
  7. import torchvision
  8. import torch
  9. import time
  10.  
  11. ONNX_MODEL = './model/yolov5n.onnx'
  12. RKNN_MODEL = './model/yolov5n.rknn'
  13.  
  14. def main():
  15.     """Create RKNN object"""
  16.     rknn = RKNN()
  17.     if not os.path.exists(ONNX_MODEL):
  18.         print('model not exist')
  19.         exit(-1)
  20.         
  21.     """pre-process config"""
  22.     print('--> Config model')
  23.     rknn.config(reorder_channel='0 1 2',
  24.                 mean_values=[[0, 0, 0]],
  25.                 std_values=[[255, 255, 255]],
  26.                 optimization_level=0,
  27.                 target_platform = ['rv1126'],
  28.                 output_optimize=1,
  29.                 quantize_input_node=True)
  30.     print('done')
  31.     
  32.     """Load ONNX model"""
  33.     print('--> Loading model')
  34.     ret = rknn.load_onnx(model=ONNX_MODEL,
  35.                         inputs=['images'],
  36.                         input_size_list = [[3, 640, 640]],
  37.                         outputs=['output0'])
  38.     if ret != 0:
  39.         print('Load yolov5 failed!')
  40.         exit(ret)
  41.     print('done')
  42.  
  43.     """Build model"""
  44.     print('--> Building model')
  45.     #ret = rknn.build(do_quantization=True,dataset='./data/data.txt')
  46.     ret = rknn.build(do_quantization=False,pre_compile=True)
  47.     if ret != 0:
  48.         print('Build yolov5 failed!')
  49.         exit(ret)
  50.     print('done')
  51.  
  52.     """Export RKNN model"""
  53.     print('--> Export RKNN model')
  54.     ret = rknn.export_rknn(RKNN_MODEL)
  55.     if ret != 0:
  56.         print('Export yolov5rknn failed!')
  57.         exit(ret)
  58.     print('done')
  59.     
  60. if __name__=="__main__":
  61.     main()

新建test_rknn.py用于测试rknn模型

  1. import os
  2. import sys
  3. os.chdir(sys.path[0])
  4. import numpy as np
  5. import cv2
  6. from rknn.api import RKNN
  7. import torchvision
  8. import torch
  9. import time
  10.  
  11. RKNN_MODEL = './model/yolov5n.rknn'
  12. DATA='./data/bus.jpg'
  13.  
  14. def xywh2xyxy(x):
  15.     coord=[]
  16.     for x_ in x:
  17.         xl=x_[0]-x_[2]/2
  18.         yl=x_[1]-x_[3]/2
  19.         xr=x_[0]+x_[2]/2
  20.         yr=x_[1]+x_[3]/2
  21.         coord.append([xl,yl,xr,yr])
  22.     coord=torch.tensor(coord).to(x.device)
  23.     return coord
  24. def box_iou(box1, box2, eps=1e-7):
  25.     # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
  26.     (a1, a2), (b1, b2) = box1.unsqueeze(1).chunk(2, 2), box2.unsqueeze(0).chunk(2, 2)
  27.     inter = (torch.min(a2, b2) - torch.max(a1, b1)).clamp(0).prod(2)
  28.     # IoU = inter / (area1 + area2 - inter)
  29.     return inter / ((a2 - a1).prod(2) + (b2 - b1).prod(2) - inter + eps)
  30. def non_max_suppression(
  31.     prediction,
  32.     conf_thres=0.25,
  33.     iou_thres=0.45,
  34.     classes=None,
  35.     agnostic=False,
  36.     multi_label=False,
  37.     labels=(),
  38.     max_det=300,
  39.     nm=0,  # number of masks
  40. ):
  41.  
  42.     # Checks
  43.     assert 0 <= conf_thres <= 1, f"Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0"
  44.     assert 0 <= iou_thres <= 1, f"Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0"
  45.  
  46.     device = prediction.device
  47.     mps = "mps" in device.type  # Apple MPS
  48.     if mps:  # MPS not fully supported yet, convert tensors to CPU before NMS
  49.         prediction = prediction.cpu()
  50.     bs = prediction.shape[0]  # batch size
  51.     nc = prediction.shape[2] - nm - 5  # number of classes
  52.     xc = prediction[..., 4] > conf_thres  # candidates
  53.     count_true = torch.sum(xc.type(torch.int))
  54.  
  55.     # Settings
  56.     # min_wh = 2  # (pixels) minimum box width and height
  57.     max_wh = 7680  # (pixels) maximum box width and height
  58.     max_nms = 30000  # maximum number of boxes into torchvision.ops.nms()
  59.     time_limit = 0.5 + 0.05 * bs  # seconds to quit after
  60.     redundant = True  # require redundant detections
  61.     multi_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)
  62.     merge = False  # use merge-NMS
  63.  
  64.     t = time.time()
  65.     mi = 5 + nc  # mask start index
  66.     output = [torch.zeros((0, 6 + nm), device=prediction.device)] * bs
  67.     for xi, x in enumerate(prediction):  # image index, image inference
  68.         # Apply constraints
  69.         # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0  # width-height
  70.         x = x[xc[xi]]  # confidence
  71.  
  72.         # Cat apriori labels if autolabelling
  73.         if labels and len(labels[xi]):
  74.             lb = labels[xi]
  75.             v = torch.zeros((len(lb), nc + nm + 5), device=x.device)
  76.             v[:, :4] = lb[:, 1:5]  # box
  77.             v[:, 4] = 1.0  # conf
  78.             v[range(len(lb)), lb[:, 0].long() + 5] = 1.0  # cls
  79.             x = torch.cat((x, v), 0)
  80.  
  81.         # If none remain process next image
  82.         if not x.shape[0]:
  83.             continue
  84.  
  85.         # Compute conf
  86.         x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf
  87.  
  88.         # Box/Mask
  89.         box = xywh2xyxy(x[:, :4])  # center_x, center_y, width, height) to (x1, y1, x2, y2)
  90.         mask = x[:, mi:]  # zero columns if no masks
  91.  
  92.         # Detections matrix nx6 (xyxy, conf, cls)
  93.         if multi_label:
  94.             i, j = (x[:, 5:mi] > conf_thres).nonzero(as_tuple=False).T
  95.             x = torch.cat((box[i], x[i, 5 + j, None], j[:, None].float(), mask[i]), 1)
  96.         else:  # best class only
  97.             conf, j = x[:, 5:mi].max(1, keepdim=True)
  98.             x = torch.cat((box, conf, j.float(), mask), 1)[conf.view(-1) > conf_thres]
  99.  
  100.         # Filter by class
  101.         if classes is not None:
  102.             x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
  103.  
  104.         # Apply finite constraint
  105.         # if not torch.isfinite(x).all():
  106.         #     x = x[torch.isfinite(x).all(1)]
  107.  
  108.         # Check shape
  109.         n = x.shape[0]  # number of boxes
  110.         if not n:  # no boxes
  111.             continue
  112.         x = x[x[:, 4].argsort(descending=True)[:max_nms]]  # sort by confidence and remove excess boxes
  113.  
  114.         # Batched NMS
  115.         c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
  116.         boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
  117.         i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
  118.         i = i[:max_det]  # limit detections
  119.         if merge and (1 < n < 3e3):  # Merge NMS (boxes merged using weighted mean)
  120.             # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
  121.             iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
  122.             weights = iou * scores[None]  # box weights
  123.             x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
  124.             if redundant:
  125.                 i = i[iou.sum(1) > 1]  # require redundancy
  126.  
  127.         output[xi] = x[i]
  128.         if mps:
  129.             output[xi] = output[xi].to(device)
  130.         if (time.time() - t) > time_limit:
  131.             break  # time limit exceeded
  132.  
  133.     return output
  134. def draw_bbox(image, result, color=(0, 0, 255), thickness=2):
  135.     # img_path = cv2.cvtColor(img_path, cv2.COLOR_BGR2RGB)
  136.     image = image.copy()
  137.     for point in result:
  138.         x1,y1,x2,y2=point
  139.         cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness)
  140.     return image
  141.  
  142. def main():
  143.     # Create RKNN object
  144.     rknn = RKNN()
  145.     rknn.list_devices()
  146.     #load rknn model
  147.     ret = rknn.load_rknn(path=RKNN_MODEL)
  148.     if ret != 0:
  149.         print('load rknn failed')
  150.         exit(ret)
  151.     # init runtime environment
  152.     print('--> Init runtime environment')
  153.     ret = rknn.init_runtime(target='rv1126', device_id='86d4fdeb7f3af5b1',perf_debug=True,eval_mem=True)
  154.     if ret != 0:
  155.         print('Init runtime environment failed')
  156.         exit(ret)
  157.     print('done')
  158.     # Set inputs
  159.     image=cv2.imread('./data/bus.jpg')
  160.     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
  161.     # Inference
  162.     print('--> Running model')
  163.     outputs = rknn.inference(inputs=[image])
  164.     #post process
  165.     out_tensor = torch.tensor(outputs)
  166.     pred = non_max_suppression(out_tensor,0.25,0.45,classes=None,agnostic=False,max_det=1000)
  167.     # Process predictions
  168.     for i, det in enumerate(pred):  # per image
  169.         im0_=cv2.imread(DATA)
  170.         im0=im0_.reshape(1,3,640,640)
  171.         coord=[]
  172.         image=im0.reshape(640,640,3)
  173.         if len(det):
  174.             """Write results"""
  175.             for *xyxy, conf, cls in reversed(det):
  176.                 c = int(cls)  # integer class
  177.                 coord.append([int(xyxy[0].item()), int(xyxy[1].item()),int(xyxy[2].item()), int(xyxy[3].item())])
  178.                 print(f"[{coord[0][0]},{coord[0][1]},{coord[0][2]},{coord[0][3]}]:ID is {c}")
  179.         image=draw_bbox(image,coord)
  180.         # Stream results
  181.         save_success =cv2.imwrite('result.jpg', image)
  182.         print(f"save image end {save_success}")
  183.     
  184.     rknn.release()
  185.  
  186. if __name__=="__main__":
  187.     main()

板端cpp推理代码编写

拷贝一份template改名为yolov5，目录结构如下

前处理代码

  1. void PreProcess(cv::Mat *image)
  2. {
  3.     cv::cvtColor(*image, *image, cv::COLOR_BGR2RGB);
  4. }

后处理代码

1：输出维度为[1，25200，85],其中85的前四个为中心点的x，y以及框的宽和高，第五个为框的置信度，后面80个为类别的置信度（有80个类别）；

2：25200=(80∗80+40∗40+20∗20)∗3，stride为8、16、32，640/8=80,640/16=40,640/32=20

3：NMS删除冗余候选框

1：IOU交并比：检测两个框重叠程度=交集面积/并集面积

2：主要步骤：

  1. 首先筛选出大于阈值的所有候选框（>0.4）
  2. 接着针对每一个种类将候选框进行分类
  3. 找到第n个类别进行循环操作
  4. 先找到置信度最大的框，放到保留区
  5. 和候选区的其他框计算交并比（IOU），若大于iou阈值则删除
  6. 再从候选区找到第二大的候选框放到保留区
  7. 重复4操作，直至候选区没有框
  8. 重复3操作，直至所有类别
  1. float iou(Bbox box1, Bbox box2) {
  2.     /*  
  3.     iou=交并比
  4.     */
  5.     int x1 = max(box1.x, box2.x);
  6.     int y1 = max(box1.y, box2.y);
  7.     int x2 = min(box1.x + box1.w, box2.x + box2.w);
  8.     int y2 = min(box1.y + box1.h, box2.y + box2.h);
  9.     int w = max(0, x2 - x1);
  10.     int h = max(0, y2 - y1);
  11.     float over_area = w * h;
  12.     return over_area / (box1.w * box1.h + box2.w * box2.h - over_area);
  13. }
  14.  
  15. bool judge_in_lst(int index, vector<int> index_lst) {
  16.     //若index在列表index_lst中则返回true，否则返回false
  17.     if (index_lst.size() > 0) {
  18.         for (int i = 0; i < int(index_lst.size()); i++) {
  19.             if (index == index_lst.at(i)) {
  20.                 return true;
  21.             }
  22.         }
  23.     }
  24.     return false;
  25. }
  26.  
  27. int get_max_index(vector<Detection> pre_detection) {
  28.     //返回最大置信度值对应的索引值
  29.     int index;
  30.     float conf;
  31.     if (pre_detection.size() > 0) {
  32.         index = 0;
  33.         conf = pre_detection.at(0).conf;
  34.         for (int i = 0; i < int(pre_detection.size()); i++) {
  35.             if (conf < pre_detection.at(i).conf) {
  36.                 index = i;
  37.                 conf = pre_detection.at(i).conf;
  38.             }
  39.         }
  40.         return index;
  41.     }
  42.     else {
  43.         return -1;
  44.     }
  45. }
  46.  
  47. vector<int> nms(vector<Detection> pre_detection, float iou_thr)
  48. {
  49.     /*
  50.     返回需保存box的pre_detection对应位置索引值
  51.     */
  52.     int index;
  53.     vector<Detection> pre_detection_new;
  54.     //Detection det_best;
  55.     Bbox box_best, box;
  56.     float iou_value;
  57.     vector<int> keep_index;
  58.     vector<int> del_index;
  59.     bool keep_bool;
  60.     bool del_bool;
  61.  
  62.     if (pre_detection.size() > 0) {
  63.         pre_detection_new.clear();
  64.         // 循环将预测结果建立索引
  65.         for (int i = 0; i < int(pre_detection.size()); i++) {
  66.             pre_detection.at(i).index = i;
  67.             pre_detection_new.push_back(pre_detection.at(i));
  68.         }
  69.         //循环遍历获得保留box位置索引-相对输入pre_detection位置
  70.         while (pre_detection_new.size() > 0) {
  71.             index = get_max_index(pre_detection_new);
  72.             if (index >= 0) {
  73.                 keep_index.push_back(pre_detection_new.at(index).index); //保留索引位置
  74.  
  75.                 // 更新最佳保留box
  76.                 box_best.x = pre_detection_new.at(index).bbox[0];
  77.                 box_best.y = pre_detection_new.at(index).bbox[1];
  78.                 box_best.w = pre_detection_new.at(index).bbox[2];
  79.                 box_best.h = pre_detection_new.at(index).bbox[3];
  80.  
  81.                 for (int j = 0; j < int(pre_detection.size()); j++) {
  82.                     keep_bool = judge_in_lst(pre_detection.at(j).index, keep_index);
  83.                     del_bool = judge_in_lst(pre_detection.at(j).index, del_index);
  84.                     if ((!keep_bool) && (!del_bool)) { //不在keep_index与del_index才计算iou
  85.                         box.x = pre_detection.at(j).bbox[0];
  86.                         box.y = pre_detection.at(j).bbox[1];
  87.                         box.w = pre_detection.at(j).bbox[2];
  88.                         box.h = pre_detection.at(j).bbox[3];
  89.                         iou_value = iou(box_best, box);
  90.                         if (iou_value > iou_thr) {
  91.                             del_index.push_back(j); //记录大于阈值将删除对应的位置
  92.                         }
  93.                     }
  94.  
  95.                 }
  96.                 //更新pre_detection_new
  97.                 pre_detection_new.clear();
  98.                 for (int j = 0; j < int(pre_detection.size()); j++) {
  99.                     keep_bool = judge_in_lst(pre_detection.at(j).index, keep_index);
  100.                     del_bool = judge_in_lst(pre_detection.at(j).index, del_index);
  101.                     if ((!keep_bool) && (!del_bool)) {
  102.                         pre_detection_new.push_back(pre_detection.at(j));
  103.                     }
  104.                 }
  105.             }
  106.         }
  107.     }
  108.  
  109.     del_index.clear();
  110.     del_index.shrink_to_fit();
  111.     pre_detection_new.clear();
  112.     pre_detection_new.shrink_to_fit();
  113.  
  114.     return  keep_index;
  115.  
  116. }
  117.  
  118.  
  119. vector<Detection> PostProcess(float* prob,float conf_thr=0.3,float nms_thr=0.5)
  120. {
  121.     vector<Detection> pre_results;
  122.     vector<int> nms_keep_index;
  123.     vector<Detection> results;
  124.     bool keep_bool;
  125.     Detection pre_res;
  126.     float conf;
  127.     int tmp_idx;
  128.     float tmp_cls_score;
  129.     for (int i = 0; i < 25200; i++) {
  130.         tmp_idx = i * (CLSNUM + 5);
  131.         pre_res.bbox[0] = prob[tmp_idx + 0];  //cx
  132.         pre_res.bbox[1] = prob[tmp_idx + 1];  //cy
  133.         pre_res.bbox[2] = prob[tmp_idx + 2];  //w
  134.         pre_res.bbox[3] = prob[tmp_idx + 3];  //h
  135.         conf = prob[tmp_idx + 4];  // 是为目标的置信度
  136.         tmp_cls_score = prob[tmp_idx + 5] * conf; //conf_thr*nms_thr
  137.         pre_res.class_id = 0;
  138.         pre_res.conf = 0;
  139.         // 这个过程相当于从除了前面5列，在后面的cla_num个数据中找出score最大的值作为pre_res.conf，对应的列作为类id
  140.         for (int j = 1; j < CLSNUM; j++) {     
  141.             tmp_idx = i * (CLSNUM + 5) + 5 + j; //获得对应类别索引
  142.             if (tmp_cls_score < prob[tmp_idx] * conf){
  143.                 tmp_cls_score = prob[tmp_idx] * conf;
  144.                 pre_res.class_id = j;
  145.                 pre_res.conf = tmp_cls_score;
  146.             }
  147.         }
  148.         if (conf >= conf_thr) {
  149.             pre_results.push_back(pre_res);
  150.         }
  151.     }
  152.     //使用nms，返回对应结果的索引
  153.     nms_keep_index=nms(pre_results,nms_thr);
  154.     // 茛据nms找到的索引，将结果取出来作为最终结果
  155.     for (int i = 0; i < int(pre_results.size()); i++) {
  156.         keep_bool = judge_in_lst(i, nms_keep_index);
  157.         if (keep_bool) {
  158.             results.push_back(pre_results.at(i));
  159.         }
  160.     }
  161.  
  162.     pre_results.clear();
  163.     pre_results.shrink_to_fit();
  164.     nms_keep_index.clear();
  165.     nms_keep_index.shrink_to_fit();
  166.  
  167.     return results; 
  168. }

结果展示

至此板端部署结束，接下来进行优化；

优化加速

可以看到模型推理的时间近2s，对于实时处理来说是远远不够的，因此需要对模型进行加速

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

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

  

闽ICP备14008679号