@@ -52,6 +52,13 @@ struct GridAndStride
52
52
int stride;
53
53
};
54
54
55
+ typedef struct Colormap_
56
+ {
57
+ int id;
58
+ std::string name;
59
+ std::vector<unsigned char > color;
60
+ } Colormap;
61
+
55
62
/* *
56
63
* @class TrtYoloX
57
64
* @brief TensorRT YOLOX for faster inference
@@ -85,7 +92,7 @@ class TrtYoloX
85
92
const bool use_gpu_preprocess = false , std::string calibration_image_list_file = std::string(),
86
93
const double norm_factor = 1.0 , [[maybe_unused]] const std::string & cache_dir = " " ,
87
94
const tensorrt_common::BatchConfig & batch_config = {1 , 1 , 1 },
88
- const size_t max_workspace_size = (1 << 30 ));
95
+ const size_t max_workspace_size = (1 << 30 ), const std::string & color_map_path = " " );
89
96
/* *
90
97
* @brief Deconstruct TrtYoloX
91
98
*/
@@ -96,7 +103,9 @@ class TrtYoloX
96
103
* @param[out] objects results for object detection
97
104
* @param[in] images batched images
98
105
*/
99
- bool doInference (const std::vector<cv::Mat> & images, ObjectArrays & objects);
106
+ bool doInference (
107
+ const std::vector<cv::Mat> & images, ObjectArrays & objects, std::vector<cv::Mat> & masks,
108
+ std::vector<cv::Mat> & color_masks);
100
109
101
110
/* *
102
111
* @brief run inference including pre-process and post-process
@@ -130,6 +139,22 @@ class TrtYoloX
130
139
*/
131
140
void printProfiling (void );
132
141
142
+ /* *
143
+ * @brief get num for multitask heads
144
+ */
145
+ int getMultitaskNum (void );
146
+
147
+ /* *
148
+ * @brief get colorized masks from index using specific colormap
149
+ * @param[out] cmask colorized mask
150
+ * @param[in] index multitask index
151
+ * @param[in] colormap colormap for masks
152
+ */
153
+ void getColorizedMask (
154
+ const std::vector<tensorrt_yolox::Colormap> & colormap, const cv::Mat & mask,
155
+ cv::Mat & colorized_mask);
156
+ inline std::vector<Colormap> getColorMap () { return sematic_color_map_; }
157
+
133
158
private:
134
159
/* *
135
160
* @brief run preprocess including resizing, letterbox, NHWC2NCHW and toFloat on CPU
@@ -177,7 +202,9 @@ class TrtYoloX
177
202
const cv::Mat & images, int batch_size, ObjectArrays & objects);
178
203
179
204
bool feedforward (const std::vector<cv::Mat> & images, ObjectArrays & objects);
180
- bool feedforwardAndDecode (const std::vector<cv::Mat> & images, ObjectArrays & objects);
205
+ bool feedforwardAndDecode (
206
+ const std::vector<cv::Mat> & images, ObjectArrays & objects, std::vector<cv::Mat> & masks,
207
+ std::vector<cv::Mat> & color_masks);
181
208
void decodeOutputs (float * prob, ObjectArray & objects, float scale, cv::Size & img_size) const ;
182
209
void generateGridsAndStride (
183
210
const int target_w, const int target_h, const std::vector<int > & strides,
@@ -206,6 +233,26 @@ class TrtYoloX
206
233
void nmsSortedBboxes (
207
234
const ObjectArray & face_objects, std::vector<int > & picked, float nms_threshold) const ;
208
235
236
+ /* *
237
+ * @brief get a mask image for a segmentation head
238
+ * @param[out] argmax argmax results
239
+ * @param[in] prob probability map
240
+ * @param[in] dims dimension for probability map
241
+ * @param[in] out_w mask width excluding letterbox
242
+ * @param[in] out_h mask height excluding letterbox
243
+ */
244
+ cv::Mat getMaskImage (float * prob, nvinfer1::Dims dims, int out_w, int out_h);
245
+
246
+ /* *
247
+ * @brief get a mask image on GPUs for a segmentation head
248
+ * @param[out] mask image
249
+ * @param[in] prob probability map on device
250
+ * @param[in] out_w mask width excluding letterbox
251
+ * @param[in] out_h mask height excluding letterbox
252
+ * @param[in] b current batch
253
+ */
254
+ cv::Mat getMaskImageGpu (float * d_prob, nvinfer1::Dims dims, int out_w, int out_h, int b);
255
+
209
256
std::unique_ptr<tensorrt_common::TrtCommon> trt_common_;
210
257
211
258
std::vector<float > input_h_;
@@ -249,6 +296,20 @@ class TrtYoloX
249
296
CudaUniquePtrHost<Roi[]> roi_h_;
250
297
// device pointer for ROI
251
298
CudaUniquePtr<Roi[]> roi_d_;
299
+
300
+ // flag whether model has multitasks
301
+ int multitask_;
302
+ // buff size for segmentation heads
303
+ CudaUniquePtr<float []> segmentation_out_prob_d_;
304
+ CudaUniquePtrHost<float []> segmentation_out_prob_h_;
305
+ size_t segmentation_out_elem_num_;
306
+ size_t segmentation_out_elem_num_per_batch_;
307
+ std::vector<cv::Mat> segmentation_masks_;
308
+ // host buffer for argmax postprocessing on GPU
309
+ CudaUniquePtrHost<unsigned char []> argmax_buf_h_;
310
+ // device buffer for argmax postprocessing on GPU
311
+ CudaUniquePtr<unsigned char []> argmax_buf_d_;
312
+ std::vector<tensorrt_yolox::Colormap> sematic_color_map_;
252
313
};
253
314
254
315
} // namespace tensorrt_yolox
0 commit comments