From 6d83a8038525cf64b0edd866ebbd94906eb5691d Mon Sep 17 00:00:00 2001
From: RAJKIRAN NATARAJAN <saskatchewancatch@gmail.com>
Date: Thu, 29 Mar 2018 06:47:31 -0700
Subject: [PATCH] Merge pull request #11176 from saskatchewancatch:issue-11165

* fix cv::imwrite to allow cv::UMat and vectors of it to be written

* cv::imwrite failure handling improvements

* Unit test for regression testing imwrite support for and fix whitespace
---
 modules/imgcodecs/src/loadsave.cpp         |  7 +++++--
 modules/imgcodecs/test/test_read_write.cpp | 20 ++++++++++++++++++++
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/modules/imgcodecs/src/loadsave.cpp b/modules/imgcodecs/src/loadsave.cpp
index c2042b8ae9..12852d1c48 100644
--- a/modules/imgcodecs/src/loadsave.cpp
+++ b/modules/imgcodecs/src/loadsave.cpp
@@ -717,11 +717,14 @@ bool imwrite( const String& filename, InputArray _img,
     CV_TRACE_FUNCTION();
     std::vector<Mat> img_vec;
     //Did we get a Mat or a vector of Mats?
-    if (_img.isMat())
+    if (_img.isMat() || _img.isUMat())
         img_vec.push_back(_img.getMat());
-    else if (_img.isMatVector())
+    else if (_img.isMatVector() || _img.isUMatVector())
         _img.getMatVector(img_vec);
+    else
+        CV_ErrorNoReturn(Error::StsBadArg, "Unknown/unsupported input encountered");
 
+    CV_Assert(!img_vec.empty());
     return imwrite_(filename, img_vec, params, false);
 }
 
diff --git a/modules/imgcodecs/test/test_read_write.cpp b/modules/imgcodecs/test/test_read_write.cpp
index 41765b9234..5056f4d632 100644
--- a/modules/imgcodecs/test/test_read_write.cpp
+++ b/modules/imgcodecs/test/test_read_write.cpp
@@ -127,4 +127,24 @@ TEST(Imgcodecs_Image, regression_9376)
     EXPECT_EQ(32, m.rows);
 }
 
+//==================================================================================================
+
+TEST(Imgcodecs_Image, write_umat)
+{
+    const string src_name = TS::ptr()->get_data_path() + "../python/images/baboon.bmp";
+    const string dst_name = cv::tempfile(".bmp");
+
+    Mat image1 = imread(src_name);
+    ASSERT_FALSE(image1.empty());
+
+    UMat image1_umat = image1.getUMat(ACCESS_RW);
+
+    imwrite(dst_name, image1_umat);
+
+    Mat image2 = imread(dst_name);
+    ASSERT_FALSE(image2.empty());
+
+    EXPECT_PRED_FORMAT2(cvtest::MatComparator(0, 0), image1, image2);
+}
+
 }} // namespace