From ba87bacf8860bcc2fac620984509b24c32074ecd Mon Sep 17 00:00:00 2001 From: lltcggie Date: Sun, 6 Dec 2015 18:48:37 +0900 Subject: [PATCH] =?UTF-8?q?=E3=83=91=E3=82=B9=E3=81=8C=E3=83=AF=E3=82=A4?= =?UTF-8?q?=E3=83=89=E6=96=87=E5=AD=97=E3=81=AE=E5=A0=B4=E5=90=88=E3=81=AB?= =?UTF-8?q?=E5=AF=BE=E5=BF=9C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- common/waifu2x.cpp | 241 ++++++++++++++++---- common/waifu2x.h | 21 +- waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj | 4 +- waifu2x-caffe-gui/waifu2x-caffe-gui.vcxproj | 4 +- waifu2x-caffe/waifu2x-caffe.vcxproj | 4 +- 5 files changed, 210 insertions(+), 64 deletions(-) diff --git a/common/waifu2x.cpp b/common/waifu2x.cpp index cbc0769..d3d25b5 100644 --- a/common/waifu2x.cpp +++ b/common/waifu2x.cpp @@ -10,6 +10,16 @@ #include #include +#include +#include +#include +#include +#include +#include +#ifdef _MSC_VER +#include +#endif + #define STB_IMAGE_IMPLEMENTATION #include #define STB_IMAGE_WRITE_IMPLEMENTATION @@ -19,6 +29,7 @@ #include #endif +#ifdef _MSC_VER #ifdef _DEBUG #pragma comment(lib, "caffe-d.lib") #pragma comment(lib, "proto-d.lib") @@ -42,6 +53,7 @@ #pragma comment(lib, "opencv_core249d.lib") #pragma comment(lib, "opencv_highgui249d.lib") #pragma comment(lib, "opencv_imgproc249d.lib") +#pragma comment(lib, "libboost_iostreams-vc120-mt-gd-1_59.lib") #else #pragma comment(lib, "caffe.lib") #pragma comment(lib, "proto.lib") @@ -65,6 +77,8 @@ #pragma comment(lib, "opencv_core249.lib") #pragma comment(lib, "opencv_highgui249.lib") #pragma comment(lib, "opencv_imgproc249.lib") +#pragma comment(lib, "libboost_iostreams-vc120-mt-1_59.lib") +#endif #endif // 入力画像のオフセット @@ -82,6 +96,8 @@ const int MinCudaDriverVersion = 6050; // https://github.com/nagadomi/waifu2x/commit/797b45ae23665a1c5e3c481c018e48e6f0d0e383 const double clip_eps8 = (1.0 / 255.0) * 0.5 - (1.0e-7 * (1.0 / 255.0) * 0.5); +const int kProtoReadBytesLimit = INT_MAX; // Max size of 2 GB minus 1 byte. + static std::once_flag waifu2x_once_flag; static std::once_flag waifu2x_cudnn_once_flag; static std::once_flag waifu2x_cuda_once_flag; @@ -132,6 +148,110 @@ namespace IgnoreErrorCV g_IgnoreErrorCV; } +template +static bool writeFile(boost::iostreams::stream &os, const std::vector &buf) +{ + if (!os) + return false; + + const auto WriteSize = sizeof(BufType) * buf.size(); + os.write((const char *)buf.data(), WriteSize); + if (os.fail()) + return false; + + return true; +} + +template +static bool writeFile(const boost::filesystem::path &path, std::vector &buf) +{ + boost::iostreams::stream os(path, std::ios_base::out | std::ios_base::binary | std::ios_base::trunc); + + return writeFile(os, buf); +} + +template +static bool readFile(boost::iostreams::stream &is, std::vector &buf) +{ + if (!is) + return false; + + const auto size = is.seekg(0, std::ios::end).tellg(); + is.seekg(0, std::ios::beg); + + buf.resize((size / sizeof(BufType)) + (size % sizeof(BufType))); + is.read(buf.data(), size); + if (is.gcount() != size) + return false; + + return true; +} + +template +static bool readFile(const boost::filesystem::path &path, std::vector &buf) +{ + boost::iostreams::stream is(path, std::ios_base::in | std::ios_base::binary); + + return readFile(is, buf); +} + +static Waifu2x::eWaifu2xError readProtoText(const boost::filesystem::path &path, ::google::protobuf::Message* proto) +{ + boost::iostreams::stream is(path, std::ios_base::in); + + if (!is) + return Waifu2x::eWaifu2xError_FailedOpenModelFile; + + std::vector tmp; + if (!readFile(is, tmp)) + return Waifu2x::eWaifu2xError_FailedParseModelFile; + + google::protobuf::io::ArrayInputStream input(tmp.data(), tmp.size()); + const bool success = google::protobuf::TextFormat::Parse(&input, proto); + + if (!success) + return Waifu2x::eWaifu2xError_FailedParseModelFile; + + return Waifu2x::eWaifu2xError_OK; +} + +static Waifu2x::eWaifu2xError readProtoBinary(const boost::filesystem::path &path, ::google::protobuf::Message* proto) +{ + boost::iostreams::stream is(path, std::ios_base::in | std::ios_base::binary); + + if (!is) + return Waifu2x::eWaifu2xError_FailedParseModelFile; + + std::vector tmp; + if (!readFile(is, tmp)) + return Waifu2x::eWaifu2xError_FailedParseModelFile; + + google::protobuf::io::ArrayInputStream input(tmp.data(), tmp.size()); + + google::protobuf::io::CodedInputStream coded_input(&input); + coded_input.SetTotalBytesLimit(kProtoReadBytesLimit, 536870912); + + const bool success = proto->ParseFromCodedStream(&coded_input); + if (!success) + return Waifu2x::eWaifu2xError_FailedParseModelFile; + + return Waifu2x::eWaifu2xError_OK; +} + +static Waifu2x::eWaifu2xError writeProtoBinary(const ::google::protobuf::Message& proto, const boost::filesystem::path &path) +{ + boost::iostreams::stream os(path, std::ios_base::out | std::ios_base::binary | std::ios_base::trunc); + + if (!os) + return Waifu2x::eWaifu2xError_FailedWriteModelFile; + + if (!proto.SerializePartialToOstream(&os)) + return Waifu2x::eWaifu2xError_FailedWriteModelFile; + + return Waifu2x::eWaifu2xError_OK; +} + + Waifu2x::Waifu2x() : is_inited(false), isCuda(false), input_block(nullptr), dummy_data(nullptr), output_block(nullptr) { } @@ -227,7 +347,7 @@ void Waifu2x::quit_liblary() { } -cv::Mat Waifu2x::LoadMat(const std::string &path) +cv::Mat Waifu2x::LoadMat(const boost::filesystem::path &path) { cv::Mat mat; LoadMat(mat, path); @@ -290,14 +410,24 @@ Waifu2x::eWaifu2xError Waifu2x::AlphaMakeBorder(std::vector &planes, co } // 画像を読み込んで値を0.0f〜1.0fの範囲に変換 -Waifu2x::eWaifu2xError Waifu2x::LoadMat(cv::Mat &float_image, const std::string &input_file) +Waifu2x::eWaifu2xError Waifu2x::LoadMat(cv::Mat &float_image, const boost::filesystem::path &input_file) { - cv::Mat original_image = cv::imread(input_file, cv::IMREAD_UNCHANGED); - if (original_image.empty()) + cv::Mat original_image; + { - const eWaifu2xError ret = LoadMatBySTBI(original_image, input_file); - if (ret != eWaifu2xError_OK) - return ret; + std::vector img_data; + if (!readFile(input_file, img_data)) + return Waifu2x::eWaifu2xError_FailedOpenInputFile; + + cv::Mat im(img_data.size(), 1, CV_8U, img_data.data()); + original_image = cv::imdecode(im, cv::IMREAD_UNCHANGED); + + if (original_image.empty()) + { + const eWaifu2xError ret = LoadMatBySTBI(original_image, img_data); + if (ret != eWaifu2xError_OK) + return ret; + } } cv::Mat convert; @@ -326,10 +456,10 @@ Waifu2x::eWaifu2xError Waifu2x::LoadMat(cv::Mat &float_image, const std::string return eWaifu2xError_OK; } -Waifu2x::eWaifu2xError Waifu2x::LoadMatBySTBI(cv::Mat &float_image, const std::string &input_file) +Waifu2x::eWaifu2xError Waifu2x::LoadMatBySTBI(cv::Mat &float_image, const std::vector &img_data) { int x, y, comp; - stbi_uc *data = stbi_load(input_file.c_str(), &x, &y, &comp, 4); + stbi_uc *data = stbi_load_from_memory((const stbi_uc *)img_data.data(), img_data.size(), &x, &y, &comp, 4); if (!data) return eWaifu2xError_FailedOpenInputFile; @@ -454,34 +584,38 @@ Waifu2x::eWaifu2xError Waifu2x::CreateZoomColorImage(const cv::Mat &float_image, // モデルファイルからネットワークを構築 // processでcudnnが指定されなかった場合はcuDNNが呼び出されないように変更する -Waifu2x::eWaifu2xError Waifu2x::ConstractNet(boost::shared_ptr> &net, const std::string &model_path, const std::string ¶m_path, const std::string &process) +Waifu2x::eWaifu2xError Waifu2x::ConstractNet(boost::shared_ptr> &net, const boost::filesystem::path &model_path, const boost::filesystem::path ¶m_path, const std::string &process) { - const std::string caffemodel_path = param_path + ".caffemodel"; - const std::string modelbin_path = model_path + ".protobin"; - - FILE *fp = fopen(caffemodel_path.c_str(), "rb"); - const bool isModelExist = fp != nullptr; - if (fp) fclose(fp); - - fp = fopen(modelbin_path.c_str(), "rb"); - const bool isModelBinExist = fp != nullptr; - if (fp) fclose(fp); + boost::filesystem::path caffemodel_path = param_path; + caffemodel_path += ".caffemodel"; + boost::filesystem::path modelbin_path = model_path; + modelbin_path += ".protobin"; caffe::NetParameter param; - if (isModelExist && isModelBinExist && caffe::ReadProtoFromBinaryFile(modelbin_path, ¶m)) + if (readProtoBinary(modelbin_path, ¶m) == eWaifu2xError_OK) { - const auto ret = SetParameter(param, process); + Waifu2x::eWaifu2xError ret; + + ret = SetParameter(param, process); if (ret != eWaifu2xError_OK) return ret; + caffe::NetParameter param_caffemodel; + ret = readProtoBinary(caffemodel_path, ¶m_caffemodel); + if (ret != eWaifu2xError_OK) + return ret; + + if (!caffe::UpgradeNetAsNeeded(caffemodel_path.string(), ¶m_caffemodel)) + return Waifu2x::eWaifu2xError_FailedParseModelFile; + net = boost::shared_ptr>(new caffe::Net(param)); - net->CopyTrainedLayersFrom(caffemodel_path); + net->CopyTrainedLayersFrom(param_caffemodel); input_plane = param.input_dim(1); } else { - const auto ret = LoadParameterFromJson(net, model_path, param_path, process); + const auto ret = LoadParameterFromJson(net, model_path, param_path, modelbin_path, caffemodel_path, process); if (ret != eWaifu2xError_OK) return ret; } @@ -527,18 +661,21 @@ Waifu2x::eWaifu2xError Waifu2x::SetParameter(caffe::NetParameter ¶m, const s return eWaifu2xError_OK; } -Waifu2x::eWaifu2xError Waifu2x::LoadParameterFromJson(boost::shared_ptr> &net, const std::string &model_path, const std::string ¶m_path, const std::string &process) +Waifu2x::eWaifu2xError Waifu2x::LoadParameterFromJson(boost::shared_ptr> &net, const boost::filesystem::path &model_path, const boost::filesystem::path ¶m_path + , const boost::filesystem::path &modelbin_path, const boost::filesystem::path &caffemodel_path, const std::string &process) { - const std::string caffemodel_path = param_path + ".caffemodel"; - const std::string modelbin_path = model_path + ".protobin"; + Waifu2x::eWaifu2xError ret; caffe::NetParameter param; - if (!caffe::ReadProtoFromTextFile(model_path, ¶m)) - return eWaifu2xError_FailedOpenModelFile; + ret = readProtoText(model_path, ¶m); + if (ret != eWaifu2xError_OK) + return ret; - caffe::WriteProtoToBinaryFile(param, modelbin_path); + ret = writeProtoBinary(param, modelbin_path); + if (ret != eWaifu2xError_OK) + return ret; - const auto ret = SetParameter(param, process); + ret = SetParameter(param, process); if (ret != eWaifu2xError_OK) return ret; @@ -549,18 +686,15 @@ Waifu2x::eWaifu2xError Waifu2x::LoadParameterFromJson(boost::shared_ptr is(param_path, std::ios_base::in | std::ios_base::binary); + if(!is) return eWaifu2xError_FailedOpenModelFile; - fseek(fp, 0, SEEK_END); - const auto size = ftell(fp); - fseek(fp, 0, SEEK_SET); + const size_t size = is.seekg(0, std::ios::end).tellg(); + is.seekg(0, std::ios::beg); jsonBuf.resize(size + 1); - fread(jsonBuf.data(), 1, size, fp); - - fclose(fp); + is.read(jsonBuf.data(), jsonBuf.size()); jsonBuf[jsonBuf.size() - 1] = '\0'; @@ -677,7 +811,9 @@ Waifu2x::eWaifu2xError Waifu2x::LoadParameterFromJson(boost::shared_ptrToProto(¶m); - caffe::WriteProtoToBinaryFile(param, caffemodel_path); + ret = writeProtoBinary(param, caffemodel_path); + if (ret != eWaifu2xError_OK) + return ret; } catch (...) { @@ -895,7 +1031,7 @@ Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr os(output_file, std::ios_base::out | std::ios_base::binary | std::ios_base::trunc); + if(!os) + return eWaifu2xError_FailedOpenOutputFile; + + if (!stbi_write_tga(os, im.size().width, im.size().height, im.channels(), data)) return eWaifu2xError_FailedOpenOutputFile; return eWaifu2xError_OK; @@ -1099,7 +1239,10 @@ Waifu2x::eWaifu2xError Waifu2x::WriteMat(const cv::Mat &im, const std::string &o try { - if (cv::imwrite(output_file, im)) + std::vector buf; + cv::imencode(ext, im, buf); + + if (writeFile(output_file, buf)) return eWaifu2xError_OK; } @@ -1344,7 +1487,7 @@ Waifu2x::eWaifu2xError Waifu2x::AfterReconstructFloatMatProcess(const bool isRec return eWaifu2xError_OK; } -Waifu2x::eWaifu2xError Waifu2x::waifu2x(const std::string &input_file, const std::string &output_file, +Waifu2x::eWaifu2xError Waifu2x::waifu2x(const boost::filesystem::path &input_file, const boost::filesystem::path &output_file, const waifu2xCancelFunc cancel_func) { Waifu2x::eWaifu2xError ret; diff --git a/common/waifu2x.h b/common/waifu2x.h index cf83fb1..af6462c 100644 --- a/common/waifu2x.h +++ b/common/waifu2x.h @@ -6,6 +6,7 @@ #include #include #include +#include #include #define CUDNN_DLL_NAME "cudnn64_70" @@ -31,6 +32,7 @@ public: eWaifu2xError_FailedOpenOutputFile, eWaifu2xError_FailedOpenModelFile, eWaifu2xError_FailedParseModelFile, + eWaifu2xError_FailedWriteModelFile, eWaifu2xError_FailedConstructModel, eWaifu2xError_FailedProcessCaffe, eWaifu2xError_FailedCudaCheck, @@ -73,7 +75,7 @@ private: std::string mode; int noise_level; double scale_ratio; - std::string model_dir; + boost::filesystem::path model_dir; std::string process; int inner_padding; @@ -95,18 +97,19 @@ private: bool use_tta; private: - static eWaifu2xError LoadMat(cv::Mat &float_image, const std::string &input_file); - static eWaifu2xError LoadMatBySTBI(cv::Mat &float_image, const std::string &input_file); + static eWaifu2xError LoadMat(cv::Mat &float_image, const boost::filesystem::path &input_file); + static eWaifu2xError LoadMatBySTBI(cv::Mat &float_image, const std::vector &img_data); static eWaifu2xError AlphaMakeBorder(std::vector &planes, const cv::Mat &alpha, const int offset); eWaifu2xError CreateBrightnessImage(const cv::Mat &float_image, cv::Mat &im); eWaifu2xError PaddingImage(const cv::Mat &input, cv::Mat &output); eWaifu2xError Zoom2xAndPaddingImage(const cv::Mat &input, cv::Mat &output, cv::Size_ &zoom_size); eWaifu2xError CreateZoomColorImage(const cv::Mat &float_image, const cv::Size_ &zoom_size, std::vector &cubic_planes); - eWaifu2xError ConstractNet(boost::shared_ptr> &net, const std::string &model_path, const std::string ¶m_path, const std::string &process); - eWaifu2xError LoadParameterFromJson(boost::shared_ptr> &net, const std::string &model_path, const std::string ¶m_path, const std::string &process); + eWaifu2xError ConstractNet(boost::shared_ptr> &net, const boost::filesystem::path &model_path, const boost::filesystem::path ¶m_path, const std::string &process); + eWaifu2xError LoadParameterFromJson(boost::shared_ptr> &net, const boost::filesystem::path &model_path, const boost::filesystem::path ¶m_path + , const boost::filesystem::path &modelbin_path, const boost::filesystem::path &caffemodel_path, const std::string &process); eWaifu2xError SetParameter(caffe::NetParameter ¶m, const std::string &process) const; eWaifu2xError ReconstructImage(boost::shared_ptr> net, cv::Mat &im); - eWaifu2xError WriteMat(const cv::Mat &im, const std::string &output_file); + eWaifu2xError WriteMat(const cv::Mat &im, const boost::filesystem::path &output_file); eWaifu2xError BeforeReconstructFloatMatProcess(const cv::Mat &in, cv::Mat &out, bool &convertBGRflag); eWaifu2xError ReconstructFloatMat(const bool isReconstructNoise, const bool isReconstructScale, const waifu2xCancelFunc cancel_func, const cv::Mat &in, cv::Mat &out); @@ -125,15 +128,15 @@ 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 boost::filesystem::path &model_dir, const std::string &process, const bool use_tta = false, const int crop_size = 128, const int batch_size = 1); void destroy(); - eWaifu2xError waifu2x(const std::string &input_file, const std::string &output_file, + eWaifu2xError waifu2x(const boost::filesystem::path &input_file, const boost::filesystem::path &output_file, const waifu2xCancelFunc cancel_func = nullptr); const std::string& used_process() const; - static cv::Mat LoadMat(const std::string &path); + static cv::Mat LoadMat(const boost::filesystem::path &path); }; diff --git a/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj b/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj index 2c5d72c..0bc580c 100644 --- a/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj +++ b/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj @@ -42,14 +42,14 @@ true D:\caffe-build\install\include;D:\caffe-build\install\include\boost-1_59;$(CUDA_PATH_V7_5)\include;C:\opencv249\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)stb;$(SolutionDir)include;$(IncludePath) - D:\caffe-build\install\lib;$(CUDA_PATH_V7_5)\lib\$(PlatformName);C:\opencv249\build\x64\vc12\lib;$(LibraryPath) + D:\caffe-build\install\lib;$(CUDA_PATH_V7_5)\lib\$(PlatformName);C:\opencv249\build\x64\vc12\lib;D:\program\boost_1_59_0\stage\x64\lib;$(LibraryPath) waifu2x-caffed $(SolutionDir)bin false D:\caffe-build\install\include;D:\caffe-build\install\include\boost-1_59;$(CUDA_PATH_V7_5)\include;C:\opencv249\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)stb;$(SolutionDir)include;$(IncludePath) - D:\caffe-build\install\lib;$(CUDA_PATH_V7_5)\lib\$(PlatformName);C:\opencv249\build\x64\vc12\lib;$(LibraryPath) + D:\caffe-build\install\lib;$(CUDA_PATH_V7_5)\lib\$(PlatformName);C:\opencv249\build\x64\vc12\lib;D:\program\boost_1_59_0\stage\x64\lib;$(LibraryPath) waifu2x-caffe $(SolutionDir)bin diff --git a/waifu2x-caffe-gui/waifu2x-caffe-gui.vcxproj b/waifu2x-caffe-gui/waifu2x-caffe-gui.vcxproj index 275ac31..d8b944e 100644 --- a/waifu2x-caffe-gui/waifu2x-caffe-gui.vcxproj +++ b/waifu2x-caffe-gui/waifu2x-caffe-gui.vcxproj @@ -42,14 +42,14 @@ true D:\caffe-build\install\include;D:\caffe-build\install\include\boost-1_59;$(CUDA_PATH_V7_5)\include;C:\opencv249\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)stb;$(SolutionDir)include;$(IncludePath) - D:\caffe-build\install\lib;$(CUDA_PATH_V7_5)\lib\$(PlatformName);C:\opencv249\build\x64\vc12\lib;$(LibraryPath) + D:\caffe-build\install\lib;$(CUDA_PATH_V7_5)\lib\$(PlatformName);C:\opencv249\build\x64\vc12\lib;D:\program\boost_1_59_0\stage\x64\lib;$(LibraryPath) $(SolutionDir)bin $(ProjectName)d false D:\caffe-build\install\include;D:\caffe-build\install\include\boost-1_59;$(CUDA_PATH_V7_5)\include;C:\opencv249\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)stb;$(SolutionDir)include;$(IncludePath) - D:\caffe-build\install\lib;$(CUDA_PATH_V7_5)\lib\$(PlatformName);C:\opencv249\build\x64\vc12\lib;$(LibraryPath) + D:\caffe-build\install\lib;$(CUDA_PATH_V7_5)\lib\$(PlatformName);C:\opencv249\build\x64\vc12\lib;D:\program\boost_1_59_0\stage\x64\lib;$(LibraryPath) $(SolutionDir)bin diff --git a/waifu2x-caffe/waifu2x-caffe.vcxproj b/waifu2x-caffe/waifu2x-caffe.vcxproj index 6bf4bce..6d9f29e 100644 --- a/waifu2x-caffe/waifu2x-caffe.vcxproj +++ b/waifu2x-caffe/waifu2x-caffe.vcxproj @@ -42,14 +42,14 @@ true D:\caffe-build\install\include;D:\caffe-build\install\include\boost-1_59;$(CUDA_PATH_V7_5)\include;C:\opencv249\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)stb;$(SolutionDir)include;$(IncludePath) - D:\caffe-build\install\lib;$(CUDA_PATH_V7_5)\lib\$(PlatformName);C:\opencv249\build\x64\vc12\lib;$(LibraryPath) + D:\caffe-build\install\lib;$(CUDA_PATH_V7_5)\lib\$(PlatformName);C:\opencv249\build\x64\vc12\lib;D:\program\boost_1_59_0\stage\x64\lib;$(LibraryPath) $(SolutionDir)bin $(ProjectName)d false D:\caffe-build\install\include;D:\caffe-build\install\include\boost-1_59;$(CUDA_PATH_V7_5)\include;C:\opencv249\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)stb;$(SolutionDir)include;$(IncludePath) - D:\caffe-build\install\lib;$(CUDA_PATH_V7_5)\lib\$(PlatformName);C:\opencv249\build\x64\vc12\lib;$(LibraryPath) + D:\caffe-build\install\lib;$(CUDA_PATH_V7_5)\lib\$(PlatformName);C:\opencv249\build\x64\vc12\lib;D:\program\boost_1_59_0\stage\x64\lib;$(LibraryPath) $(SolutionDir)bin