From 1e1cf6a6a0b66e1f4f7fd5561c5b87ba430ebbbe Mon Sep 17 00:00:00 2001
From: lltcggie <lltcggie@gmail.com>
Date: Sun, 6 Dec 2015 05:22:47 +0900
Subject: [PATCH] =?UTF-8?q?=E3=82=A2=E3=83=AB=E3=83=95=E3=82=A1=E3=83=81?=
 =?UTF-8?q?=E3=83=A3=E3=83=B3=E3=83=8D=E3=83=AB=E4=BB=98=E3=81=8D=E7=94=BB?=
 =?UTF-8?q?=E5=83=8F=E3=81=A7=E9=80=8F=E6=98=8E=E9=A0=98=E5=9F=9F=E5=86=85?=
 =?UTF-8?q?=E3=81=AE=E5=BC=95=E3=81=8D=E4=BC=B8=E3=81=B0=E3=81=97=E3=81=9F?=
 =?UTF-8?q?=E9=83=A8=E5=88=86=E3=81=8C=E3=81=86=E3=81=BE=E3=81=8F=E6=B6=88?=
 =?UTF-8?q?=E3=81=88=E3=81=A6=E3=81=84=E3=81=AA=E3=81=8B=E3=81=A3=E3=81=9F?=
 =?UTF-8?q?=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/waifu2x.cpp | 28 +++++++++++++++++++---------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/common/waifu2x.cpp b/common/waifu2x.cpp
index 9e11ed1..0a3b750 100644
--- a/common/waifu2x.cpp
+++ b/common/waifu2x.cpp
@@ -1342,21 +1342,13 @@ Waifu2x::eWaifu2xError Waifu2x::AfterReconstructFloatMatProcess(const bool isRec
 		Reconstruct(false, true, cancel_func, planes[3], alpha);
 	}
 
-	// アルファチャンネルがあったらアルファを付加して、完全透明のピクセルの色を消す(処理の都合上、完全透明のピクセルにも色を付けたから)
+	// アルファチャンネルがあったらアルファを付加する
 	if (!alpha.empty())
 	{
 		std::vector<cv::Mat> planes;
 		cv::split(process_image, planes);
 		process_image.release();
 
-		cv::Mat mask;
-		cv::threshold(alpha, mask, 0.0, 1.0, cv::THRESH_BINARY); // アルファチャンネルを二値化してマスクとして扱う
-
-		// アルファチャンネルが0のところの色を消す
-		planes[0] = planes[0].mul(mask);
-		planes[1] = planes[1].mul(mask);
-		planes[2] = planes[2].mul(mask);
-
 		planes.push_back(alpha);
 
 		cv::merge(planes, process_image);
@@ -1410,6 +1402,24 @@ Waifu2x::eWaifu2xError Waifu2x::waifu2x(const std::string &input_file, const std
 	process_image.convertTo(write_iamge, CV_8U, 255.0);
 	process_image.release();
 
+	// 完全透明のピクセルの色を消す(処理の都合上、完全透明のピクセルにも色を付けたから)
+	// モデルによっては画像全域の完全透明の場所にごく小さい値のアルファが広がることがある。それを消すためにuint8_tに変換してからこの処理を行うことにした
+	if (write_iamge.channels() > 3)
+	{
+		std::vector<cv::Mat> planes;
+		cv::split(write_iamge, planes);
+
+		cv::Mat mask;
+		cv::threshold(planes[3], mask, 0.0, 1.0, cv::THRESH_BINARY); // アルファチャンネルを二値化してマスクとして扱う
+
+		// アルファチャンネルが0のところの色を消す
+		planes[0] = planes[0].mul(mask);
+		planes[1] = planes[1].mul(mask);
+		planes[2] = planes[2].mul(mask);
+
+		cv::merge(planes, write_iamge);
+	}
+
 	ret = WriteMat(write_iamge, output_file);
 	if (ret != eWaifu2xError_OK)
 		return ret;