From 2dbf72d191238179f4885328dac8d87dcd3128d7 Mon Sep 17 00:00:00 2001
From: nagadomi <nagadomi@nurs.or.jp>
Date: Mon, 21 Mar 2016 18:05:31 +0900
Subject: [PATCH] Update models for latest caffe master

waifu2x-caffe(CUI) works on ubuntu.

# Conflicts:
#	.gitignore
---
 bin/models/anime_style_art/srcnn.prototxt     | 11 ++++++-----
 bin/models/anime_style_art_rgb/srcnn.prototxt | 11 ++++++-----
 bin/models/photo/srcnn.prototxt               | 11 ++++++-----
 bin/models/ukbench/srcnn.prototxt             | 11 ++++++-----
 common/waifu2x.cpp                            | 15 +++++++--------
 5 files changed, 31 insertions(+), 28 deletions(-)

diff --git a/bin/models/anime_style_art/srcnn.prototxt b/bin/models/anime_style_art/srcnn.prototxt
index b91d3b3..06adcec 100644
--- a/bin/models/anime_style_art/srcnn.prototxt
+++ b/bin/models/anime_style_art/srcnn.prototxt
@@ -1,9 +1,10 @@
 name: "srcnn"
-input: "input"
-input_dim: 1
-input_dim: 1
-input_dim: 142
-input_dim: 142
+layer {
+  name: "input"
+  type: "Input"
+  top: "input"
+  input_param { shape: { dim: 1 dim: 1 dim: 142 dim: 142 } }
+}
 layer {
   name: "conv1_layer"
   type: "Convolution"
diff --git a/bin/models/anime_style_art_rgb/srcnn.prototxt b/bin/models/anime_style_art_rgb/srcnn.prototxt
index b3c0ae7..430ce6e 100644
--- a/bin/models/anime_style_art_rgb/srcnn.prototxt
+++ b/bin/models/anime_style_art_rgb/srcnn.prototxt
@@ -1,9 +1,10 @@
 name: "srcnn"
-input: "input"
-input_dim: 1
-input_dim: 3
-input_dim: 142
-input_dim: 142
+layer {
+  name: "input"
+  type: "Input"
+  top: "input"
+  input_param { shape: { dim: 1 dim: 3 dim: 142 dim: 142 } }
+}
 layer {
   name: "conv1_layer"
   type: "Convolution"
diff --git a/bin/models/photo/srcnn.prototxt b/bin/models/photo/srcnn.prototxt
index b3c0ae7..430ce6e 100644
--- a/bin/models/photo/srcnn.prototxt
+++ b/bin/models/photo/srcnn.prototxt
@@ -1,9 +1,10 @@
 name: "srcnn"
-input: "input"
-input_dim: 1
-input_dim: 3
-input_dim: 142
-input_dim: 142
+layer {
+  name: "input"
+  type: "Input"
+  top: "input"
+  input_param { shape: { dim: 1 dim: 3 dim: 142 dim: 142 } }
+}
 layer {
   name: "conv1_layer"
   type: "Convolution"
diff --git a/bin/models/ukbench/srcnn.prototxt b/bin/models/ukbench/srcnn.prototxt
index b3c0ae7..430ce6e 100644
--- a/bin/models/ukbench/srcnn.prototxt
+++ b/bin/models/ukbench/srcnn.prototxt
@@ -1,9 +1,10 @@
 name: "srcnn"
-input: "input"
-input_dim: 1
-input_dim: 3
-input_dim: 142
-input_dim: 142
+layer {
+  name: "input"
+  type: "Input"
+  top: "input"
+  input_param { shape: { dim: 1 dim: 3 dim: 142 dim: 142 } }
+}
 layer {
   name: "conv1_layer"
   type: "Convolution"
diff --git a/common/waifu2x.cpp b/common/waifu2x.cpp
index 9b8966a..e814b76 100644
--- a/common/waifu2x.cpp
+++ b/common/waifu2x.cpp
@@ -715,7 +715,7 @@ Waifu2x::eWaifu2xError Waifu2x::ConstractNet(boost::shared_ptr<caffe::Net<float>
 		net = boost::shared_ptr<caffe::Net<float>>(new caffe::Net<float>(param_model));
 		net->CopyTrainedLayersFrom(param_caffemodel);
 
-		input_plane = param_model.input_dim(1);
+		input_plane = param_model.layer(0).input_param().shape().Get(0).dim(1);
 	}
 	else
 	{
@@ -732,14 +732,13 @@ Waifu2x::eWaifu2xError Waifu2x::SetParameter(caffe::NetParameter &param, const s
 	param.mutable_state()->set_phase(caffe::TEST);
 
 	{
-		auto mid = param.mutable_input_dim();
-
-		if (mid->size() != 4)
+		auto input_layer = param.mutable_layer(0);
+		auto mid = input_layer->mutable_input_param()->mutable_shape();
+		if (mid->size() > 0 && mid->Mutable(0)->dim_size() != 4)
 			return eWaifu2xError_FailedParseModelFile;
-
-		*mid->Mutable(0) = batch_size;
-		*mid->Mutable(2) = input_block_size;
-		*mid->Mutable(3) = input_block_size;
+		mid->Mutable(0)->set_dim(0, batch_size);
+		mid->Mutable(0)->set_dim(2, input_block_size);
+		mid->Mutable(0)->set_dim(3, input_block_size);
 	}
 
 	for (int i = 0; i < param.layer_size(); i++)