• yolov7使用onnx推理

    官方代码：

    GitHub - WongKinYiu/yolov7: Implementation of paper - YOLOv7: Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors

            上个月，官方放了个使用onnx推理的ipynb文件，过了几天上去看，官方又给删了，不知道是不是要更新波大的，还好手快保存了一份，这个可以作为备忘，懒得再重新写（不过这得是多懒，都在这码字了不去写代码。。。）

    不带NMS 

            先贴代码吧：

    1. import cv2
    2. import time
    3. import requests
    4. import random
    5. import numpy as np
    6. import onnxruntime as ort
    7. from PIL import Image
    8. from pathlib import Path
    9. from collections import OrderedDict,namedtuple
    10.  
    11. cuda = False
    12. w = "yolov7.onnx"
    13. providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
    14. session = ort.InferenceSession(w, providers=providers)
    15.  
    16. def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleup=True, stride=32):
    17.     # Resize and pad image while meeting stride-multiple constraints
    18.     shape = im.shape[:2]  # current shape [height, width]
    19.     if isinstance(new_shape, int):
    20.         new_shape = (new_shape, new_shape)
    21.  
    22.     # Scale ratio (new / old)
    23.     r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    24.     if not scaleup:  # only scale down, do not scale up (for better val mAP)
    25.         r = min(r, 1.0)
    26.  
    27.     # Compute padding
    28.     new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    29.     dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    30.  
    31.     if auto:  # minimum rectangle
    32.         dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
    33.  
    34.     dw /= 2  # divide padding into 2 sides
    35.     dh /= 2
    36.  
    37.     if shape[::-1] != new_unpad:  # resize
    38.         im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    39.     top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    40.     left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    41.     im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    42.     return im, r, (dw, dh)
    43.  
    44. names = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
    45.          'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
    46.          'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
    47.          'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
    48.          'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
    49.          'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
    50.          'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
    51.          'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
    52.          'hair drier', 'toothbrush']
    53. colors = {name:[random.randint(0, 255) for _ in range(3)] for i,name in enumerate(names)}
    54.  
    55. # url = 'https://oneflow-static.oss-cn-beijing.aliyuncs.com/tripleMu/image1.jpg'
    56. # file = requests.get(url)
    57.  
    58. img_path = r'\inference\images\image3.jpg'
    59. # img = cv2.imdecode(np.frombuffer(file.content, np.uint8), 1)
    60. img = cv2.imread(img_path)
    61. # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    62.  
    63. image = img.copy()
    64. image, ratio, dwdh = letterbox(image, auto=False)
    65. image = image.transpose((2, 0, 1))
    66. image = np.expand_dims(image, 0)
    67. image = np.ascontiguousarray(image)
    68.  
    69. im = image.astype(np.float32)
    70. im /= 255
    71.  
    72.  
    73. outname = [i.name for i in session.get_outputs()]
    74. inname = [i.name for i in session.get_inputs()]
    75.  
    76. inp = {inname[0]:im}
    77. t1 = time.time()
    78. outputs = session.run(outname, inp)[0]
    79. print('inference time :%.4f'%(time.time()-t1))
    80. # print(outputs)
    81.  
    82. ori_images = [img.copy()]
    83.  
    84. for i,(batch_id,x0,y0,x1,y1,cls_id,score) in enumerate(outputs):
    85.     image = ori_images[int(batch_id)]
    86.     box = np.array([x0,y0,x1,y1])
    87.     box -= np.array(dwdh*2)
    88.     box /= ratio
    89.     box = box.round().astype(np.int32).tolist()
    90.     cls_id = int(cls_id)
    91.     score = round(float(score),3)
    92.     name = names[cls_id]
    93.     color = colors[name]
    94.     name += ' '+str(score)
    95.     cv2.rectangle(image,box[:2],box[2:],color,2)
    96.     cv2.putText(image,name,(box[0], box[1] - 2),cv2.FONT_HERSHEY_SIMPLEX,0.75,[225, 255, 255],thickness=2)
    97.  
    98. cv2.imshow('dddd',ori_images[0])
    99. cv2.waitKey(0)
    100. cv2.destroyAllWindows()
    101.  
    102. # Image.fromarray(ori_images[0])

            熟悉yolo系列的朋友应该看出上面的问题了，没有NMS，这是因为官方代码在导出onnx的时候做了简化和端到端的处理。

            导出指令：

    python export.py --weights yolov7.pt --grid --end2end --simplify \
        --topk-all 100 --iou-thres 0.65 --conf-thres 0.35 --img-size 640 640 --max-wh 640

            如果单纯运行export.py导出的onnx是运行不了上面的代码的，在for循环的时候会报错。

    在控制台执行以上指令：

            可以看到模型最后是导出成功的，过程会有些警告，忽视即可，不过简化失败，这个主要是应为onnx的版本问题，我的是1.9，降级到1.8.1之后就可以了。

            直接运行最初的onnx推理代码，不过需要注意的是图片的路径可能需要修改一下

    img_path = r'改成本地图片路径'

            推理结果：

     

    带NMS

            这里是带有nms的版本

            在yolov7的根目录下创建detect_onnx.py，并复制如下代码：

    1. import argparse
    2. import time
    3. from pathlib import Path
    4.  
    5. import cv2
    6. import torch
    7. import torch.backends.cudnn as cudnn
    8. from numpy import random
    9. from utils.datasets import LoadStreams, LoadImages
    10. from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
    11.     scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path
    12. from utils.plots import plot_one_box
    13. from utils.torch_utils import select_device, load_classifier, time_synchronized
    14. import onnxruntime
    15.  
    16. def detect(save_img=False):
    17.     source, weights, view_img, save_txt, imgsz, trace = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size, not opt.no_trace
    18.     save_img = not opt.nosave and not source.endswith('.txt')  # save inference images
    19.     webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
    20.         ('rtsp://', 'rtmp://', 'http://', 'https://'))
    21.  
    22.     # Directories
    23.     save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))  # increment run
    24.     (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
    25.  
    26.     # Initialize
    27.     set_logging()
    28.     device = select_device(opt.device)
    29.     half = device.type != 'cpu'  # half precision only supported on CUDA
    30.  
    31.     cuda = torch.cuda.is_available()
    32.     check_requirements(('onnx', 'onnxruntime-gpu' if cuda else 'onnxruntime'))
    33.     providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
    34.     session = onnxruntime.InferenceSession(weights, providers=providers)
    35.  
    36.     # Load model
    37.     model = session
    38.     # stride = int(model.stride.max())  # model stride
    39.     stride = 32
    40.     imgsz = check_img_size(imgsz, s=stride)  # check img_size
    41.  
    42.     # Second-stage classifier
    43.     classify = False
    44.     if classify:
    45.         modelc = load_classifier(name='resnet101', n=2)  # initialize
    46.         modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()
    47.  
    48.     # Set Dataloader
    49.     vid_path, vid_writer = None, None
    50.     if webcam:
    51.         view_img = check_imshow()
    52.         cudnn.benchmark = True  # set True to speed up constant image size inference
    53.         dataset = LoadStreams(source, img_size=imgsz, stride=stride)
    54.     else:
    55.         dataset = LoadImages(source, img_size=imgsz, stride=stride,auto=False)
    56.  
    57.     # Get names and colors
    58.     names = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
    59.          'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
    60.          'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
    61.          'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
    62.          'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
    63.          'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
    64.          'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
    65.          'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
    66.          'hair drier', 'toothbrush']
    67.     colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
    68.  
    69.     # Run inference
    70.     if device.type != 'cpu':
    71.         model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
    72.     t0 = time.time()
    73.     for path, img, im0s, vid_cap in dataset:
    74.         img = torch.from_numpy(img).to(device)
    75.         img = img.half() if half else img.float()  # uint8 to fp16/32
    76.         img /= 255.0  # 0 - 255 to 0.0 - 1.0
    77.         if img.ndimension() == 3:
    78.             img = img.unsqueeze(0)
    79.  
    80.         # Inference
    81.         t1 = time_synchronized()
    82.         ts = time.time()
    83.  
    84.         im = img.cpu().numpy()  # torch to numpy
    85.         pred = model.run([session.get_outputs()[0].name], {session.get_inputs()[0].name: im})[0]
    86.         # pred = model.run([i.name for i in session.get_outputs()], {session.get_inputs()[0].name: im})[0]
    87.         te = time.time()
    88.         print('inference time : %.4f s'%(te-ts))
    89.         # Apply NMS
    90.         pred = torch.from_numpy(pred).reshape(1,-1,85)
    91.         pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
    92.         t2 = time_synchronized()
    93.  
    94.         # Apply Classifier
    95.         if classify:
    96.             pred = apply_classifier(pred, modelc, img, im0s)
    97.  
    98.         # Process detections
    99.         for i, det in enumerate(pred):  # detections per image
    100.             if webcam:  # batch_size >= 1
    101.                 p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
    102.             else:
    103.                 p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
    104.  
    105.             p = Path(p)  # to Path
    106.             save_path = str(save_dir / p.name)  # img.jpg
    107.             txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # img.txt
    108.             s += '%gx%g ' % img.shape[2:]  # print string
    109.             gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
    110.             if len(det):
    111.                 # Rescale boxes from img_size to im0 size
    112.                 det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
    113.  
    114.                 # Print results
    115.                 for c in det[:, -1].unique():
    116.                     n = (det[:, -1] == c).sum()  # detections per class
    117.                     s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
    118.  
    119.                 # Write results
    120.                 for *xyxy, conf, cls in reversed(det):
    121.                     if save_txt:  # Write to file
    122.                         xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
    123.                         line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh)  # label format
    124.                         with open(txt_path + '.txt', 'a') as f:
    125.                             f.write(('%g ' * len(line)).rstrip() % line + '\n')
    126.  
    127.                     if save_img or view_img:  # Add bbox to image
    128.                         label = f'{names[int(cls)]} {conf:.2f}'
    129.                         plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
    130.  
    131.             # Print time (inference + NMS)
    132.             print(f'{s}Done. ({t2 - t1:.3f}s)')
    133.  
    134.             # Stream results
    135.             if view_img:
    136.                 cv2.imshow(str(p), im0)
    137.                 cv2.waitKey(0)  # 1 millisecond
    138.  
    139.             # Save results (image with detections)
    140.             if save_img:
    141.                 if dataset.mode == 'image':
    142.                     cv2.imwrite(save_path, im0)
    143.                     print(f" The image with the result is saved in: {save_path}")
    144.                 else:  # 'video' or 'stream'
    145.                     if vid_path != save_path:  # new video
    146.                         vid_path = save_path
    147.                         if isinstance(vid_writer, cv2.VideoWriter):
    148.                             vid_writer.release()  # release previous video writer
    149.                         if vid_cap:  # video
    150.                             fps = vid_cap.get(cv2.CAP_PROP_FPS)
    151.                             w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    152.                             h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    153.                         else:  # stream
    154.                             fps, w, h = 30, im0.shape[1], im0.shape[0]
    155.                             save_path += '.mp4'
    156.                         vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
    157.                     vid_writer.write(im0)
    158.  
    159.     if save_txt or save_img:
    160.         s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
    161.         #print(f"Results saved to {save_dir}{s}")
    162.  
    163.     print(f'Done. ({time.time() - t0:.3f}s)')
    164.  
    165.  
    166. if __name__ == '__main__':
    167.     parser = argparse.ArgumentParser()
    168.     parser.add_argument('--weights', nargs='+', type=str, default='yolov7.onnx', help='model.onnx path(s)')
    169.     parser.add_argument('--source', type=str, default='inference/images', help='source')  # file/folder, 0 for webcam
    170.     parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    171.     parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
    172.     parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
    173.     parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    174.     parser.add_argument('--view-img', action='store_true', help='display results')
    175.     parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    176.     parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    177.     parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
    178.     parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    179.     parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    180.     parser.add_argument('--augment', action='store_true', help='augmented inference')
    181.     parser.add_argument('--update', action='store_true', help='update all models')
    182.     parser.add_argument('--project', default='runs/detect', help='save results to project/name')
    183.     parser.add_argument('--name', default='exp', help='save results to project/name')
    184.     parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    185.     parser.add_argument('--no-trace', action='store_true', help='don`t trace model')
    186.     opt = parser.parse_args()
    187.     print(opt)
    188.     #check_requirements(exclude=('pycocotools', 'thop'))
    189.  
    190.     with torch.no_grad():
    191.         if opt.update:  # update all models (to fix SourceChangeWarning)
    192.             for opt.weights in ['yolov7.pt']:
    193.                 detect()
    194.                 strip_optimizer(opt.weights)
    195.         else:
    196.             detect()

    这里有两点要注意：

    1、dataset = LoadImages(source, img_size=imgsz, stride=stride,auto=False)
需要加auto参数，并设置为False，主要是为了把参数传到以下函数里，不然自动padding的时候图像大小可能会变化，导致报错
    img = letterbox(img0, self.img_size, stride=self.stride,auto=self.auto)[0]

    2、导出onnx时的语句：python export.py --weights yolov7.pt --grid --img-size 640 640

    推理结果：

    ​​​​​​​ 

     

     

  • 相关阅读:
    Java基础进阶-序列化
    优思学院｜看板方式与传统生产方式的对比
    N个元素进栈 出栈情况种数
    [c++刷题]贪心算法.N01
    JS面向对象---原型链继承
    想要拿到手软的大厂offer必须要刷5遍这份5000页Java 最全技术栈手册
    【校招VIP】产品深度理解之热点事件分析
    【沐风老师】3DMAX一键生成圣诞树建模插件使用教程
    【Linux】操作系统安装详解
    （尊享版）22年国内最牛的Java面试八股文合集，不接受反驳
  • 原文地址：https://blog.csdn.net/athrunsunny/article/details/126306363