Mirror/waifu2x-caffe
1
0
Fork 0
mirror of https://github.com/lltcggie/waifu2x-caffe.git synced 2025-06-26 13:42:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

crop_sizeとbatch_sizeが指定出来るようになった

Browse Source
This commit is contained in:
lltcggie 2015-06-03 03:24:01 +09:00
parent 4ddeaee0e6
commit 1320f126f9
4 changed files with 53 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
common/waifu2x.cpp
View File

@ -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<caffe::Net<float>
{
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<caffe::Net<fl
const int BlockNum = WidthNum * HeightNum;
const int input_block_plane_size = block_width_height * block_width_height;
const int output_block_plane_size = block_size * block_size;
const int input_block_plane_size = block_size * block_size;
const int output_block_plane_size = crop_size * crop_size;
std::vector<float> block(input_block_plane_size * batch_size, 0.0f);
std::vector<float> dummy_data(block.size(), 0.0f);
@ -404,10 +394,10 @@ Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr<caffe::Net<fl
const int w = wn * output_size;
const int h = hn * output_size;
if (w + block_size <= Width && h + block_size <= Height)
if (w + crop_size <= Width && h + crop_size <= Height)
{
{
cv::Mat someimg = im(cv::Rect(w, h, block_size, block_size));
cv::Mat someimg = im(cv::Rect(w, h, crop_size, crop_size));
cv::Mat someborderimg;
// 画像を中央にパディング。余白はcv::BORDER_REPLICATEで埋める
cv::copyMakeBorder(someimg, someborderimg, layer_num, layer_num, layer_num, layer_num, cv::BORDER_REPLICATE);
@ -420,12 +410,12 @@ Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr<caffe::Net<fl
const auto Line = someborderimg.step1();
if (block_width_height == Line)
memcpy(fptr, uptr, block_width_height * block_width_height * sizeof(float));
if (block_size == Line)
memcpy(fptr, uptr, block_size * block_size * sizeof(float));
else
{
for (int i = 0; i < block_width_height; i++)
memcpy(fptr + i * block_width_height, uptr + i * Line, block_width_height * sizeof(float));
for (int i = 0; i < block_size; i++)
memcpy(fptr + i * block_size, uptr + i * Line, block_size * sizeof(float));
}
}
}
@ -462,8 +452,8 @@ Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr<caffe::Net<fl
const float *fptr = block.data() + (output_block_plane_size * n);
// 結果を入力画像にコピー(後に処理する部分とここで上書きする部分は被らないから、入力画像を上書きしても大丈夫)
for (int i = 0; i < block_size; i++)
caffe::caffe_copy(block_size, fptr + i * block_size, imptr + (h + i) * Line + w);
for (int i = 0; i < crop_size; i++)
caffe::caffe_copy(crop_size, fptr + i * crop_size, imptr + (h + i) * Line + w);
}
}
}
@ -475,7 +465,8 @@ Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr<caffe::Net<fl
return eWaifu2xError_OK;
}
Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &Mode, const int NoiseLevel, const double ScaleRatio, const std::string &ModelDir, const std::string &Process)
Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &Mode, const int NoiseLevel, const double ScaleRatio, const std::string &ModelDir, const std::string &Process,
const int CropSize, const int BatchSize)
{
Waifu2x::eWaifu2xError ret;
@ -491,6 +482,13 @@ Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &M
model_dir = ModelDir;
process = Process;
crop_size = CropSize;
batch_size = BatchSize;
output_size = crop_size - offset * 2;
block_size = crop_size + layer_num * 2;
original_width_height = 128 + layer_num * 2;
std::call_once(waifu2x_once_flag, [argc, argv]()
{
assert(argc >= 1);

17
common/waifu2x.h
View File

@ -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();

8
waifu2x-caffe-gui/Source.cpp
View File

@ -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)
{
}

11
waifu2x-caffe/Source.cpp
View File

@ -99,6 +99,14 @@ int main(int argc, char** argv)
TCLAP::ValueArg<std::string> cmdProcess("p", "process", "process mode",
false, "gpu", &cmdProcessConstraint, cmd);
TCLAP::ValueArg<int> cmdCropSizeFile("c", "crop_size",
"input image split size", false,
128, "int", cmd);
TCLAP::ValueArg<int> 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: