From 8478249ed8112444d2aef5e58cb57a80b045d5a6 Mon Sep 17 00:00:00 2001 From: Chen Xin Date: Wed, 13 Sep 2023 10:07:43 +0800 Subject: [PATCH] fix img load for super-resolution in sdk (#2426) * fix load for super-resolution in sdk * fix lint --- csrc/mmdeploy/preprocess/transform/load.cpp | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/csrc/mmdeploy/preprocess/transform/load.cpp b/csrc/mmdeploy/preprocess/transform/load.cpp index c386957222..88782d5c59 100644 --- a/csrc/mmdeploy/preprocess/transform/load.cpp +++ b/csrc/mmdeploy/preprocess/transform/load.cpp @@ -22,6 +22,7 @@ class PrepareImage : public Transform { explicit PrepareImage(const Value& args) { to_float32_ = args.value("to_float32", to_float32_); color_type_ = args.value("color_type", color_type_); + channel_order_ = args.value("channel_order", channel_order_); cvt_color_ = operation::Managed::Create(); to_float_ = operation::Managed::Create(); @@ -59,7 +60,11 @@ class PrepareImage : public Transform { Mat src_mat = data["ori_img"].get(); Mat dst_mat; if (color_type_ == "color" || color_type_ == "color_ignore_orientation") { - OUTCOME_TRY(cvt_color_.Apply(src_mat, dst_mat, PixelFormat::kBGR)); + if (channel_order_ == "bgr") { + OUTCOME_TRY(cvt_color_.Apply(src_mat, dst_mat, PixelFormat::kBGR)); + } else { + OUTCOME_TRY(cvt_color_.Apply(src_mat, dst_mat, PixelFormat::kRGB)); + } } else { OUTCOME_TRY(cvt_color_.Apply(src_mat, dst_mat, PixelFormat::kGRAYSCALE)); } @@ -93,6 +98,7 @@ class PrepareImage : public Transform { operation::Managed to_float_; bool to_float32_{false}; std::string color_type_{"color"}; + std::string channel_order_{"bgr"}; }; MMDEPLOY_REGISTER_TRANSFORM2(PrepareImage, (LoadImageFromFile, 0));