├── README.md ├── cf_v6.onnx ├── imags_14.jpg ├── myToolkit.py └── onnxruntime_例子.py /README.md: -------------------------------------------------------------------------------- 1 | ## yolov5 onnxruntime框架示例源码，如果安装 onnxruntime-gpu则使用gpu推理， 安装方法： pip install onnxruntime-gpu 2 | -------------------------------------------------------------------------------- /cf_v6.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AMXZzzz/yolov5-onnxruntime/df456eb72a27d4475404202d21de5cdba756ef3b/cf_v6.onnx -------------------------------------------------------------------------------- /imags_14.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AMXZzzz/yolov5-onnxruntime/df456eb72a27d4475404202d21de5cdba756ef3b/imags_14.jpg -------------------------------------------------------------------------------- /myToolkit.py: -------------------------------------------------------------------------------- 1 | # --------------------- 2 | # Bilibili: 随风而息 3 | # Time：2022/5/25 23:20 4 | 5 | import torchvision 6 | import cv2 7 | import numpy as np 8 | import torch 9 | 10 | ''' 11 | 工具包,所有的预处理和后处理函数都在这 12 | ''' 13 | 14 | '''坐标转换''' 15 | # 中心点xy,宽高wh转换到左上角xy和右下角xy 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 | 26 | # 左上角xy,右下角xy转换至左上角xy,宽高wh 27 | def xyxy2xywh(x): 28 | # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right 29 | y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x) 30 | y[:, 0] = (x[:, 0] + x[:, 2]) / 2 # x center 31 | y[:, 1] = (x[:, 1] + x[:, 3]) / 2 # y center 32 | y[:, 2] = x[:, 2] - x[:, 0] # width 33 | y[:, 3] = x[:, 3] - x[:, 1] # height 34 | return y 35 | 36 | '''--------预处理--------''' 37 | def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32): 38 | # 在满足跨步约束的同时调整图像大小和填充图像 39 | shape = im.shape[:2] # 当前形状[高度，宽度] #注意顺序 40 | # print(f'图片尺寸:{shape}') 41 | if isinstance(new_shape, int): 42 | new_shape = (new_shape, new_shape) # new_shape:(640, 640) 43 | 44 | # 计算图片尺寸的比例 45 | r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) # 找出图片的(高宽)最小的比例，[0]是高，[1]是宽 46 | if not scaleup: # 只缩小，不放大（为了更好的 val mAP）默认跳过 47 | r = min(r, 1.0) # # 若有大于1的则用1比例，若有小于1的则选最小，更新r 48 | 49 | # 计算填充边缘 50 | ratio = r, r # 高宽比，用上面计算的最小r作为宽高比 51 | new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) # round:四舍五入,宽高,注意顺序 new_unpad:(640, 361) 52 | # print(f'按比例需要缩放到:{new_unpad}') 53 | dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding (640-640) ,(640-361) 54 | # print(f'填充的大小,dw:{dw},dh:{dh}') 55 | if auto: # 最小矩形,为False 56 | dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding # mod:计算两数组对应位置元素的余数。 57 | # print(f'最小矩形dw:{dw},dh:{dh}') 58 | elif scaleFill: # 缩放，一般为False 59 | dw, dh = 0.0, 0.0 60 | new_unpad = (new_shape[1], new_shape[0]) 61 | ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios 62 | dw /= 2 # 将填充分为 2 条边 63 | dh /= 2 64 | if shape[::-1] != new_unpad: # 裁剪 shape[::-1]:(1376, 776) new_unpad:(640, 361) 65 | im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR) 66 | top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) # 防止过填充 67 | left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) 68 | im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border 69 | # print(f'填充后的图片尺寸:{im.shape}') 70 | return im, ratio, (dw, dh) 71 | 72 | '''------后处理------''' 73 | # NMS 74 | def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False, 75 | max_det=300): 76 | """ 77 | 78 | 返回：检测列表，每个图像的 (n,6) tensor [xyxy, conf, cls] # [左上角坐标xy右下角坐标xy,置信度,类别] 79 | """ 80 | 81 | nc = prediction.shape[2]- 5 # 类别数量 82 | # print(prediction.shape) 83 | xc = prediction[..., 4] > conf_thres # 候选框 84 | 85 | # 设置 86 | min_wh, max_wh = 2, 4096 # （像素）最小和最大盒子宽度和高度 87 | max_nms = 3000 # torchvision.ops.nms() 中的最大框数 88 | multi_label &= nc > 1 # 每个候选框的多标签设置（增加 0.5ms/img） 89 | 90 | 91 | 92 | output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0] 93 | for xi, x in enumerate(prediction): # 图像索引xi，图像推断x # enumerate可遍历的数据对象组合为一个索引序列 94 | # 应用约束 95 | x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # 宽高 96 | x = x[xc[xi]] # 置信度 97 | 98 | # 如果没有数据就处理下一个图像 99 | if not x.shape[0]: 100 | continue 101 | 102 | # 计算配置 103 | x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf 104 | # 框（中心x，中心y，宽度，高度）到（x1，y1，x2，y2） 105 | box = xywh2xyxy(x[:, :4]) 106 | 107 | # 检测矩阵 nx6 (xyxy, conf, cls) 108 | if multi_label: 109 | i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T 110 | x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1) 111 | else: # 只有最好类 112 | conf, j = x[:, 5:].max(1, keepdim=True) 113 | x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres] 114 | 115 | # 按类别过滤 116 | if classes is not None: 117 | x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)] 118 | 119 | # Check shape 120 | n = x.shape[0] # 框的个数 121 | if not n: # 没有锚框 122 | continue 123 | elif n > max_nms: # 多余的锚框 124 | x = x[x[:, 4].argsort(descending=True)[:max_nms]] # 按置信度排序 125 | 126 | # 批量NMS 127 | c = x[:, 5:6] * (0 if agnostic else max_wh) # 类别 128 | boxes, scores = x[:, :4] + c, x[:, 4] # 框（类别偏移），分数 129 | i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS 130 | if i.shape[0] > max_det: # 极限检测 131 | i = i[:max_det] 132 | 133 | output[xi] = x[i] 134 | 135 | return output 136 | 137 | 138 | def clip_coords(boxes, shape): 139 | # 将边界 xyxy 框裁剪为图像形状（高度、宽度） 140 | if isinstance(boxes, torch.Tensor): # tensor类型 141 | boxes[:, 0].clamp_(0, shape[1]) # x1 142 | boxes[:, 1].clamp_(0, shape[0]) # y1 143 | boxes[:, 2].clamp_(0, shape[1]) # x2 144 | boxes[:, 3].clamp_(0, shape[0]) # y2 145 | else: # np.array 类型 146 | boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, shape[1]) # x1, x2 147 | boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, shape[0]) # y1, y2 148 | 149 | # 转换 150 | def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None): 151 | # 将坐标 (xyxy) 从 img_shape 重新缩放为 img0_shape 152 | if ratio_pad is None: # 从 img0_shape 计算 153 | gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # 比例 = 旧 / 新 154 | pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh 填充大小 155 | else: 156 | gain = ratio_pad[0][0] 157 | pad = ratio_pad[1] 158 | coords[:, [0, 2]] -= pad[0] # x 填充 159 | coords[:, [1, 3]] -= pad[1] # y 填充 160 | coords[:, :4] /= gain 161 | clip_coords(coords, img0_shape) 162 | return coords 163 | -------------------------------------------------------------------------------- /onnxruntime_例子.py: -------------------------------------------------------------------------------- 1 | # --------------------- 2 | # Bilibili: 随风而息 3 | # Time：2022/5/25 23:20 4 | #---------------------- 5 | 6 | import cv2 7 | import numpy as np 8 | import onnxruntime 9 | import torch 10 | import mss 11 | from myToolkit import letterbox, non_max_suppression, scale_coords 12 | 13 | ''' 14 | onnxruntime引擎例子,可以用cpu和gpu版本 15 | ''' 16 | 17 | cap = mss.mss() 18 | def grab_screen_mss(monitor): 19 | return cv2.cvtColor(np.array(cap.grab(monitor)), cv2.COLOR_BGRA2BGR) 20 | if __name__ == '__main__': 21 | # 参数定义 22 | img_size = (640, 640) # 训练权重的传入尺寸 23 | cuda = torch.cuda.is_available() 24 | device = 'cuda' if cuda else 'cpu' # 根据pytorch是否支持gpu选择设备 25 | 26 | conf_thres = 0.25 # 置信度 27 | iou_thres = 0.45 # iou 28 | 29 | # 加载onnx引擎 30 | onnx_path = f'cf_3000img.onnx' # onnx模型的路径 31 | providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider'] # 选择onnxruntime 32 | sess = onnxruntime.InferenceSession(onnx_path, providers=providers) # 加载模型 33 | 34 | # # 加载单张图片 35 | # img_path = 'imags_14.jpg' # 图片路径 36 | # img0 = cv2.imread(img_path) # cv2读取图片 37 | 38 | # ----截屏---- 39 | scr = mss.mss() # 实例化mss 40 | game_left, game_tap, game_x, game_y = 0, 0, 1920, 1080 # 截图范围 41 | monitor = { 42 | 'left': game_left, # 起始点 43 | 'top': game_tap, # 起始点 44 | 'width': game_x, # 长度 45 | 'height': game_y } # 高度 46 | cv2.namedWindow('img', cv2.WINDOW_NORMAL) # 创建窗口 47 | cv2.resizeWindow('img', 1920//3, 1080//3) # 裁剪窗口 48 | #--------------- 49 | 50 | while True: 51 | if not cv2.getWindowProperty('img', cv2.WND_PROP_VISIBLE): 52 | cv2.destroyAllWindows() 53 | exit('程序结束...') 54 | break 55 | # 截图 56 | img0 = grab_screen_mss(monitor) 57 | 58 | # 预处理 59 | img = letterbox(img0, img_size, stride=64, auto=False)[0] 60 | # 转tensor 61 | img = img.transpose((2, 0, 1))[::-1] # HWC 转 CHW，BGR 转 RGB 62 | # 返回一个连续的array，其内存是连续的。 63 | img = np.ascontiguousarray(img) 64 | # 放入设备 65 | img = torch.from_numpy(img).to(device) 66 | # uint8 转 fp32 67 | img = img.float() 68 | # 归一化 69 | img /= 255 70 | # 扩大批量调暗 71 | if len(img.shape): 72 | img = img[None] 73 | 74 | # 推理 75 | img = img.cpu().numpy() # 传入cpu并转成numpy格式 76 | pred_onnx = torch.tensor(sess.run([sess.get_outputs()[0].name], {sess.get_inputs()[0].name: img})[0]) 77 | 78 | # nms 79 | pred = non_max_suppression(pred_onnx, conf_thres, iou_thres, classes=None, agnostic=False, max_det=1000) 80 | 81 | # 转换 82 | aims = [] 83 | for i, det in enumerate(pred): 84 | if len(det): 85 | # 将坐标 (xyxy) 从 img_shape 重新缩放为 img0_shape 86 | det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round() 87 | for *xyxy, conf, cls in reversed(det): # 从末尾遍历 88 | # 将xyxy合并至一个维度,锚框的左上角和右下角 89 | xyxy = (torch.tensor(xyxy).view(1, 4)).view(-1) 90 | # 将类别和坐标合并 91 | line = (cls, *xyxy) 92 | # 提取tensor类型里的坐标数据 93 | aim = ('%g ' * len( 94 | line)).rstrip() % line # %g 格式为浮点数 .rstrip()删除tring字符串末尾的指定字符,默认为空白符包括空格,即删除2个坐标之间的空格 95 | # 划分元素 96 | aim = aim.split(' ') # 将一个元素按空格符分为多个元素,获得单个目标信息列表 97 | # 所有目标的类别和锚框的坐标(类别,左上角x,左上角y,右下角x,右下角y) 98 | aims.append(aim) # 添加至列表 99 | print(aims) 100 | 101 | # 绘制 102 | for det in aims: 103 | _, x, y, x0, y0 = det # 每个框左上角坐标和右下角坐标,类型为str 104 | cv2.rectangle(img0, (int(x), int(y)), (int(x0), int(y0)), (0, 255, 0), thickness=2) 105 | # (左上角坐标,右下角坐标),str转为int # (0, 255, 0)代表框的颜色 # 代表线条粗细,最低为1 106 | 107 | cv2.imshow('img', img0) 108 | cv2.waitKey(1) 109 | --------------------------------------------------------------------------------