From 8e785648daf9ca336b4ea2d6db9d261be0fb03eb Mon Sep 17 00:00:00 2001
From: jiangzhiwen <jiangzhiwen8@huawei.com>
Date: Wed, 21 Oct 2020 14:44:22 +0800
Subject: [PATCH] add merge image processing

---
 .../kernels/image/lite_cv/image_process.cc    | 50 +++++++++++++++++++
 .../kernels/image/lite_cv/image_process.h     |  3 ++
 tests/ut/cpp/dataset/image_process_test.cc    | 28 +++++++++++
 3 files changed, 81 insertions(+)

diff --git a/mindspore/ccsrc/minddata/dataset/kernels/image/lite_cv/image_process.cc b/mindspore/ccsrc/minddata/dataset/kernels/image/lite_cv/image_process.cc
index 081d7f2122..d910fc5e44 100644
--- a/mindspore/ccsrc/minddata/dataset/kernels/image/lite_cv/image_process.cc
+++ b/mindspore/ccsrc/minddata/dataset/kernels/image/lite_cv/image_process.cc
@@ -699,6 +699,56 @@ bool Split(const LiteMat &src, std::vector<LiteMat> &mv) {
   return false;
 }
 
+template <typename T>
+inline void MergeImpl(const std::vector<LiteMat> &mv, T *dst_ptr, int height, int width, int channel) {
+  T *mv_ptr[4];
+  int area = height * width;
+  for (int c = 0; c < channel; c++) {
+    mv_ptr[c] = reinterpret_cast<T *>(mv[c].data_ptr_);
+  }
+  for (int i = 0; i < area; i++) {
+    for (int c = 0; c < channel; c++) {
+      dst_ptr[c] = *mv_ptr[c];
+      mv_ptr[c]++;
+    }
+    dst_ptr += channel;
+  }
+}
+
+bool Merge(const std::vector<LiteMat> &mv, LiteMat &dst) {
+  if (mv.empty() || mv.size() > 4) return false;
+
+  int width = mv[0].width_;
+  int height = mv[0].height_;
+  int channel = mv.size();
+  LDataType data_type = mv[0].data_type_;
+
+  // The arrays in list must be single-channel
+  for (int i = 0; i < mv.size(); i++) {
+    if (mv[i].channel_ != 1) return false;
+  }
+
+  for (int i = 1; i < mv.size(); i++) {
+    if (width != mv[i].width_ || height != mv[i].height_ || data_type != mv[i].data_type_) {
+      return false;
+    }
+  }
+
+  if (dst.IsEmpty() || dst.width_ != width || dst.height_ != height || dst.channel_ != channel ||
+      dst.data_type_ != data_type) {
+    dst.Init(width, height, channel, data_type);
+  }
+
+  if (dst.data_type_ == LDataType::FLOAT32) {
+    MergeImpl<float>(mv, dst, height, width, channel);
+  } else if (dst.data_type_ == LDataType::UINT8) {
+    MergeImpl<uint8_t>(mv, dst, height, width, channel);
+  } else {
+    return false;
+  }
+  return true;
+}
+
 bool Pad(const LiteMat &src, LiteMat &dst, int top, int bottom, int left, int right, PaddBorderType pad_type,
          uint8_t fill_b_or_gray, uint8_t fill_g, uint8_t fill_r) {
   if (top <= 0 || bottom <= 0 || left <= 0 || right <= 0) {
diff --git a/mindspore/ccsrc/minddata/dataset/kernels/image/lite_cv/image_process.h b/mindspore/ccsrc/minddata/dataset/kernels/image/lite_cv/image_process.h
index 1816fe837a..ff00080451 100644
--- a/mindspore/ccsrc/minddata/dataset/kernels/image/lite_cv/image_process.h
+++ b/mindspore/ccsrc/minddata/dataset/kernels/image/lite_cv/image_process.h
@@ -85,6 +85,9 @@ bool ExtractChannel(const LiteMat &src, LiteMat &dst, int col);
 
 bool Split(const LiteMat &src, std::vector<LiteMat> &mv);
 
+/// \brief Create a multi-channel image out of several single-channel arrays.
+bool Merge(const std::vector<LiteMat> &mv, LiteMat &dst);
+
 /// \brief Apply affine transformation for 1 channel image
 bool Affine(LiteMat &src, LiteMat &out_img, double M[6], std::vector<size_t> dsize, UINT8_C1 borderValue);
 
diff --git a/tests/ut/cpp/dataset/image_process_test.cc b/tests/ut/cpp/dataset/image_process_test.cc
index 32859e87e2..95322f66bd 100644
--- a/tests/ut/cpp/dataset/image_process_test.cc
+++ b/tests/ut/cpp/dataset/image_process_test.cc
@@ -241,6 +241,34 @@ TEST_F(MindDataImageProcess, testSplit) {
   cv::Mat dst_imageR(lite_r.height_, lite_r.width_, CV_8UC1, lite_r.data_ptr_);
 }
 
+TEST_F(MindDataImageProcess, testMerge) {
+  std::string filename = "data/dataset/apple.jpg";
+  cv::Mat src_image = cv::imread(filename, cv::ImreadModes::IMREAD_COLOR);
+  std::vector<cv::Mat> dst_images;
+  cv::split(src_image, dst_images);
+  // convert to RGBA for Android bitmap(rgba)
+  cv::Mat rgba_mat;
+  cv::cvtColor(src_image, rgba_mat, CV_BGR2RGBA);
+
+  bool ret = false;
+  LiteMat lite_mat_bgr;
+  ret =
+    InitFromPixel(rgba_mat.data, LPixelType::RGBA2BGR, LDataType::UINT8, rgba_mat.cols, rgba_mat.rows, lite_mat_bgr);
+  ASSERT_TRUE(ret == true);
+  std::vector<LiteMat> lite_all;
+  ret = Split(lite_mat_bgr, lite_all);
+  ASSERT_TRUE(ret == true);
+  ASSERT_TRUE(lite_all.size() == 3);
+  LiteMat lite_r = lite_all[2];
+  cv::Mat dst_imageR(lite_r.height_, lite_r.width_, CV_8UC1, lite_r.data_ptr_);
+
+  LiteMat merge_mat;
+  EXPECT_TRUE(Merge(lite_all, merge_mat));
+  EXPECT_EQ(merge_mat.height_, lite_mat_bgr.height_);
+  EXPECT_EQ(merge_mat.width_, lite_mat_bgr.width_);
+  EXPECT_EQ(merge_mat.channel_, lite_mat_bgr.channel_);
+}
+
 void Lite1CImageProcess(LiteMat &lite_mat_bgr, LiteMat &lite_norm_mat_cut) {
   LiteMat lite_mat_resize;
   int ret = ResizeBilinear(lite_mat_bgr, lite_mat_resize, 256, 256);