From b880441b027e7dfb2523d93e6a4a8142072b0d20 Mon Sep 17 00:00:00 2001
From: lltcggie <lltcggie@gmail.com>
Date: Fri, 8 Jul 2016 23:35:51 +0900
Subject: [PATCH] =?UTF-8?q?=E3=83=87=E3=83=BC=E3=82=BF=E3=82=92=E7=9B=B4?=
 =?UTF-8?q?=E6=8E=A5=E4=B8=8E=E3=81=88=E3=82=8BAPI=E3=81=8C=E3=81=86?=
 =?UTF-8?q?=E3=81=BE=E3=81=8F=E5=8B=95=E3=81=8B=E3=81=AA=E3=81=8B=E3=81=A3?=
 =?UTF-8?q?=E3=81=9F=E3=81=AE=E3=82=92=E4=BF=AE=E6=AD=A3?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 common/stImage.cpp |  6 ++++++
 common/waifu2x.cpp | 22 +++++++++++++++++++++-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/common/stImage.cpp b/common/stImage.cpp
index 530b301..8cc55b5 100644
--- a/common/stImage.cpp
+++ b/common/stImage.cpp
@@ -380,6 +380,8 @@ Waifu2x::eWaifu2xError stImage::Load(const boost::filesystem::path &input_file)
 
 Waifu2x::eWaifu2xError stImage::Load(const void* source, const int width, const int height, const int channel, const int stride)
 {
+	Clear();
+
 	cv::Mat original_image(cv::Size(width, height), CV_MAKETYPE(CV_8U, channel), (void *)source, stride);
 
 	if (original_image.channels() >= 3) // RGBなのでBGRにする
@@ -393,6 +395,10 @@ Waifu2x::eWaifu2xError stImage::Load(const void* source, const int width, const
 	}
 
 	mOrgFloatImage = original_image;
+	mOrgChannel = original_image.channels();
+	mOrgSize = original_image.size();
+
+	mIsRequestDenoise = false;
 
 	return Waifu2x::eWaifu2xError_OK;
 }
diff --git a/common/waifu2x.cpp b/common/waifu2x.cpp
index 6325e5d..6e22a1c 100644
--- a/common/waifu2x.cpp
+++ b/common/waifu2x.cpp
@@ -5,6 +5,7 @@
 #include <cudnn.h>
 #include <mutex>
 #include <opencv2/core.hpp>
+#include <opencv2/imgproc.hpp>
 #include <tclap/CmdLine.h>
 #include <boost/filesystem.hpp>
 #include <boost/algorithm/string.hpp>
@@ -790,6 +791,18 @@ Waifu2x::eWaifu2xError Waifu2x::waifu2x(const double factor, const void* source,
 	if (!mIsInited)
 		return Waifu2x::eWaifu2xError_NotInitialized;
 
+	int cvrSetting = -1;
+	if (in_channel == 3 && out_channel == 3)
+		cvrSetting = CV_BGR2RGB;
+	else if (in_channel == 4 && out_channel == 4)
+		cvrSetting = CV_BGRA2RGBA;
+	else if (in_channel == 3 && out_channel == 4)
+		cvrSetting = CV_BGR2RGBA;
+	else if (in_channel == 4 && out_channel == 3)
+		cvrSetting = CV_BGRA2RGB;
+	else if (!(in_channel == 1 && out_channel == 1))
+		return Waifu2x::eWaifu2xError_InvalidParameter;
+
 	stImage image;
 	ret = image.Load(source, width, height, in_channel, in_stride);
 	if (ret != Waifu2x::eWaifu2xError_OK)
@@ -811,9 +824,16 @@ Waifu2x::eWaifu2xError Waifu2x::waifu2x(const double factor, const void* source,
 
 	image.Postprocess(mInputPlane, Factor, 8);
 
-	cv::Mat out_image = image.GetEndImage();
+	cv::Mat out_bgr_image = image.GetEndImage();
 	image.Clear();
 
+	cv::Mat out_image;
+	if (cvrSetting >= 0)
+		cv::cvtColor(out_bgr_image, out_image, cvrSetting); // BGRからRGBに戻す
+	else
+		out_image = out_bgr_image;
+	out_bgr_image.release();
+
 	// 出力配列へ書き込み
 	{
 		const auto width = out_image.size().width;