yolov7-onnxruntime
司徒正信
yolov7-onnxruntime调用示例
参考文章
https://cloud.tencent.com/developer/article/2205367
https://github.com/hpc203/yolov7-opencv-onnxrun-cpp-py
官方文档 https://onnxruntime.ai/docs/api/c/index.html
python示例
这里测试可正常推理自定义模型，模型导出参考文章yolov7-opencv
import cv2
import numpy as np
import onnxruntime
import argparse
class YOLOv7:
def __init__(self, path, conf_thres=0.7, iou_thres=0.5):
self.conf_threshold = conf_thres
self.iou_threshold = iou_thres
#self.class_names = list(map(lambda x: x.strip(), open('coco.names', 'r').readlines()))
self.class_names =["ceshi1","ceshi2","ceshi3"]
# Initialize model
self.session = onnxruntime.InferenceSession(path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
model_inputs = self.session.get_inputs()
self.input_names = [model_inputs[i].name for i in range(len(model_inputs))]
self.input_shape = model_inputs[0].shape
self.input_height = self.input_shape[2]
self.input_width = self.input_shape[3]
model_outputs = self.session.get_outputs()
self.output_names = [model_outputs[i].name for i in range(len(model_outputs))]
self.has_postprocess = 'score' in self.output_names
def detect(self, image):
input_tensor = self.prepare_input(image)
# Perform inference on the image
outputs = self.session.run(self.output_names, {self.input_names[0]: input_tensor})
if self.has_postprocess:
boxes, scores, class_ids = self.parse_processed_output(outputs)
else:
# Process output data
boxes, scores, class_ids = self.process_output(outputs)
return boxes, scores, class_ids
def prepare_input(self, image):
self.img_height, self.img_width = image.shape[:2]
input_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# Resize input image
input_img = cv2.resize(input_img, (self.input_width, self.input_height))
# Scale input pixel values to 0 to 1
input_img = input_img / 255.0
input_img = input_img.transpose(2, 0, 1)
input_tensor = input_img[np.newaxis, :, :, :].astype(np.float32)
return input_tensor
def process_output(self, output):
predictions = np.squeeze(output[0])
# Filter out object confidence scores below threshold
obj_conf = predictions[:, 4]
predictions = predictions[obj_conf > self.conf_threshold]
obj_conf = obj_conf[obj_conf > self.conf_threshold]
# Multiply class confidence with bounding box confidence
predictions[:, 5:] *= obj_conf[:, np.newaxis]
# Get the scores
scores = np.max(predictions[:, 5:], axis=1)
# Filter out the objects with a low score
valid_scores = scores > self.conf_threshold
predictions = predictions[valid_scores]
scores = scores[valid_scores]
# Get the class with the highest confidence
class_ids = np.argmax(predictions[:, 5:], axis=1)
# Get bounding boxes for each object
boxes = self.extract_boxes(predictions)
# Apply non-maxima suppression to suppress weak, overlapping bounding boxes
# indices = nms(boxes, scores, self.iou_threshold)
indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), self.conf_threshold, self.iou_threshold).flatten()
return boxes[indices], scores[indices], class_ids[indices]
def parse_processed_output(self, outputs):
scores = np.squeeze(outputs[self.output_names.index('score')])
predictions = outputs[self.output_names.index('batchno_classid_x1y1x2y2')]
# Filter out object scores below threshold
valid_scores = scores > self.conf_threshold
predictions = predictions[valid_scores, :]
scores = scores[valid_scores]
# Extract the boxes and class ids
# TODO: Separate based on batch number
batch_number = predictions[:, 0]
class_ids = predictions[:, 1]
boxes = predictions[:, 2:]
# In postprocess, the x,y are the y,x
boxes = boxes[:, [1, 0, 3, 2]]
# Rescale boxes to original image dimensions
boxes = self.rescale_boxes(boxes)
return boxes, scores, class_ids
def extract_boxes(self, predictions):
# Extract boxes from predictions
boxes = predictions[:, :4]
# Scale boxes to original image dimensions
boxes = self.rescale_boxes(boxes)
# Convert boxes to xywh format
boxes_ = np.copy(boxes)
boxes_[..., 0] = boxes[..., 0] - boxes[..., 2] * 0.5
boxes_[..., 1] = boxes[..., 1] - boxes[..., 3] * 0.5
return boxes_
def rescale_boxes(self, boxes):
# Rescale boxes to original image dimensions
input_shape = np.array([self.input_width, self.input_height, self.input_width, self.input_height])
boxes = np.divide(boxes, input_shape, dtype=np.float32)
boxes *= np.array([self.img_width, self.img_height, self.img_width, self.img_height])
return boxes
def draw_detections(self, image, boxes, scores, class_ids):
for box, score, class_id in zip(boxes, scores, class_ids):
x, y, w, h = box.astype(int)
print("当前找到%s left %d top %d width %d height %d" %(self.class_names[class_id],x,y,w,h))
# Draw rectangle
cv2.rectangle(image, (x, y), (x+w, y+h), (0, 0, 255), thickness=2)
label = self.class_names[class_id]
label = f'{label} {int(score * 100)}%'
labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
# top = max(y1, labelSize[1])
# cv.rectangle(frame, (left, top - round(1.5 * labelSize[1])), (left + round(1.5 * labelSize[0]), top + baseLine), (255,255,255), cv.FILLED)
cv2.putText(image, label, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), thickness=2)
return image
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--imgpath', type=str, default='datasets/test.jpg', help="image path")
parser.add_argument('--modelpath', type=str, default='runs/train/exp/weights/best.onnx',
choices=["models/yolov7_640x640.onnx", "models/yolov7-tiny_640x640.onnx",
"models/yolov7_736x1280.onnx", "models/yolov7-tiny_384x640.onnx",
"models/yolov7_480x640.onnx", "models/yolov7_384x640.onnx",
"models/yolov7-tiny_256x480.onnx", "models/yolov7-tiny_256x320.onnx",
"models/yolov7_256x320.onnx", "models/yolov7-tiny_256x640.onnx",
"models/yolov7_256x640.onnx", "models/yolov7-tiny_480x640.onnx",
"models/yolov7-tiny_736x1280.onnx", "models/yolov7_256x480.onnx"],
help="onnx filepath")
parser.add_argument('--confThreshold', default=0.3, type=float, help='class confidence')
parser.add_argument('--nmsThreshold', default=0.5, type=float, help='nms iou thresh')
args = parser.parse_args()
# Initialize YOLOv7 object detector
yolov7_detector = YOLOv7(args.modelpath, conf_thres=args.confThreshold, iou_thres=args.nmsThreshold)
srcimg = cv2.imread(args.imgpath)
# Detect Objects
boxes, scores, class_ids = yolov7_detector.detect(srcimg)
# Draw detections
dstimg = yolov7_detector.draw_detections(srcimg, boxes, scores, class_ids)
winName = 'Deep learning object detection in ONNXRuntime'
cv2.namedWindow(winName, 0)
cv2.imshow(winName, dstimg)
cv2.waitKey(0)
cv2.destroyAllWindows()
这里onnxruntime版本使用的是1.13.1
>>> import onnxruntime
>>> onnxruntime.__version__
'1.13.1'
>>>
cpp版本调用示例
onnxruntime下载地址 https://github.com/microsoft/onnxruntime/releases/tag/v1.12.1
这里测试1.13.1版本执行推理就报内存错误，在多次debug的过程中神奇的有两次成功显示结果
后来改为下载1.12.1版本，推理成功
#include <fstream>
#include <sstream>
#include <iostream>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
//#include <cuda_provider_factory.h>
#include <onnxruntime_cxx_api.h>
using namespace std;
using namespace cv;
using namespace Ort;
struct Net_config
{
float confThreshold; // Confidence threshold
float nmsThreshold; // Non-maximum suppression threshold
string modelpath;
};
typedef struct BoxInfo
{
float x1;
float y1;
float x2;
float y2;
float score;
int label;
} BoxInfo;
class YOLOV7
{
public:
YOLOV7(Net_config config);
void detect(Mat& frame);
private:
int inpWidth;
int inpHeight;
int nout;
int num_proposal;
vector<string> class_names;
int num_class;
float confThreshold;
float nmsThreshold;
vector<float> input_image_;
void normalize_(Mat img);
void nms(vector<BoxInfo>& input_boxes);
Env env = Env(ORT_LOGGING_LEVEL_ERROR, "YOLOV7");
Ort::Session* ort_session = nullptr;
SessionOptions sessionOptions = SessionOptions();
vector<char *> input_names;
vector<char *> output_names;
vector<vector<int64_t>> input_node_dims; // >=1 outputs
vector<vector<int64_t>> output_node_dims; // >=1 outputs
};
YOLOV7::YOLOV7(Net_config config)
{
this->confThreshold = config.confThreshold;
this->nmsThreshold = config.nmsThreshold;
//string classesFile = "coco.names";
string model_path = config.modelpath;
std::wstring widestr = std::wstring(model_path.begin(), model_path.end());
//OrtStatus* status = OrtSessionOptionsAppendExecutionProvider_CUDA(sessionOptions, 0);
sessionOptions.SetGraphOptimizationLevel(ORT_ENABLE_BASIC);
ort_session = new Session(env, widestr.c_str(), sessionOptions);
size_t numInputNodes = ort_session->GetInputCount();
size_t numOutputNodes = ort_session->GetOutputCount();
AllocatorWithDefaultOptions allocator;
for (int i = 0; i < numInputNodes; i++)
{
//input_names.push_back(ort_session->GetInputNameAllocated(i, allocator).get()); //1.13.1 用法
input_names.push_back(ort_session->GetInputName(i, allocator)); //1.12.1 用法
//input_names.push_back();
Ort::TypeInfo input_type_info = ort_session->GetInputTypeInfo(i);
auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo();
auto input_dims = input_tensor_info.GetShape();
input_node_dims.push_back(input_dims);
}
for (int i = 0; i < numOutputNodes; i++)
{
// output_names.push_back(ort_session->GetOutputNameAllocated(i, allocator).get()); //1.13.1 用法
output_names.push_back(ort_session->GetOutputName(i, allocator)); // 1.12.1
Ort::TypeInfo output_type_info = ort_session->GetOutputTypeInfo(i);
auto output_tensor_info = output_type_info.GetTensorTypeAndShapeInfo();
auto output_dims = output_tensor_info.GetShape();
output_node_dims.push_back(output_dims);
}
this->inpHeight = input_node_dims[0][2];
this->inpWidth = input_node_dims[0][3];
this->nout = output_node_dims[0][2];
this->num_proposal = output_node_dims[0][1];
//ifstream ifs(classesFile.c_str());
//string line;
//while (getline(ifs, line)) this->class_names.push_back(line);
this->class_names = { "ceshi1","ceshi2","ceshi3" }; //模型中的类别
this->num_class = class_names.size();
}
void YOLOV7::normalize_(Mat img)
{
// img.convertTo(img, CV_32F);
int row = img.rows;
int col = img.cols;
this->input_image_.resize(row * col * img.channels());
for (int c = 0; c < 3; c++)
{
for (int i = 0; i < row; i++)
{
for (int j = 0; j < col; j++)
{
float pix = img.ptr<uchar>(i)[j * 3 + 2 - c];
this->input_image_[c * row * col + i * col + j] = pix / 255.0;
}
}
}
}
void YOLOV7::nms(vector<BoxInfo>& input_boxes)
{
sort(input_boxes.begin(), input_boxes.end(), [](BoxInfo a, BoxInfo b) { return a.score > b.score; });
vector<float> vArea(input_boxes.size());
for (int i = 0; i < int(input_boxes.size()); ++i)
{
vArea[i] = (input_boxes.at(i).x2 - input_boxes.at(i).x1 + 1)
* (input_boxes.at(i).y2 - input_boxes.at(i).y1 + 1);
}
vector<bool> isSuppressed(input_boxes.size(), false);
for (int i = 0; i < int(input_boxes.size()); ++i)
{
if (isSuppressed[i]) { continue; }
for (int j = i + 1; j < int(input_boxes.size()); ++j)
{
if (isSuppressed[j]) { continue; }
float xx1 = (max)(input_boxes[i].x1, input_boxes[j].x1);
float yy1 = (max)(input_boxes[i].y1, input_boxes[j].y1);
float xx2 = (min)(input_boxes[i].x2, input_boxes[j].x2);
float yy2 = (min)(input_boxes[i].y2, input_boxes[j].y2);
float w = (max)(float(0), xx2 - xx1 + 1);
float h = (max)(float(0), yy2 - yy1 + 1);
float inter = w * h;
float ovr = inter / (vArea[i] + vArea[j] - inter);
if (ovr >= this->nmsThreshold)
{
isSuppressed[j] = true;
}
}
}
// return post_nms;
int idx_t = 0;
input_boxes.erase(remove_if(input_boxes.begin(), input_boxes.end(), [&idx_t, &isSuppressed](const BoxInfo& f) { return isSuppressed[idx_t++]; }), input_boxes.end());
}
void YOLOV7::detect(Mat& frame)
{
Mat dstimg;
resize(frame, dstimg, Size(this->inpWidth, this->inpHeight));
this->normalize_(dstimg);
array<int64_t, 4> input_shape_{ 1, 3, this->inpHeight, this->inpWidth };
auto allocator_info = MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
Value input_tensor_ = Value::CreateTensor<float>(allocator_info, input_image_.data(), input_image_.size(), input_shape_.data(), input_shape_.size());
// 开始推理
vector<Value> ort_outputs = ort_session->Run(RunOptions{nullptr}, input_names.data(), &input_tensor_, 1, output_names.data(), output_names.size()); // 开始推理
/generate proposals
vector<BoxInfo> generate_boxes;
float ratioh = (float)frame.rows / this->inpHeight, ratiow = (float)frame.cols / this->inpWidth;
int n = 0, k = 0; ///cx,cy,w,h,box_score, class_score
const float* pdata = ort_outputs[0].GetTensorMutableData<float>();
for (n = 0; n < this->num_proposal; n++) ///特征图尺度
{
float box_score = pdata[4];
if (box_score > this->confThreshold)
{
int max_ind = 0;
float max_class_socre = 0;
for (k = 0; k < num_class; k++)
{
if (pdata[k + 5] > max_class_socre)
{
max_class_socre = pdata[k + 5];
max_ind = k;
}
}
max_class_socre *= box_score;
if (max_class_socre > this->confThreshold)
{
float cx = pdata[0] * ratiow; ///cx
float cy = pdata[1] * ratioh; ///cy
float w = pdata[2] * ratiow; ///w
float h = pdata[3] * ratioh; ///h
float xmin = cx - 0.5 * w;
float ymin = cy - 0.5 * h;
float xmax = cx + 0.5 * w;
float ymax = cy + 0.5 * h;
generate_boxes.push_back(BoxInfo{ xmin, ymin, xmax, ymax, max_class_socre, max_ind });
}
}
pdata += nout;
}
// Perform non maximum suppression to eliminate redundant overlapping boxes with
// lower confidences
nms(generate_boxes);
for (size_t i = 0; i < generate_boxes.size(); ++i)
{
int xmin = int(generate_boxes[i].x1);
int ymin = int(generate_boxes[i].y1);
rectangle(frame, Point(xmin, ymin), Point(int(generate_boxes[i].x2), int(generate_boxes[i].y2)), Scalar(0, 0, 255), 2);
string label = format("%.2f", generate_boxes[i].score);
label = this->class_names[generate_boxes[i].label] + ":" + label;
putText(frame, label, Point(xmin, ymin - 5), FONT_HERSHEY_SIMPLEX, 0.75, Scalar(0, 255, 0), 1);
}
}
int main()
{
Net_config YOLOV7_nets = { 0.3, 0.5, "D:/models/best.onnx" }; choices=["models/yolov7_640x640.onnx", "models/yolov7-tiny_640x640.onnx", "models/yolov7_736x1280.onnx", "models/yolov7-tiny_384x640.onnx", "models/yolov7_480x640.onnx", "models/yolov7_384x640.onnx", "models/yolov7-tiny_256x480.onnx", "models/yolov7-tiny_256x320.onnx", "models/yolov7_256x320.onnx", "models/yolov7-tiny_256x640.onnx", "models/yolov7_256x640.onnx", "models/yolov7-tiny_480x640.onnx", "models/yolov7-tiny_736x1280.onnx", "models/yolov7_256x480.onnx"]
YOLOV7 net(YOLOV7_nets);
string imgpath = "D:/images/7.jpg";
Mat srcimg = imread(imgpath);
net.detect(srcimg);
static const string kWinName = "Deep learning object detection in ONNXRuntime";
namedWindow(kWinName, WINDOW_NORMAL);
imshow(kWinName, srcimg);
waitKey(0);
destroyAllWindows();
}
类似资料：