diff --git a/README.md b/README.md index 7bd6d58..1af6294 100644 --- a/README.md +++ b/README.md @@ -167,8 +167,18 @@ TRAIN: ``` # Deploy -## NCNN -* Waiting for update +## Export onnx +* You can export .onnx by adding the --onnx option when executing test.py + ``` + python3 test.py --yaml configs/config.yaml --weight weights/weight_AP05\:0.278_280-epoch.pth --img data/3.jpg --onnx + ``` +## onnx-runtime +* You can learn about the pre and post-processing methods of FastestDet in this Sample + ``` + cd example/onnx-runtime + pip install onnx-runtime + python3 runtime.py + ``` # Citation * If you find this project useful in your research, please consider cite: ``` diff --git a/example/ncnn/.gitkeep b/example/ncnn/.gitkeep new file mode 100644 index 0000000..af634f0 --- /dev/null +++ b/example/ncnn/.gitkeep @@ -0,0 +1 @@ +# pytorch-detector \ No newline at end of file diff --git a/example/onnx-runtime/1.jpg b/example/onnx-runtime/1.jpg new file mode 100644 index 0000000..7912379 Binary files /dev/null and b/example/onnx-runtime/1.jpg differ diff --git a/example/onnx-runtime/FastestDet.onnx b/example/onnx-runtime/FastestDet.onnx new file mode 100644 index 0000000..f82f538 Binary files /dev/null and b/example/onnx-runtime/FastestDet.onnx differ diff --git a/example/onnx-runtime/coco.names b/example/onnx-runtime/coco.names new file mode 100644 index 0000000..ca76c80 --- /dev/null +++ b/example/onnx-runtime/coco.names @@ -0,0 +1,80 @@ +person +bicycle +car +motorbike +aeroplane +bus +train +truck +boat +traffic light +fire hydrant +stop sign +parking meter +bench +bird +cat +dog +horse +sheep +cow +elephant +bear +zebra +giraffe +backpack +umbrella +handbag +tie +suitcase +frisbee +skis +snowboard +sports ball +kite +baseball bat +baseball glove +skateboard +surfboard +tennis racket +bottle +wine glass +cup +fork +knife +spoon +bowl +banana +apple +sandwich +orange +broccoli +carrot +hot dog +pizza +donut +cake +chair +sofa +pottedplant +bed +diningtable +toilet +tvmonitor +laptop +mouse +remote +keyboard +cell phone +microwave +oven +toaster +sink +refrigerator +book +clock +vase +scissors +teddy bear +hair drier +toothbrush diff --git a/example/onnx-runtime/result.jpg b/example/onnx-runtime/result.jpg new file mode 100644 index 0000000..04d02e1 Binary files /dev/null and b/example/onnx-runtime/result.jpg differ diff --git a/example/onnx-runtime/runtime.py b/example/onnx-runtime/runtime.py new file mode 100755 index 0000000..6e38afd --- /dev/null +++ b/example/onnx-runtime/runtime.py @@ -0,0 +1,142 @@ +import cv2 +import numpy as np +import onnxruntime + +# sigmoid函数 +def sigmoid(x): + return 1. / (1 + np.exp(-x)) + +# tanh函数 +def tanh(x): + return 2. / (1 + np.exp(-2 * x)) - 1 + +# 数据预处理 +def preprocess(src_img, size): + output = cv2.resize(src_img,(size[0], size[1]),interpolation=cv2.INTER_AREA) + output = output.transpose(2,0,1) + output = output.reshape((1, 3, size[1], size[0])) / 255 + + return output.astype('float32') + +# nms算法 +def nms(dets, thresh=0.35): + # dets:N*M,N是bbox的个数，M的前4位是对应的（x1,y1,x2,y2），第5位是对应的分数 + # #thresh:0.3,0.5.... + x1 = dets[:, 0] + y1 = dets[:, 1] + x2 = dets[:, 2] + y2 = dets[:, 3] + scores = dets[:, 4] + areas = (x2 - x1 + 1) * (y2 - y1 + 1) # 求每个bbox的面积 + order = scores.argsort()[::-1] # 对分数进行倒排序 + keep = [] # 用来保存最后留下来的bboxx下标 + + while order.size > 0: + i = order[0] # 无条件保留每次迭代中置信度最高的bbox + keep.append(i) + + # 计算置信度最高的bbox和其他剩下bbox之间的交叉区域 + xx1 = np.maximum(x1[i], x1[order[1:]]) + yy1 = np.maximum(y1[i], y1[order[1:]]) + xx2 = np.minimum(x2[i], x2[order[1:]]) + yy2 = np.minimum(y2[i], y2[order[1:]]) + + # 计算置信度高的bbox和其他剩下bbox之间交叉区域的面积 + w = np.maximum(0.0, xx2 - xx1 + 1) + h = np.maximum(0.0, yy2 - yy1 + 1) + inter = w * h + + # 求交叉区域的面积占两者（置信度高的bbox和其他bbox）面积和的必烈 + ovr = inter / (areas[i] + areas[order[1:]] - inter) + + # 保留ovr小于thresh的bbox，进入下一次迭代。 + inds = np.where(ovr <= thresh)[0] + + # 因为ovr中的索引不包括order[0]所以要向后移动一位 + order = order[inds + 1] + + output = [] + for i in keep: + output.append(dets[i].tolist()) + + return output + +# 人脸检测 +def detection(session, img, input_width, input_height, thresh): + pred = [] + + # 输入图像的原始宽高 + H, W, _ = img.shape + + # 数据预处理: resize, 1/255 + data = preprocess(img, [input_width, input_height]) + + # 模型推理 + input_name = session.get_inputs()[0].name + feature_map = session.run([], {input_name: data})[0][0] + + # 输出特征图转置: CHW, HWC + feature_map = feature_map.transpose(1, 2, 0) + # 输出特征图的宽高 + feature_map_height = feature_map.shape[0] + feature_map_width = feature_map.shape[1] + + # 特征图后处理 + for h in range(feature_map_height): + for w in range(feature_map_width): + data = feature_map[h][w] + + # 解析检测框置信度 + obj_score, cls_score = data[0], data[5:].max() + score = obj_score * cls_score + + # 阈值筛选 + if score > thresh: + # 检测框类别 + cls_index = np.argmax(data[5:]) + # 检测框中心点偏移 + x_offset, y_offset = tanh(data[1]), tanh(data[2]) + # 检测框归一化后的宽高 + box_width, box_height = sigmoid(data[3]), sigmoid(data[4]) + # 检测框归一化后中心点 + box_cx = (w + x_offset) / feature_map_width + box_cy = (h + y_offset) / feature_map_height + + # cx,cy,w,h => x1, y1, x2, y2 + x1, y1 = box_cx - 0.5 * box_width, box_cy - 0.5 * box_height + x2, y2 = box_cx + 0.5 * box_width, box_cy + 0.5 * box_height + x1, y1, x2, y2 = int(x1 * W), int(y1 * H), int(x2 * W), int(y2 * H) + + pred.append([x1, y1, x2, y2, score, cls_index]) + + return nms(np.array(pred)) + +if __name__ == '__main__': + # 读取图片 + img = cv2.imread("1.jpg") + # 模型输入的宽高 + input_width, input_height = 512, 512 + # 加载模型 + session = onnxruntime.InferenceSession('FastestDet.onnx') + # 目标检测 + bboxes = detection(session, img, input_width, input_height, 0.8) + + # 加载label names + names = [] + with open("coco.names", 'r') as f: + for line in f.readlines(): + names.append(line.strip()) + + for b in bboxes: + print(b) + obj_score, cls_index = b[4], int(b[5]) + x1, y1, x2, y2 = int(b[0]), int(b[1]), int(b[2]), int(b[3]) + + #绘制检测框 + cv2.rectangle(img, (x1,y1), (x2, y2), (255, 255, 0), 2) + cv2.putText(img, '%.2f' % obj_score, (x1, y1 - 5), 0, 0.7, (0, 255, 0), 2) + cv2.putText(img, names[cls_index], (x1, y1 - 25), 0, 0.7, (0, 255, 0), 2) + + cv2.imwrite("result.jpg", img) + +