├── README.md
├── images
    ├── 0.jpg
    ├── 1.jpg
    ├── 2.jpg
    └── 3.jpg
├── log_space.bin
├── lstr_360x640.onnx
├── main.cpp
└── main.py


/README.md:
--------------------------------------------------------------------------------
1 | # LSTR-lane-detect-onnxrun-cpp-py
2 | 使用ONNXRuntime部署LSTR基于Transformer的端到端实时车道线检测，包含C++和Python两个版本的程序。
3 | 
4 | onnx文件的大小只有2.93M，可以做到实时性轻量化部署。
5 | 起初，我想使用opencv做部署的，但是opencv的dnn模块读取onnx文件出错， 无赖只能使用onnxruntime做部署了。
6 | 


--------------------------------------------------------------------------------
/images/0.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/hpc203/LSTR-lane-detect-onnxrun-cpp-py/8d82a8a273dc71e73e03d7060a7479820a8e54eb/images/0.jpg


--------------------------------------------------------------------------------
/images/1.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/hpc203/LSTR-lane-detect-onnxrun-cpp-py/8d82a8a273dc71e73e03d7060a7479820a8e54eb/images/1.jpg


--------------------------------------------------------------------------------
/images/2.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/hpc203/LSTR-lane-detect-onnxrun-cpp-py/8d82a8a273dc71e73e03d7060a7479820a8e54eb/images/2.jpg


--------------------------------------------------------------------------------
/images/3.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/hpc203/LSTR-lane-detect-onnxrun-cpp-py/8d82a8a273dc71e73e03d7060a7479820a8e54eb/images/3.jpg


--------------------------------------------------------------------------------
/log_space.bin:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/hpc203/LSTR-lane-detect-onnxrun-cpp-py/8d82a8a273dc71e73e03d7060a7479820a8e54eb/log_space.bin


--------------------------------------------------------------------------------
/lstr_360x640.onnx:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/hpc203/LSTR-lane-detect-onnxrun-cpp-py/8d82a8a273dc71e73e03d7060a7479820a8e54eb/lstr_360x640.onnx


--------------------------------------------------------------------------------
/main.cpp:
--------------------------------------------------------------------------------
  1 | #define _CRT_SECURE_NO_WARNINGS
  2 | #include <iostream>
  3 | #include <fstream>
  4 | #include <string>
  5 | #include <opencv2/imgproc.hpp>
  6 | #include <opencv2/highgui.hpp>
  7 | //#include <cuda_provider_factory.h>
  8 | #include <onnxruntime_cxx_api.h>
  9 | 
 10 | using namespace cv;
 11 | using namespace std;
 12 | using namespace Ort;
 13 | 
 14 | class LSTR
 15 | {
 16 | public:
 17 | 	LSTR();
 18 | 	Mat detect(Mat& cv_image);
 19 | 	~LSTR();  // 析构函数, 释放内存
 20 | 	
 21 | private:
 22 | 	void normalize_(Mat img);
 23 | 	int inpWidth;
 24 | 	int inpHeight;
 25 | 	vector<float> input_image_;
 26 | 	vector<float> mask_tensor;
 27 | 	float mean[3] = { 0.485, 0.456, 0.406 };
 28 | 	float std[3] = { 0.229, 0.224, 0.225 };
 29 | 	const int len_log_space = 50;
 30 | 	float* log_space;
 31 | 	const Scalar lane_colors[8] = { Scalar(68,65,249), Scalar(44,114,243),Scalar(30,150,248),Scalar(74,132,249),Scalar(79,199,249),Scalar(109,190,144),Scalar(142, 144, 77),Scalar(161, 125, 39) };
 32 | 
 33 | 	Env env = Env(ORT_LOGGING_LEVEL_ERROR, "LSTR");
 34 | 	Ort::Session *ort_session = nullptr;
 35 | 	SessionOptions sessionOptions = SessionOptions();
 36 | 	vector<char*> input_names;
 37 | 	vector<char*> output_names;
 38 | 	vector<vector<int64_t>> input_node_dims; // >=1 outputs
 39 | 	vector<vector<int64_t>> output_node_dims; // >=1 outputs
 40 | };
 41 | 
 42 | LSTR::LSTR()
 43 | {
 44 | 	string model_path = "lstr_360x640.onnx";
 45 | 	std::wstring widestr = std::wstring(model_path.begin(), model_path.end());
 46 | 	//OrtStatus* status = OrtSessionOptionsAppendExecutionProvider_CUDA(sessionOptions, 0);
 47 | 	sessionOptions.SetGraphOptimizationLevel(ORT_ENABLE_BASIC);
 48 | 	ort_session = new Session(env, widestr.c_str(), sessionOptions);
 49 | 	size_t numInputNodes = ort_session->GetInputCount();
 50 | 	size_t numOutputNodes = ort_session->GetOutputCount();
 51 | 	AllocatorWithDefaultOptions allocator;
 52 | 	for (int i = 0; i < numInputNodes; i++)
 53 | 	{
 54 | 		input_names.push_back(ort_session->GetInputName(i, allocator));
 55 | 		Ort::TypeInfo input_type_info = ort_session->GetInputTypeInfo(i);
 56 | 		auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo();
 57 | 		auto input_dims = input_tensor_info.GetShape();
 58 | 		input_node_dims.push_back(input_dims);
 59 | 	}
 60 | 	for (int i = 0; i < numOutputNodes; i++)
 61 | 	{
 62 | 		output_names.push_back(ort_session->GetOutputName(i, allocator));
 63 | 		Ort::TypeInfo output_type_info = ort_session->GetOutputTypeInfo(i);
 64 | 		auto output_tensor_info = output_type_info.GetTensorTypeAndShapeInfo();
 65 | 		auto output_dims = output_tensor_info.GetShape();
 66 | 		output_node_dims.push_back(output_dims);
 67 | 	}
 68 | 	this->inpHeight = input_node_dims[0][2];
 69 | 	this->inpWidth = input_node_dims[0][3];
 70 | 	this->mask_tensor.resize(this->inpHeight * this->inpWidth, 0.0);
 71 | 	log_space = new float[len_log_space];
 72 | 	FILE* fp = fopen("log_space.bin", "rb");
 73 | 	fread(log_space, sizeof(float), len_log_space, fp);//导入数据
 74 | 	fclose(fp);//关闭文件。
 75 | }
 76 | 
 77 | LSTR::~LSTR()
 78 | {
 79 | 	delete[] log_space;
 80 | 	log_space = NULL;
 81 | }
 82 | 
 83 | void LSTR::normalize_(Mat img)
 84 | {
 85 | 	//img.convertTo(img, CV_32F);
 86 | 	int row = img.rows;
 87 | 	int col = img.cols;
 88 | 	this->input_image_.resize(row * col * img.channels());
 89 | 	for (int c = 0; c < 3; c++)
 90 | 	{
 91 | 		for (int i = 0; i < row; i++)
 92 | 		{
 93 | 			for (int j = 0; j < col; j++)
 94 | 			{
 95 | 				float pix = img.ptr<uchar>(i)[j * 3 + c];
 96 | 				this->input_image_[c * row * col + i * col + j] = (pix / 255.0 - mean[c]) / std[c];
 97 | 			}
 98 | 		}
 99 | 	}
100 | }
101 | 
102 | Mat LSTR::detect(Mat& srcimg)
103 | {
104 | 	const int img_height = srcimg.rows;
105 | 	const int img_width = srcimg.cols;
106 | 	Mat dstimg;
107 | 	resize(srcimg, dstimg, Size(this->inpWidth, this->inpHeight), INTER_LINEAR);
108 | 	this->normalize_(dstimg);
109 | 	array<int64_t, 4> input_shape_{ 1, 3, this->inpHeight, this->inpWidth };
110 | 	array<int64_t, 4> mask_shape_{ 1, 1, this->inpHeight, this->inpWidth };
111 | 
112 | 	auto allocator_info = MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
113 | 	vector<Value> ort_inputs;
114 | 	ort_inputs.push_back(Value::CreateTensor<float>(allocator_info, input_image_.data(), input_image_.size(), input_shape_.data(), input_shape_.size()));
115 | 	ort_inputs.push_back(Value::CreateTensor<float>(allocator_info, mask_tensor.data(), mask_tensor.size(), mask_shape_.data(), mask_shape_.size()));
116 | 	// ¿ªÊ¼ÍÆÀí
117 | 	vector<Value> ort_outputs = ort_session->Run(RunOptions{ nullptr }, input_names.data(), ort_inputs.data(), 2, output_names.data(), output_names.size());
118 | 	const float* pred_logits = ort_outputs[0].GetTensorMutableData<float>();
119 | 	const float* pred_curves = ort_outputs[1].GetTensorMutableData<float>();
120 | 	const int logits_h = output_node_dims[0][1];
121 | 	const int logits_w = output_node_dims[0][2];
122 | 	const int curves_w = output_node_dims[1][2];
123 | 	vector<int> good_detections;
124 | 	vector< vector<Point>> lanes;
125 | 	for (int i = 0; i < logits_h; i++)
126 | 	{
127 | 		float max_logits = -10000;
128 | 		int max_id = -1;
129 | 		for (int j = 0; j < logits_w; j++)
130 | 		{
131 | 			const float data = pred_logits[i*logits_w + j];
132 | 			if (data > max_logits)
133 | 			{
134 | 				max_logits = data;
135 | 				max_id = j;
136 | 			}
137 | 		}
138 | 		if (max_id == 1)
139 | 		{
140 | 			good_detections.push_back(i);
141 | 			const float *p_lane_data = pred_curves + i * curves_w;
142 | 			vector<Point> lane_points(len_log_space);
143 | 			for (int k = 0; k < len_log_space; k++)
144 | 			{
145 | 				const float y = p_lane_data[0] + log_space[k] * (p_lane_data[1] - p_lane_data[0]);
146 | 				const float x = p_lane_data[2] / powf(y - p_lane_data[3], 2.0) + p_lane_data[4] / (y - p_lane_data[3]) + p_lane_data[5] + p_lane_data[6] * y - p_lane_data[7];
147 | 				lane_points[k] = Point(int(x*img_width), int(y*img_height));
148 | 			}
149 | 			lanes.push_back(lane_points);
150 | 		}
151 | 	}
152 | 
153 | 	/// draw lines
154 | 	vector<int> right_lane;
155 | 	vector<int> left_lane;
156 | 	for (int i = 0; i < good_detections.size(); i++)
157 | 	{
158 | 		if (good_detections[i] == 0)
159 | 		{
160 | 			right_lane.push_back(i);
161 | 		}
162 | 		if (good_detections[i] == 5)
163 | 		{
164 | 			left_lane.push_back(i);
165 | 		}
166 | 	}
167 | 	Mat visualization_img = srcimg.clone();
168 | 	if (right_lane.size() == left_lane.size())
169 | 	{
170 | 		Mat lane_segment_img = visualization_img.clone();
171 | 		vector<Point> points = lanes[right_lane[0]];
172 | 		reverse(points.begin(), points.end());
173 | 		points.insert(points.begin(), lanes[left_lane[0]].begin(), lanes[left_lane[0]].end());
174 | 		fillConvexPoly(lane_segment_img, points, Scalar(0, 191, 255));
175 | 		addWeighted(visualization_img, 0.7, lane_segment_img, 0.3, 0, visualization_img);
176 | 	}
177 | 	for (int i = 0; i < lanes.size(); i++)
178 | 	{
179 | 		for (int j = 0; j < lanes[i].size(); j++)
180 | 		{
181 | 			circle(visualization_img, lanes[i][j], 3, lane_colors[good_detections[i]], -1);
182 | 		}
183 | 	}
184 | 	return visualization_img;
185 | }
186 | 
187 | int main()
188 | {
189 | 	LSTR mynet;
190 | 	string imgpath = "images/0.jpg";
191 | 	Mat srcimg = imread(imgpath);
192 | 	Mat dstimg = mynet.detect(srcimg);
193 | 
194 | 	static const string kWinName = "Deep learning lane detection in ONNXRuntime";
195 | 	namedWindow(kWinName, WINDOW_NORMAL);
196 | 	imshow(kWinName, dstimg);
197 | 	waitKey(0);
198 | 	destroyAllWindows();
199 | }


--------------------------------------------------------------------------------
/main.py:
--------------------------------------------------------------------------------
  1 | import argparse
  2 | import cv2
  3 | import numpy as np
  4 | import onnxruntime as ort
  5 | 
  6 | lane_colors = [(68,65,249),(44,114,243),(30,150,248),(74,132,249),(79,199,249),(109,190,144),(142, 144, 77),(161, 125, 39)]
  7 | log_space = np.logspace(0,2, 50, base=1/10, endpoint=True)
  8 | 
  9 | class LSTR():
 10 |     def __init__(self):
 11 |         so = ort.SessionOptions()
 12 |         so.log_severity_level = 3
 13 |         self.session = ort.InferenceSession("lstr_360x640.onnx", so)
 14 |         model_inputs = self.session.get_inputs()
 15 |         self.rgb_input_name = model_inputs[0].name
 16 |         self.mask_input_name = model_inputs[1].name
 17 |         self.input_shape = model_inputs[0].shape
 18 |         self.input_height = int(self.input_shape[2])
 19 |         self.input_width = int(self.input_shape[3])
 20 | 
 21 |         model_outputs = self.session.get_outputs()
 22 |         self.output_names = [model_outputs[i].name for i in range(len(model_outputs))]
 23 | 
 24 |         self.mean = [0.485, 0.456, 0.406]
 25 |         self.std = [0.229, 0.224, 0.225]
 26 |         self.mask_tensor = np.zeros((1, 1, self.input_height, self.input_width), dtype=np.float32)
 27 | 
 28 |     def softmax(self, x):
 29 |         """Compute softmax values for each sets of scores in x."""
 30 |         e_x = np.exp(x - np.max(x))
 31 |         return e_x / e_x.sum(axis=-1).T
 32 | 
 33 |     def draw_lanes(self, input_img, detected_lanes, good_lanes):
 34 |         # Write the detected line points in the image
 35 |         visualization_img = input_img.copy()
 36 | 
 37 |         # Draw a mask for the current lane
 38 |         right_lane = np.where(good_lanes == 0)[0]
 39 |         left_lane = np.where(good_lanes == 5)[0]
 40 | 
 41 |         if (len(left_lane) and len(right_lane)):
 42 |             lane_segment_img = visualization_img.copy()
 43 | 
 44 |             points = np.vstack((detected_lanes[left_lane[0]].T, np.flipud(detected_lanes[right_lane[0]].T)))
 45 |             cv2.fillConvexPoly(lane_segment_img, points, color=(0, 191, 255))
 46 |             visualization_img = cv2.addWeighted(visualization_img, 0.7, lane_segment_img, 0.3, 0)
 47 | 
 48 |         for lane_num, lane_points in zip(good_lanes, detected_lanes):
 49 |             for lane_point in lane_points.T:
 50 |                 cv2.circle(visualization_img, (lane_point[0], lane_point[1]), 3, lane_colors[lane_num], -1)
 51 | 
 52 |         return visualization_img
 53 | 
 54 |     def detect(self, frame):
 55 |         img_height, img_width = frame.shape[:2]
 56 |         img = cv2.resize(frame, (self.input_width, self.input_height))
 57 | 
 58 |         img = (img.astype(np.float32) / 255.0 - self.mean) / self.std
 59 | 
 60 |         img = img.transpose(2, 0, 1)
 61 |         input_tensor = img[np.newaxis, :, :, :].astype(np.float32)
 62 | 
 63 |         # Inference
 64 |         outputs = self.session.run(self.output_names, {self.rgb_input_name: input_tensor, self.mask_input_name: self.mask_tensor})
 65 | 
 66 |         ## process outputs
 67 |         pred_logits = outputs[0]
 68 |         pred_curves = outputs[1]
 69 | 
 70 |         # Filter good lanes based on the probability
 71 |         prob = self.softmax(pred_logits)
 72 |         good_detections = np.where(np.argmax(prob, axis=-1) == 1)
 73 |         pred_logits = pred_logits[good_detections]
 74 |         pred_curves = pred_curves[good_detections]
 75 | 
 76 |         lanes = []
 77 |         for lane_data in pred_curves:
 78 |             bounds = lane_data[:2]
 79 |             k_2, f_2, m_2, n_1, b_2, b_3 = lane_data[2:]
 80 | 
 81 |             # Calculate the points for the lane
 82 |             # Note: the logspace is used for a visual effect, np.linspace would also work as in the original repository
 83 |             y_norm = bounds[0] + log_space * (bounds[1] - bounds[0])
 84 |             x_norm = (k_2 / (y_norm - f_2) ** 2 + m_2 / (y_norm - f_2) + n_1 + b_2 * y_norm - b_3)
 85 |             lane_points = np.vstack((x_norm * img_width, y_norm * img_height)).astype(int)
 86 | 
 87 |             lanes.append(lane_points)
 88 | 
 89 |         return lanes, good_detections[1]
 90 | 
 91 | if __name__ == '__main__':
 92 |     parser = argparse.ArgumentParser()
 93 |     parser.add_argument("--imgpath", type=str, default='images/2.jpg', help="image path")
 94 |     args = parser.parse_args()
 95 | 
 96 |     net = LSTR()
 97 |     srcimg = cv2.imread(args.imgpath)
 98 |     detected_lanes, lane_ids = net.detect(srcimg)
 99 |     dstimg = net.draw_lanes(srcimg, detected_lanes, lane_ids)
100 | 
101 |     winName = 'Deep learning lane detection in ONNXRuntime'
102 |     cv2.namedWindow(winName, cv2.WINDOW_NORMAL)
103 |     cv2.imshow(winName, dstimg)
104 |     cv2.waitKey(0)
105 |     cv2.destroyAllWindows()
106 | 


--------------------------------------------------------------------------------