├── README.md ├── coco.names ├── images ├── bus.jpg ├── dog.jpg ├── person.jpg └── zidane.jpg ├── main.cpp └── main.py /README.md: -------------------------------------------------------------------------------- 1 | # DAMO-YOLO-detect-onnxrun-cpp-py 2 | 使用ONNXRuntime部署DAMO-YOLO目标检测，包含C++和Python两个版本的程序。 3 | 起初，我想使用opencv做部署的，但是opencv的dnn模块读取onnx文件出错， 4 | 无赖只能使用onnxruntime做部署了。本套程序一共提供了27个onnx模型， 5 | onnx文件需要从百度云盘下载， 6 | 链接：https://pan.baidu.com/s/10-5ke_fs2omqUMSgKTJV0Q 7 | 提取码：w9kp 8 | 9 | 其中在百度云盘里一共有30个onnx模型文件，但是逐个读取onnx文件做推理时， 10 | 发现有3个onnx文件在onnxruntime读取时出错了，在程序里的choices参数里声明了 11 | 27个模型文件的名称。 12 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /images/bus.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/DAMO-YOLO-detect-onnxrun-cpp-py/3fe8df1df5aa1ad0e4a7b211839a98a5d56abc10/images/bus.jpg -------------------------------------------------------------------------------- /images/dog.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/DAMO-YOLO-detect-onnxrun-cpp-py/3fe8df1df5aa1ad0e4a7b211839a98a5d56abc10/images/dog.jpg -------------------------------------------------------------------------------- /images/person.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/DAMO-YOLO-detect-onnxrun-cpp-py/3fe8df1df5aa1ad0e4a7b211839a98a5d56abc10/images/person.jpg -------------------------------------------------------------------------------- /images/zidane.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/DAMO-YOLO-detect-onnxrun-cpp-py/3fe8df1df5aa1ad0e4a7b211839a98a5d56abc10/images/zidane.jpg -------------------------------------------------------------------------------- /main.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/DAMO-YOLO-detect-onnxrun-cpp-py/3fe8df1df5aa1ad0e4a7b211839a98a5d56abc10/main.cpp -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import cv2 3 | import numpy as np 4 | import onnxruntime as ort 5 | 6 | 7 | class DAMO_YOLO(): 8 | def __init__(self, model_path, confThreshold=0.5, nmsThreshold=0.5): 9 | self.classes = list(map(lambda x: x.strip(), open('coco.names', 'r').readlines())) 10 | self.num_class = len(self.classes) 11 | so = ort.SessionOptions() 12 | so.log_severity_level = 3 13 | self.session = ort.InferenceSession(model_path, so) 14 | model_inputs = self.session.get_inputs() 15 | self.input_name = model_inputs[0].name 16 | self.input_shape = model_inputs[0].shape 17 | self.input_height = int(self.input_shape[2]) 18 | self.input_width = int(self.input_shape[3]) 19 | 20 | self.confThreshold = confThreshold 21 | self.nmsThreshold = nmsThreshold 22 | 23 | def drawPred(self, frame, classId, conf, left, top, right, bottom): 24 | # Draw a bounding box. 25 | cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), thickness=2) 26 | 27 | label = '%.2f' % conf 28 | label = '%s:%s' % (self.classes[classId], label) 29 | 30 | # Display the label at the top of the bounding box 31 | labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) 32 | top = max(top, labelSize[1]) 33 | # cv.rectangle(frame, (left, top - round(1.5 * labelSize[1])), (left + round(1.5 * labelSize[0]), top + baseLine), (255,255,255), cv.FILLED) 34 | cv2.putText(frame, label, (left, top - 10), 0, 0.7, (0, 255, 0), thickness=2) 35 | return frame 36 | 37 | def detect(self, frame): 38 | temp_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) 39 | padded_image = np.ones((self.input_height, self.input_width, 3), dtype=np.uint8) 40 | ratio = min(self.input_height / temp_image.shape[0], self.input_width / temp_image.shape[1]) 41 | neww, newh = int(temp_image.shape[1] * ratio), int(temp_image.shape[0] * ratio) 42 | temp_image = cv2.resize(temp_image, (neww, newh), interpolation=cv2.INTER_LINEAR) 43 | padded_image[:newh, :neww, :] = temp_image 44 | 45 | padded_image = padded_image.transpose(2, 0, 1) 46 | padded_image = np.expand_dims(padded_image, axis=0).astype(np.float32) 47 | 48 | # Inference 49 | results = self.session.run(None, {self.input_name: padded_image}) 50 | 51 | scores, bboxes = results[0].squeeze(axis=0), results[1].squeeze(axis=0) 52 | bboxes /= ratio 53 | 54 | boxes, confidences, classIds = [], [], [] 55 | for i in range(bboxes.shape[0]): 56 | score = np.max(scores[i, :]) 57 | if score < self.confThreshold: 58 | continue 59 | 60 | class_id = np.argmax(scores[i, :]) 61 | x, y, xmax, ymax = bboxes[i, :].astype(np.int32) 62 | width, height = xmax - x, ymax - y 63 | 64 | boxes.append([x, y, width, height]) 65 | classIds.append(class_id) 66 | confidences.append(score) 67 | indices = cv2.dnn.NMSBoxes(boxes, confidences, self.confThreshold, self.nmsThreshold) 68 | for i in indices: 69 | box = boxes[i] 70 | left = box[0] 71 | top = box[1] 72 | width = box[2] 73 | height = box[3] 74 | frame = self.drawPred(frame, classIds[i], confidences[i], left, top, left + width, top + height) 75 | return frame 76 | 77 | 78 | if __name__ == '__main__': 79 | parser = argparse.ArgumentParser() 80 | parser.add_argument("--modelpath", type=str, default='weights/damoyolo_tinynasL20_T_192x320.onnx', 81 | choices=["weights/damoyolo_tinynasL20_T_192x320.onnx", 82 | "weights/damoyolo_tinynasL20_T_256x320.onnx", 83 | "weights/damoyolo_tinynasL20_T_256x416.onnx", 84 | "weights/damoyolo_tinynasL20_T_288x480.onnx", 85 | "weights/damoyolo_tinynasL20_T_384x640.onnx", 86 | "weights/damoyolo_tinynasL20_T_480x640.onnx", 87 | "weights/damoyolo_tinynasL20_T_480x800.onnx", 88 | "weights/damoyolo_tinynasL20_T_640x640.onnx", 89 | "weights/damoyolo_tinynasL20_T_736x1280.onnx", 90 | "weights/damoyolo_tinynasL25_S_192x320.onnx", 91 | "weights/damoyolo_tinynasL25_S_256x320.onnx", 92 | "weights/damoyolo_tinynasL25_S_256x416.onnx", 93 | "weights/damoyolo_tinynasL25_S_288x480.onnx", 94 | "weights/damoyolo_tinynasL25_S_384x640.onnx", 95 | "weights/damoyolo_tinynasL25_S_480x640.onnx", 96 | "weights/damoyolo_tinynasL25_S_480x800.onnx", 97 | "weights/damoyolo_tinynasL25_S_640x640.onnx", 98 | "weights/damoyolo_tinynasL25_S_736x1280.onnx", 99 | "weights/damoyolo_tinynasL35_M_192x320.onnx", 100 | "weights/damoyolo_tinynasL35_M_256x320.onnx", 101 | "weights/damoyolo_tinynasL35_M_256x416.onnx", 102 | "weights/damoyolo_tinynasL35_M_288x480.onnx", 103 | "weights/damoyolo_tinynasL35_M_384x640.onnx", 104 | "weights/damoyolo_tinynasL35_M_480x640.onnx", 105 | "weights/damoyolo_tinynasL35_M_480x800.onnx", 106 | "weights/damoyolo_tinynasL35_M_640x640.onnx", 107 | "weights/damoyolo_tinynasL35_M_736x1280.onnx"], help="model path") 108 | parser.add_argument("--imgpath", type=str, default='images/dog.jpg', help="image path") 109 | parser.add_argument("--confThreshold", default=0.5, type=float, help='class confidence') 110 | parser.add_argument("--nmsThreshold", default=0.85, type=float, help='iou thresh') 111 | args = parser.parse_args() 112 | 113 | net = DAMO_YOLO(args.modelpath, confThreshold=args.confThreshold, nmsThreshold=args.nmsThreshold) 114 | srcimg = cv2.imread(args.imgpath) 115 | srcimg = net.detect(srcimg) 116 | 117 | winName = 'Deep learning object detection in ONNXRuntime' 118 | cv2.namedWindow(winName, cv2.WINDOW_NORMAL) 119 | cv2.imshow(winName, srcimg) 120 | cv2.waitKey(0) 121 | cv2.destroyAllWindows() 122 | --------------------------------------------------------------------------------