noise_scaleとnoiseを両方持っているモデルに対応

lltcggie · Nov 23, 2018 · e5c310f · e5c310f
1 parent 8b53d0f
commit e5c310f
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 5 deletions.
diff --git a/bin/models/cunet/info.json b/bin/models/cunet/info.json
@@ -1,4 +1,4 @@
-{"name":"CUnet","arch_name":"upcunet","has_noise_scale":true,"channels":3,
+{"name":"CUnet","arch_name":"upcunet","has_noise_scale":true,"has_noise_only":true, "channels":3,
 "scale_factor":2,"offset":36,
 "scale_factor_noise":1,"offset_noise":28
 }
diff --git a/common/cNet.cpp b/common/cNet.cpp
@@ -190,11 +190,13 @@ Waifu2x::eWaifu2xError cNet::GetInfo(const boost::filesystem::path & info_path,
 		const auto name = d["name"].GetString();
 		const auto arch_name = d["arch_name"].GetString();
 		const bool has_noise_scale = d.HasMember("has_noise_scale") && d["has_noise_scale"].GetBool() ? true : false;
+		const bool has_noise_only = d.HasMember("has_noise_only") && d["has_noise_only"].GetBool() ? true : false;
 		const int channels = d["channels"].GetInt();
 
 		info.name = name;
 		info.arch_name = arch_name;
 		info.has_noise_scale = has_noise_scale;
+		info.has_noise_only = has_noise_only;
 		info.channels = channels;
 
 		if (d.HasMember("offset"))

diff --git a/common/waifu2x.cpp b/common/waifu2x.cpp
@@ -655,7 +655,7 @@ Waifu2x::eWaifu2xError Waifu2x::Init(const eWaifu2xModelType mode, const int noi
 		if (ret != Waifu2x::eWaifu2xError_OK)
 			return ret;
 
-		mHasNoiseScale = info.has_noise_scale;
+		mHasNoiseScaleOnly = info.has_noise_scale && !info.has_noise_only;
 		mInputPlane = info.channels;
 
 		if (mode == eWaifu2xModelTypeNoise || mode == eWaifu2xModelTypeNoiseScale || mode == eWaifu2xModelTypeAutoScale)
@@ -665,7 +665,7 @@ Waifu2x::eWaifu2xError Waifu2x::Init(const eWaifu2xModelType mode, const int noi
 			mNoiseNet.reset(new cNet);
 
 			eWaifu2xModelType Mode = mode;
-			if (info.has_noise_scale) // ノイズ除去と拡大を同時に行う
+			if (mHasNoiseScaleOnly) // ノイズ除去と拡大を同時に行う
 			{
 				// ノイズ除去拡大ネットの構築はeWaifu2xModelTypeNoiseScaleを指定する必要がある
 				Mode = eWaifu2xModelTypeNoiseScale;
@@ -905,7 +905,7 @@ Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(const Factor factor, const int
 
 	if (isReconstructNoise)
 	{
-		if (!mHasNoiseScale) // ノイズ除去だけ
+		if (!mHasNoiseScaleOnly) // ノイズ除去だけ
 		{
 			cv::Mat im;
 			cv::Size_<int> size;

diff --git a/common/waifu2x.h b/common/waifu2x.h
@@ -74,6 +74,7 @@ class Waifu2x
 		std::string name;
 		std::string arch_name;
 		bool has_noise_scale;
+		bool has_noise_only;
 		int channels;
 
 		stParam noise;
@@ -141,7 +142,7 @@ class Waifu2x
 
 	int mInputPlane; // ネットへの入力チャンネル数
 	int mMaxNetOffset; // ネットに入力するとどれくらい削れるか
-	bool mHasNoiseScale;
+	bool mHasNoiseScaleOnly;
 
 	float *mOutputBlock;
 	size_t mOutputBlockSize;