From abf87fabbbcd5c91c966514ff12bb42ac15a54fe Mon Sep 17 00:00:00 2001
From: lltcggie <lltcggie@gmail.com>
Date: Fri, 29 Apr 2016 15:30:02 +0900
Subject: [PATCH] =?UTF-8?q?16bit=E5=87=BA=E5=8A=9B=E3=81=A0=E3=81=A8?=
 =?UTF-8?q?=E4=BE=8B=E5=A4=96=E3=81=8C=E7=99=BA=E7=94=9F=E3=81=99=E3=82=8B?=
 =?UTF-8?q?=E3=83=90=E3=82=B0=E3=82=92=E4=BF=AE=E6=AD=A3=20#26?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 common/waifu2x.cpp | 66 +++++++++++++++++++++++++++++++++++++++++-----
 common/waifu2x.h   |  1 +
 2 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/common/waifu2x.cpp b/common/waifu2x.cpp
index 1c18c97..06cb124 100644
--- a/common/waifu2x.cpp
+++ b/common/waifu2x.cpp
@@ -1874,6 +1874,36 @@ Waifu2x::eWaifu2xError Waifu2x::AfterReconstructFloatMatProcess(const bool isRec
 	return eWaifu2xError_OK;
 }
 
+namespace
+{
+	template<typename T>
+	void AlphaZeroToZero(std::vector<cv::Mat> &planes)
+	{
+		cv::Mat alpha(planes[3]);
+
+		const T *aptr = (const T *)alpha.data;
+
+		T *ptr0 = (T *)planes[0].data;
+		T *ptr1 = (T *)planes[1].data;
+		T *ptr2 = (T *)planes[2].data;
+
+		const size_t Line = alpha.step1();
+		const size_t Width = alpha.size().width;
+		const size_t Height = alpha.size().height;
+
+		for (size_t i = 0; i < Height; i++)
+		{
+			for (size_t j = 0; j < Width; j++)
+			{
+				const size_t pos = Line * i + j;
+
+				if (aptr[pos] == (T)0)
+					ptr0[pos] = ptr1[pos] = ptr2[pos] = (T)0;
+			}
+		}
+	}
+}
+
 Waifu2x::eWaifu2xError Waifu2x::waifu2xConvetedMat(const bool isJpeg, const cv::Mat &inMat, cv::Mat &outMat, const waifu2xCancelFunc cancel_func)
 {
 	Waifu2x::eWaifu2xError ret;
@@ -1895,6 +1925,13 @@ Waifu2x::eWaifu2xError Waifu2x::waifu2xConvetedMat(const bool isJpeg, const cv::
 	const double max_val = GetValumeMaxFromCVDepth(cv_depth);
 	const double eps = GetEPS(cv_depth);
 
+	{
+		std::vector<cv::Mat> planes;
+		cv::split(process_image, planes);
+		float *ptr = (float *)planes[3].data;
+		planes[0].release();
+	}
+
 	cv::Mat write_iamge;
 	if (output_depth != 32) // 出力がfloat形式なら変換しない
 		process_image.convertTo(write_iamge, cv_depth, max_val, eps);
@@ -1906,18 +1943,35 @@ Waifu2x::eWaifu2xError Waifu2x::waifu2xConvetedMat(const bool isJpeg, const cv::
 	// 完全透明のピクセルの色を消す(処理の都合上、完全透明のピクセルにも色を付けたから)
 	// モデルによっては画像全域の完全透明の場所にごく小さい値のアルファが広がることがある。それを消すためにcv_depthへ変換してからこの処理を行うことにした
 	// (ただしcv_depthが32の場合だと意味は無いが)
+	// TODO: モデル(例えばPhoto)によっては0しかない画像を変換しても0.000114856390とかになるので、適切な値のクリッピングを行う？
 	if (write_iamge.channels() > 3)
 	{
 		std::vector<cv::Mat> planes;
 		cv::split(write_iamge, planes);
+		write_iamge.release();
 
-		cv::Mat mask;
-		cv::threshold(planes[3], mask, 0.0, 1.0, cv::THRESH_BINARY); // アルファチャンネルを二値化してマスクとして扱う
+		const auto depth = planes[0].depth();
+		switch (depth)
+		{
+		case CV_8U:
+			AlphaZeroToZero<uint8_t>(planes);
+			break;
 
-		// アルファチャンネルが0のところの色を消す
-		planes[0] = planes[0].mul(mask);
-		planes[1] = planes[1].mul(mask);
-		planes[2] = planes[2].mul(mask);
+		case CV_16U:
+			AlphaZeroToZero<uint16_t>(planes);
+			break;
+
+		case CV_32F:
+			AlphaZeroToZero<float>(planes);
+			break;
+
+		case CV_64F:
+			AlphaZeroToZero<double>(planes);
+			break;
+
+		default:
+			return eWaifu2xError_FailedUnknownType;
+		}
 
 		cv::merge(planes, write_iamge);
 	}
diff --git a/common/waifu2x.h b/common/waifu2x.h
index 0ab4ff2..aaff306 100644
--- a/common/waifu2x.h
+++ b/common/waifu2x.h
@@ -38,6 +38,7 @@ public:
 		eWaifu2xError_FailedConstructModel,
 		eWaifu2xError_FailedProcessCaffe,
 		eWaifu2xError_FailedCudaCheck,
+		eWaifu2xError_FailedUnknownType,
 	};
 
 	enum eWaifu2xCudaError