├── README.md ├── coco ├── bus.jpg ├── coco.names ├── dog.jpg ├── person.jpg ├── picodet_m_320_coco.onnx ├── picodet_m_416_coco.onnx ├── picodet_s_320_coco.onnx ├── picodet_s_416_coco.onnx ├── street.png └── zidane.jpg ├── main.cpp └── main.py /README.md: -------------------------------------------------------------------------------- 1 | # picodet-onnxruntime 2 | 使用ONNXRuntime部署PicoDet目标检测，包含C++和Python两个版本的程序 3 | 4 | 起初我打算用opencv部署PicoDet目标检测，但是程序读取onnx文件失败了。 5 | 于是我用onnxruntime做部署就成功了，不过在PicoDet官方代码仓库里提供了10个.onnx文件， 6 | 我逐个运行之后，发现只有4个.onnx文件是onnxruntime库能正常读取的 7 | -------------------------------------------------------------------------------- /coco/bus.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/picodet-onnxruntime/dfe4d6cd0a18cd4126ee381e81df8a0b1e6ad97b/coco/bus.jpg -------------------------------------------------------------------------------- /coco/coco.names: -------------------------------------------------------------------------------- 1 | person 2 | bicycle 3 | car 4 | motorbike 5 | aeroplane 6 | bus 7 | train 8 | truck 9 | boat 10 | traffic light 11 | fire hydrant 12 | stop sign 13 | parking meter 14 | bench 15 | bird 16 | cat 17 | dog 18 | horse 19 | sheep 20 | cow 21 | elephant 22 | bear 23 | zebra 24 | giraffe 25 | backpack 26 | umbrella 27 | handbag 28 | tie 29 | suitcase 30 | frisbee 31 | skis 32 | snowboard 33 | sports ball 34 | kite 35 | baseball bat 36 | baseball glove 37 | skateboard 38 | surfboard 39 | tennis racket 40 | bottle 41 | wine glass 42 | cup 43 | fork 44 | knife 45 | spoon 46 | bowl 47 | banana 48 | apple 49 | sandwich 50 | orange 51 | broccoli 52 | carrot 53 | hot dog 54 | pizza 55 | donut 56 | cake 57 | chair 58 | sofa 59 | pottedplant 60 | bed 61 | diningtable 62 | toilet 63 | tvmonitor 64 | laptop 65 | mouse 66 | remote 67 | keyboard 68 | cell phone 69 | microwave 70 | oven 71 | toaster 72 | sink 73 | refrigerator 74 | book 75 | clock 76 | vase 77 | scissors 78 | teddy bear 79 | hair drier 80 | toothbrush 81 | -------------------------------------------------------------------------------- /coco/dog.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/picodet-onnxruntime/dfe4d6cd0a18cd4126ee381e81df8a0b1e6ad97b/coco/dog.jpg -------------------------------------------------------------------------------- /coco/person.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/picodet-onnxruntime/dfe4d6cd0a18cd4126ee381e81df8a0b1e6ad97b/coco/person.jpg -------------------------------------------------------------------------------- /coco/picodet_m_320_coco.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/picodet-onnxruntime/dfe4d6cd0a18cd4126ee381e81df8a0b1e6ad97b/coco/picodet_m_320_coco.onnx -------------------------------------------------------------------------------- /coco/picodet_m_416_coco.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/picodet-onnxruntime/dfe4d6cd0a18cd4126ee381e81df8a0b1e6ad97b/coco/picodet_m_416_coco.onnx -------------------------------------------------------------------------------- /coco/picodet_s_320_coco.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/picodet-onnxruntime/dfe4d6cd0a18cd4126ee381e81df8a0b1e6ad97b/coco/picodet_s_320_coco.onnx -------------------------------------------------------------------------------- /coco/picodet_s_416_coco.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/picodet-onnxruntime/dfe4d6cd0a18cd4126ee381e81df8a0b1e6ad97b/coco/picodet_s_416_coco.onnx -------------------------------------------------------------------------------- /coco/street.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/picodet-onnxruntime/dfe4d6cd0a18cd4126ee381e81df8a0b1e6ad97b/coco/street.png -------------------------------------------------------------------------------- /coco/zidane.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/picodet-onnxruntime/dfe4d6cd0a18cd4126ee381e81df8a0b1e6ad97b/coco/zidane.jpg -------------------------------------------------------------------------------- /main.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/picodet-onnxruntime/dfe4d6cd0a18cd4126ee381e81df8a0b1e6ad97b/main.cpp -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import argparse 4 | import onnxruntime as ort 5 | import math 6 | 7 | class PicoDet(): 8 | def __init__(self, model_pb_path, label_path, prob_threshold=0.4, iou_threshold=0.3): 9 | self.classes = list(map(lambda x: x.strip(), open(label_path, 'r').readlines())) 10 | self.num_classes = len(self.classes) 11 | self.prob_threshold = prob_threshold 12 | self.iou_threshold = iou_threshold 13 | self.mean = np.array([103.53, 116.28, 123.675], dtype=np.float32).reshape(1, 1, 3) 14 | self.std = np.array([57.375, 57.12, 58.395], dtype=np.float32).reshape(1, 1, 3) 15 | so = ort.SessionOptions() 16 | so.log_severity_level = 3 17 | self.net = ort.InferenceSession(model_pb_path, so) 18 | self.input_shape = (self.net.get_inputs()[0].shape[2], self.net.get_inputs()[0].shape[3]) 19 | self.num_outs = int(len(self.net.get_outputs()) * 0.5) 20 | self.reg_max = int(self.net.get_outputs()[self.num_outs].shape[-1] / 4) - 1 21 | self.project = np.arange(self.reg_max + 1) 22 | self.strides = [int(8 * (2**i)) for i in range(self.num_outs)] 23 | self.mlvl_anchors = [] 24 | for i in range(len(self.strides)): 25 | anchors = self._make_grid((math.ceil(self.input_shape[0] / self.strides[i]), math.ceil(self.input_shape[1] / self.strides[i])), self.strides[i]) 26 | self.mlvl_anchors.append(anchors) 27 | def _make_grid(self, featmap_size, stride): 28 | feat_h, feat_w = featmap_size 29 | shift_x = np.arange(0, feat_w) * stride 30 | shift_y = np.arange(0, feat_h) * stride 31 | xv, yv = np.meshgrid(shift_x, shift_y) 32 | xv = xv.flatten() 33 | yv = yv.flatten() 34 | cx = xv + 0.5 * (stride-1) 35 | cy = yv + 0.5 * (stride - 1) 36 | return np.stack((cx, cy), axis=-1) 37 | def softmax(self,x, axis=1): 38 | x_exp = np.exp(x) 39 | # 如果是列向量，则axis=0 40 | x_sum = np.sum(x_exp, axis=axis, keepdims=True) 41 | s = x_exp / x_sum 42 | return s 43 | 44 | def _normalize(self, img): ### c++: https://blog.csdn.net/wuqingshan2010/article/details/107727909 45 | img = img.astype(np.float32) 46 | img = (img / 255.0 - self.mean / 255.0) / (self.std / 255.0) 47 | # img = (img - self.mean) / (self.std) 48 | return img 49 | def resize_image(self, srcimg, keep_ratio=True): 50 | top, left, newh, neww = 0, 0, self.input_shape[0], self.input_shape[1] 51 | if keep_ratio and srcimg.shape[0] != srcimg.shape[1]: 52 | hw_scale = srcimg.shape[0] / srcimg.shape[1] 53 | if hw_scale > 1: 54 | newh, neww = self.input_shape[0], int(self.input_shape[1] / hw_scale) 55 | img = cv2.resize(srcimg, (neww, newh), interpolation=cv2.INTER_AREA) 56 | left = int((self.input_shape[1] - neww) * 0.5) 57 | img = cv2.copyMakeBorder(img, 0, 0, left, self.input_shape[1] - neww - left, cv2.BORDER_CONSTANT, 58 | value=0) # add border 59 | else: 60 | newh, neww = int(self.input_shape[0] * hw_scale), self.input_shape[1] 61 | img = cv2.resize(srcimg, (neww, newh), interpolation=cv2.INTER_AREA) 62 | top = int((self.input_shape[0] - newh) * 0.5) 63 | img = cv2.copyMakeBorder(img, top, self.input_shape[0] - newh - top, 0, 0, cv2.BORDER_CONSTANT, value=0) 64 | else: 65 | img = cv2.resize(srcimg, self.input_shape, interpolation=cv2.INTER_AREA) 66 | return img, newh, neww, top, left 67 | def post_process(self, preds): 68 | cls_scores, bbox_preds = preds[:self.num_outs], preds[self.num_outs:] 69 | det_bboxes, det_conf, det_classid = self.get_bboxes_single(cls_scores, bbox_preds, 1, rescale=False) 70 | return det_bboxes.astype(np.int32), det_conf, det_classid 71 | def get_bboxes_single(self, cls_scores, bbox_preds, scale_factor, rescale=False): 72 | mlvl_bboxes = [] 73 | mlvl_scores = [] 74 | for stride, cls_score, bbox_pred, anchors in zip(self.strides, cls_scores, bbox_preds, self.mlvl_anchors): 75 | if cls_score.ndim==3: 76 | cls_score = cls_score.squeeze(axis=0) 77 | if bbox_pred.ndim==3: 78 | bbox_pred = bbox_pred.squeeze(axis=0) 79 | bbox_pred = self.softmax(bbox_pred.reshape(-1, self.reg_max + 1), axis=1) 80 | # bbox_pred = np.sum(bbox_pred * np.expand_dims(self.project, axis=0), axis=1).reshape((-1, 4)) 81 | bbox_pred = np.dot(bbox_pred, self.project).reshape(-1,4) 82 | bbox_pred *= stride 83 | 84 | # nms_pre = cfg.get('nms_pre', -1) 85 | nms_pre = 1000 86 | if nms_pre > 0 and cls_score.shape[0] > nms_pre: 87 | max_scores = cls_score.max(axis=1) 88 | topk_inds = max_scores.argsort()[::-1][0:nms_pre] 89 | anchors = anchors[topk_inds, :] 90 | bbox_pred = bbox_pred[topk_inds, :] 91 | cls_score = cls_score[topk_inds, :] 92 | 93 | bboxes = self.distance2bbox(anchors, bbox_pred, max_shape=self.input_shape) 94 | mlvl_bboxes.append(bboxes) 95 | mlvl_scores.append(cls_score) 96 | 97 | mlvl_bboxes = np.concatenate(mlvl_bboxes, axis=0) 98 | if rescale: 99 | mlvl_bboxes /= scale_factor 100 | mlvl_scores = np.concatenate(mlvl_scores, axis=0) 101 | 102 | bboxes_wh = mlvl_bboxes.copy() 103 | bboxes_wh[:, 2:4] = bboxes_wh[:, 2:4] - bboxes_wh[:, 0:2] ####xywh 104 | classIds = np.argmax(mlvl_scores, axis=1) 105 | confidences = np.max(mlvl_scores, axis=1) ####max_class_confidence 106 | 107 | indices = cv2.dnn.NMSBoxes(bboxes_wh.tolist(), confidences.tolist(), self.prob_threshold, self.iou_threshold) 108 | if len(indices)>0: 109 | indices=indices.flatten() 110 | mlvl_bboxes = mlvl_bboxes[indices] 111 | confidences = confidences[indices] 112 | classIds = classIds[indices] 113 | return mlvl_bboxes, confidences, classIds 114 | else: 115 | print('nothing detect') 116 | return np.array([]), np.array([]), np.array([]) 117 | def distance2bbox(self, points, distance, max_shape=None): 118 | x1 = points[:, 0] - distance[:, 0] 119 | y1 = points[:, 1] - distance[:, 1] 120 | x2 = points[:, 0] + distance[:, 2] 121 | y2 = points[:, 1] + distance[:, 3] 122 | if max_shape is not None: 123 | x1 = np.clip(x1, 0, max_shape[1]) 124 | y1 = np.clip(y1, 0, max_shape[0]) 125 | x2 = np.clip(x2, 0, max_shape[1]) 126 | y2 = np.clip(y2, 0, max_shape[0]) 127 | return np.stack([x1, y1, x2, y2], axis=-1) 128 | def detect(self, srcimg): 129 | img, newh, neww, top, left = self.resize_image(srcimg) 130 | img = self._normalize(img) 131 | blob = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0) 132 | 133 | outs = self.net.run(None, {self.net.get_inputs()[0].name: blob}) 134 | det_bboxes, det_conf, det_classid = self.post_process(outs) 135 | 136 | # results = [] 137 | ratioh,ratiow = srcimg.shape[0]/newh,srcimg.shape[1]/neww 138 | for i in range(det_bboxes.shape[0]): 139 | xmin, ymin, xmax, ymax = max(int((det_bboxes[i,0] - left) * ratiow), 0), max(int((det_bboxes[i,1] - top) * ratioh), 0), min( 140 | int((det_bboxes[i,2] - left) * ratiow), srcimg.shape[1]), min(int((det_bboxes[i,3] - top) * ratioh), srcimg.shape[0]) 141 | # results.append((xmin, ymin, xmax, ymax, self.classes[det_classid[i]], det_conf[i])) 142 | cv2.rectangle(srcimg, (xmin, ymin), (xmax, ymax), (0, 0, 255), thickness=1) 143 | print(self.classes[det_classid[i]]+': '+str(round(det_conf[i], 3))) 144 | cv2.putText(srcimg, self.classes[det_classid[i]]+': '+str(round(det_conf[i], 3)), (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), thickness=1) 145 | # cv2.imwrite('result.jpg', srcimg) 146 | return srcimg 147 | 148 | if __name__=='__main__': 149 | parser = argparse.ArgumentParser() 150 | parser.add_argument('--imgpath', type=str, default='coco/dog.jpg', help="image path") 151 | parser.add_argument('--modelpath', type=str, default='coco/picodet_s_320_coco.onnx', choices=["coco/picodet_m_320_coco.onnx", "coco/picodet_m_416_coco.onnx", "coco/picodet_s_320_coco.onnx", "coco/picodet_s_416_coco.onnx"], help="onnx filepath") 152 | parser.add_argument('--classfile', type=str, default='coco/coco.names', help="classname filepath") 153 | parser.add_argument('--confThreshold', default=0.5, type=float, help='class confidence') 154 | parser.add_argument('--nmsThreshold', default=0.6, type=float, help='nms iou thresh') 155 | args = parser.parse_args() 156 | 157 | srcimg = cv2.imread(args.imgpath) 158 | net = PicoDet(args.modelpath, args.classfile, prob_threshold=args.confThreshold, iou_threshold=args.nmsThreshold) 159 | srcimg = net.detect(srcimg) 160 | 161 | winName = 'Deep learning object detection in onnxruntime' 162 | cv2.namedWindow(winName, cv2.WINDOW_NORMAL) 163 | cv2.imshow(winName, srcimg) 164 | cv2.waitKey(0) 165 | cv2.destroyAllWindows() --------------------------------------------------------------------------------