From b6e104bed7fcbcfea2795547f72c1281b8a9f19f Mon Sep 17 00:00:00 2001
From: lltcggie <lltcggie@gmail.com>
Date: Wed, 3 Jun 2015 00:12:40 +0900
Subject: [PATCH] =?UTF-8?q?block=5Fsize=E3=82=92=E5=A4=89=E6=9B=B4?=
 =?UTF-8?q?=E5=87=BA=E6=9D=A5=E3=82=8B=E3=82=88=E3=81=86=E3=81=AB=E3=81=97?=
 =?UTF-8?q?=E3=81=A6=E3=81=BF=E3=81=9F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 common/waifu2x.cpp | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/common/waifu2x.cpp b/common/waifu2x.cpp
index 05d5f45..309f146 100644
--- a/common/waifu2x.cpp
+++ b/common/waifu2x.cpp
@@ -29,6 +29,10 @@ const auto offset = 0;
 const auto layer_num = 7;
 const auto output_size = block_size - offset * 2;
 
+const auto block_width_height = block_size + layer_num * 2;
+
+const auto original_width_height = 128 + layer_num * 2;
+
 const int ConvertMode = CV_RGB2YUV;
 const int ConvertInverseMode = CV_YUV2RGB;
 
@@ -311,6 +315,14 @@ eWaifu2xError ConstractNet(boost::shared_ptr<caffe::Net<float>> &net, const std:
 			else
 				layer_param->mutable_relu_param()->set_engine(caffe::ReLUParameter_Engine_CAFFE);
 		}
+		else if (type == "MemoryData")
+		{
+			if (layer_param->mutable_memory_data_param()->width() == original_width_height && layer_param->mutable_memory_data_param()->height() == original_width_height)
+			{
+				layer_param->mutable_memory_data_param()->set_width(block_width_height);
+				layer_param->mutable_memory_data_param()->set_height(block_width_height);
+			}
+		}
 	}
 
 	net = boost::shared_ptr<caffe::Net<float>>(new caffe::Net<float>(param));