MCPcopy Create free account
hub / github.com/FeiYull/TensorRT-Alpha / postprocess

Method postprocess

utils/yolo.cpp:205–236  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

203}
204
205void yolo::YOLO::postprocess(const std::vector<cv::Mat>& imgsBatch)
206{
207 decodeDevice(m_param, m_output_src_device, 5 + m_param.num_class, m_total_objects, m_output_area,
208 m_output_objects_device, m_output_objects_width, m_param.topK);
209
210 // nmsv1(nms faster)
211 nmsDeviceV1(m_param, m_output_objects_device, m_output_objects_width, m_param.topK, m_param.topK * m_output_objects_width + 1);
212
213 // nmsv2(nms sort)
214 //nmsDeviceV2(m_param, m_output_objects_device, m_output_objects_width, m_param.topK, m_param.topK * m_output_objects_width + 1, m_output_idx_device, m_output_conf_device);
215
216 CHECK(cudaMemcpy(m_output_objects_host, m_output_objects_device, m_param.batch_size * sizeof(float) * (1 + 7 * m_param.topK), cudaMemcpyDeviceToHost));
217 for (size_t bi = 0; bi < imgsBatch.size(); bi++)
218 {
219 int num_boxes = std::min((int)(m_output_objects_host + bi * (m_param.topK * m_output_objects_width + 1))[0], m_param.topK);
220 for (size_t i = 0; i < num_boxes; i++)
221 {
222 float* ptr = m_output_objects_host + bi * (m_param.topK * m_output_objects_width + 1) + m_output_objects_width * i + 1;
223 int keep_flag = ptr[6];
224 if (keep_flag)
225 {
226 float x_lt = m_dst2src.v0 * ptr[0] + m_dst2src.v1 * ptr[1] + m_dst2src.v2;
227 float y_lt = m_dst2src.v3 * ptr[0] + m_dst2src.v4 * ptr[1] + m_dst2src.v5;
228 float x_rb = m_dst2src.v0 * ptr[2] + m_dst2src.v1 * ptr[3] + m_dst2src.v2;
229 float y_rb = m_dst2src.v3 * ptr[2] + m_dst2src.v4 * ptr[3] + m_dst2src.v5;
230 m_objectss[bi].emplace_back(x_lt, y_lt, x_rb, y_rb, ptr[4], (int)ptr[5]);
231 }
232 }
233
234 }
235
236}
237
238std::vector<std::vector<utils::Box>> yolo::YOLO::getObjectss() const
239{

Callers 2

taskFunction · 0.45
taskFunction · 0.45

Calls

no outgoing calls

Tested by

no test coverage detected