@@ -22,6 +22,7 @@ class PrepareImage : public Transform {
2222 explicit PrepareImage (const Value& args) {
2323 to_float32_ = args.value (" to_float32" , to_float32_);
2424 color_type_ = args.value (" color_type" , color_type_);
25+ channel_order_ = args.value (" channel_order" , channel_order_);
2526
2627 cvt_color_ = operation::Managed<CvtColor>::Create ();
2728 to_float_ = operation::Managed<ToFloat>::Create ();
@@ -59,7 +60,11 @@ class PrepareImage : public Transform {
5960 Mat src_mat = data[" ori_img" ].get <Mat>();
6061 Mat dst_mat;
6162 if (color_type_ == " color" || color_type_ == " color_ignore_orientation" ) {
62- OUTCOME_TRY (cvt_color_.Apply (src_mat, dst_mat, PixelFormat::kBGR ));
63+ if (channel_order_ == " bgr" ) {
64+ OUTCOME_TRY (cvt_color_.Apply (src_mat, dst_mat, PixelFormat::kBGR ));
65+ } else {
66+ OUTCOME_TRY (cvt_color_.Apply (src_mat, dst_mat, PixelFormat::kRGB ));
67+ }
6368 } else {
6469 OUTCOME_TRY (cvt_color_.Apply (src_mat, dst_mat, PixelFormat::kGRAYSCALE ));
6570 }
@@ -93,6 +98,7 @@ class PrepareImage : public Transform {
9398 operation::Managed<ToFloat> to_float_;
9499 bool to_float32_{false };
95100 std::string color_type_{" color" };
101+ std::string channel_order_{" bgr" };
96102};
97103
98104MMDEPLOY_REGISTER_TRANSFORM2 (PrepareImage, (LoadImageFromFile, 0 ));
0 commit comments