diff --git a/common/waifu2x.cpp b/common/waifu2x.cpp index b650b69..5d078e2 100644 --- a/common/waifu2x.cpp +++ b/common/waifu2x.cpp @@ -26,21 +26,11 @@ #pragma comment(lib, "libprotoc.lib") #endif -// 一度に処理する画像の幅 -const int block_size = 128; -// 一度に何ブロック分処理するか -const int batch_size = 1; // 入力画像のオフセット const int offset = 0; // srcnn.prototxtで定義されたレイヤーの数 const int layer_num = 7; -const auto output_size = block_size - offset * 2; -// ネットワークに入力する画像のサイズ(出力画像の幅はlayer_num * 2だけ小さくなる) -const auto block_width_height = block_size + layer_num * 2; -// srcnn.prototxtで定義された入力する画像のサイズ -const auto original_width_height = 128 + layer_num * 2; - const int ConvertMode = CV_RGB2YUV; const int ConvertInverseMode = CV_YUV2RGB; @@ -336,8 +326,8 @@ Waifu2x::eWaifu2xError Waifu2x::ConstractNet(boost::shared_ptr { 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); + layer_param->mutable_memory_data_param()->set_width(block_size); + layer_param->mutable_memory_data_param()->set_height(block_size); } } } @@ -382,8 +372,8 @@ Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr block(input_block_plane_size * batch_size, 0.0f); std::vector dummy_data(block.size(), 0.0f); @@ -404,10 +394,10 @@ Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr= 1); diff --git a/common/waifu2x.h b/common/waifu2x.h index 2116125..2033d20 100644 --- a/common/waifu2x.h +++ b/common/waifu2x.h @@ -37,6 +37,20 @@ public: private: bool is_inited; + // 一度に処理する画像の幅 + int crop_size; + // 一度に何ブロック分処理するか + int batch_size; + + // ネットに入力する画像のサイズ + int block_size; + // ブロック変換後の出力サイズ + int output_size; + // ネットワークに入力する画像のサイズ(出力画像の幅はlayer_num * 2だけ小さくなる) + int block_width_height; + // srcnn.prototxtで定義された入力する画像のサイズ + int original_width_height; + std::string mode; int noise_level; double scale_ratio; @@ -64,7 +78,8 @@ public: // mode: noise or scale or noise_scale or auto_scale // process: cpu or gpu or cudnn - eWaifu2xError init(int argc, char** argv, const std::string &mode, const int noise_level, const double scale_ratio, const std::string &model_dir, const std::string &process); + eWaifu2xError init(int argc, char** argv, const std::string &mode, const int noise_level, const double scale_ratio, const std::string &model_dir, const std::string &process, + const int crop_size = 128, const int batch_size = 1); void destroy(); diff --git a/waifu2x-caffe-gui/Source.cpp b/waifu2x-caffe-gui/Source.cpp index 109bf74..9118efb 100644 --- a/waifu2x-caffe-gui/Source.cpp +++ b/waifu2x-caffe-gui/Source.cpp @@ -77,6 +77,9 @@ private: std::string outputExt; std::string inputFileExt; + int crop_size; + int batch_size; + std::thread processThread; std::atomic_bool cancelFlag; @@ -282,7 +285,7 @@ private: Waifu2x::eWaifu2xError ret; Waifu2x w; - ret = w.init(__argc, __argv, mode, noise_level, scale_ratio, "models", process); + ret = w.init(__argc, __argv, mode, noise_level, scale_ratio, "models", process, crop_size, batch_size); if(ret != Waifu2x::eWaifu2xError_OK) SendMessage(dh, WM_ON_WAIFU2X_ERROR, (WPARAM)&ret, 0); else @@ -397,7 +400,8 @@ private: } public: - DialogEvent() : dh(nullptr), mode("noise_scale"), noise_level(1), scale_ratio(2.0), process("gpu"), outputExt("png"), inputFileExt("png:jpg:jpeg:tif:tiff:bmp"), isLastError(false) + DialogEvent() : dh(nullptr), mode("noise_scale"), noise_level(1), scale_ratio(2.0), process("gpu"), outputExt("png"), inputFileExt("png:jpg:jpeg:tif:tiff:bmp"), + crop_size(128), batch_size(1), isLastError(false) { } diff --git a/waifu2x-caffe/Source.cpp b/waifu2x-caffe/Source.cpp index 3391123..f28d002 100644 --- a/waifu2x-caffe/Source.cpp +++ b/waifu2x-caffe/Source.cpp @@ -99,6 +99,14 @@ int main(int argc, char** argv) TCLAP::ValueArg cmdProcess("p", "process", "process mode", false, "gpu", &cmdProcessConstraint, cmd); + TCLAP::ValueArg cmdCropSizeFile("c", "crop_size", + "input image split size", false, + 128, "int", cmd); + + TCLAP::ValueArg cmdBatchSizeFile("b", "batch_size", + "input batch size", false, + 1, "int", cmd); + // definition of command line argument : end TCLAP::Arg::enableIgnoreMismatched(); @@ -229,7 +237,8 @@ int main(int argc, char** argv) Waifu2x::eWaifu2xError ret; Waifu2x w; - ret = w.init(argc, argv, cmdMode.getValue(), cmdNRLevel.getValue(), cmdScaleRatio.getValue(), cmdModelPath.getValue(), cmdProcess.getValue()); + ret = w.init(argc, argv, cmdMode.getValue(), cmdNRLevel.getValue(), cmdScaleRatio.getValue(), cmdModelPath.getValue(), cmdProcess.getValue(), + cmdCropSizeFile.getValue(), cmdBatchSizeFile.getValue()); switch (ret) { case Waifu2x::eWaifu2xError_InvalidParameter: