diff --git a/common/waifu2x.cpp b/common/waifu2x.cpp index abd8312..baad8c7 100644 --- a/common/waifu2x.cpp +++ b/common/waifu2x.cpp @@ -1673,38 +1673,23 @@ Waifu2x::eWaifu2xError Waifu2x::AfterReconstructFloatMatProcess(const bool isRec return eWaifu2xError_OK; } -Waifu2x::eWaifu2xError Waifu2x::waifu2x(const boost::filesystem::path &input_file, const boost::filesystem::path &output_file, - const waifu2xCancelFunc cancel_func) +Waifu2x::eWaifu2xError Waifu2x::waifu2xConvetedMat(const bool isJpeg, const cv::Mat &inMat, cv::Mat &outMat, const waifu2xCancelFunc cancel_func) { Waifu2x::eWaifu2xError ret; - if (!is_inited) - return eWaifu2xError_NotInitialized; - - const boost::filesystem::path ip(input_file); - const boost::filesystem::path ipext(ip.extension()); - - const bool isJpeg = boost::iequals(ipext.string(), ".jpg") || boost::iequals(ipext.string(), ".jpeg"); const bool isReconstructNoise = mode == "noise" || mode == "noise_scale" || (mode == "auto_scale" && isJpeg); const bool isReconstructScale = mode == "scale" || mode == "noise_scale" || mode == "auto_scale"; - cv::Mat float_image; - ret = LoadMat(float_image, input_file); - if (ret != eWaifu2xError_OK) - return ret; - cv::Mat reconstruct_image; - ret = Reconstruct(isReconstructNoise, isReconstructScale, cancel_func, float_image, reconstruct_image); + ret = Reconstruct(isReconstructNoise, isReconstructScale, cancel_func, inMat, reconstruct_image); if (ret != eWaifu2xError_OK) return ret; cv::Mat process_image; - ret = AfterReconstructFloatMatProcess(isReconstructScale, cancel_func, float_image, reconstruct_image, process_image); + ret = AfterReconstructFloatMatProcess(isReconstructScale, cancel_func, inMat, reconstruct_image, process_image); if (ret != eWaifu2xError_OK) return ret; - float_image.release(); - const int cv_depth = DepthBitToCVDepth(output_depth); const double max_val = GetValumeMaxFromCVDepth(cv_depth); const double eps = GetEPS(cv_depth); @@ -1736,11 +1721,115 @@ Waifu2x::eWaifu2xError Waifu2x::waifu2x(const boost::filesystem::path &input_fil cv::merge(planes, write_iamge); } + outMat = write_iamge; + + return eWaifu2xError_OK; +} + +Waifu2x::eWaifu2xError Waifu2x::waifu2x(const boost::filesystem::path &input_file, const boost::filesystem::path &output_file, + const waifu2xCancelFunc cancel_func) +{ + Waifu2x::eWaifu2xError ret; + + if (!is_inited) + return eWaifu2xError_NotInitialized; + + const boost::filesystem::path ip(input_file); + const boost::filesystem::path ipext(ip.extension()); + + const bool isJpeg = boost::iequals(ipext.string(), ".jpg") || boost::iequals(ipext.string(), ".jpeg"); + + cv::Mat float_image; + ret = LoadMat(float_image, input_file); + if (ret != eWaifu2xError_OK) + return ret; + + cv::Mat write_iamge; + ret = waifu2xConvetedMat(isJpeg, float_image, write_iamge, cancel_func); + if (ret != eWaifu2xError_OK) + return ret; + ret = WriteMat(write_iamge, output_file, output_quality); if (ret != eWaifu2xError_OK) return ret; - write_iamge.release(); + return eWaifu2xError_OK; +} + +Waifu2x::eWaifu2xError Waifu2x::waifu2x(double factor, const void* source, void* dest, int width, int height, int in_channel, int in_stride, int out_channel, int out_stride) +{ + Waifu2x::eWaifu2xError ret; + + if (!is_inited) + return eWaifu2xError_NotInitialized; + + if (output_depth != 8) // 出力深度は8bitだけ + return eWaifu2xError_InvalidParameter; + + cv::Mat float_image; + + // Matへ変換 + { + cv::Mat original_image(cv::Size(width, height), CV_MAKETYPE(CV_8U, in_channel), (void *)source, in_stride); + + cv::Mat convert; + switch (original_image.depth()) + { + case CV_8U: + original_image.convertTo(convert, CV_32F, 1.0 / GetValumeMaxFromCVDepth(CV_8U)); + break; + + case CV_16U: + original_image.convertTo(convert, CV_32F, 1.0 / GetValumeMaxFromCVDepth(CV_16U)); + break; + + case CV_32F: + convert = original_image; // 元から0.0〜1.0のはずなので変換は必要ない + break; + } + + original_image.release(); + + if (convert.channels() == 1) + cv::cvtColor(convert, convert, cv::COLOR_GRAY2BGR); + else if (convert.channels() == 4) + { + // アルファチャンネル付きだったら透明なピクセルのと不透明なピクセルの境界部分の色を広げる + + std::vector planes; + cv::split(convert, planes); + + cv::Mat alpha = planes[3]; + planes.resize(3); + AlphaMakeBorder(planes, alpha, layer_num); + + planes.push_back(alpha); + cv::merge(planes, convert); + } + + float_image = convert; + } + + const auto oldScale = scale_ratio; + scale_ratio = factor; + + cv::Mat write_iamge; + ret = waifu2xConvetedMat(false, float_image, write_iamge); + + scale_ratio = oldScale; + + if (ret != eWaifu2xError_OK) + return ret; + + float_image.release(); + + // 出力配列へ書き込み + { + const auto width = write_iamge.size().width; + const auto stride = write_iamge.step1(); + for (int i = 0; i < write_iamge.size().height; i++) + memcpy((uint8_t *)dest + out_stride * i, write_iamge.data + stride * i, stride); + } return eWaifu2xError_OK; } diff --git a/common/waifu2x.h b/common/waifu2x.h index 5fc292c..80ffdcb 100644 --- a/common/waifu2x.h +++ b/common/waifu2x.h @@ -133,6 +133,9 @@ private: eWaifu2xError Reconstruct(const bool isReconstructNoise, const bool isReconstructScale, const waifu2xCancelFunc cancel_func, const cv::Mat &in, cv::Mat &out); eWaifu2xError AfterReconstructFloatMatProcess(const bool isReconstructScale, const waifu2xCancelFunc cancel_func, const cv::Mat &floatim, const cv::Mat &in, cv::Mat &out); + eWaifu2xError waifu2xConvetedMat(const bool isJpeg, const cv::Mat &inMat, cv::Mat &outMat, + const waifu2xCancelFunc cancel_func = nullptr); + static int DepthBitToCVDepth(const int depth_bit); static double GetValumeMaxFromCVDepth(const int cv_depth); static double GetEPS(const int cv_depth); @@ -157,6 +160,13 @@ public: eWaifu2xError waifu2x(const boost::filesystem::path &input_file, const boost::filesystem::path &output_file, const waifu2xCancelFunc cancel_func = nullptr); + // factor: 倍率 + // source: (4チャンネルの場合は)BGRAな画素配列 + // dest: (4チャンネルの場合は)処理したBGRAな画素配列 + // in_stride: sourceのストライド(バイト単位) + // out_stride: destのストライド(バイト単位) + eWaifu2xError waifu2x(double factor, const void* source, void* dest, int width, int height, int in_channel, int in_stride, int out_channel, int out_stride); + const std::string& used_process() const; static cv::Mat LoadMat(const boost::filesystem::path &path); diff --git a/waifu2x-caffe-dll/Source.cpp b/waifu2x-caffe-dll/Source.cpp index 963606c..2712d59 100644 --- a/waifu2x-caffe-dll/Source.cpp +++ b/waifu2x-caffe-dll/Source.cpp @@ -1,16 +1,16 @@ #include #include -#include "waifu2x.h" +#include "../common/waifu2x.h" __declspec(dllexport) -void* Waifu2xInit(const char *mode, const int noise_level, const char *model_dir, const char *process, const int crop_size = 128, const int batch_size = 1) +void* Waifu2xInit(const char *mode, const int noise_level, const char *model_dir, const char *process, const int output_depth = 8, const bool use_tta = false, const int crop_size = 128, const int batch_size = 1) { Waifu2x *obj = new Waifu2x(); char *argv[] = { "" }; - if (obj->init(1, argv, mode, noise_level, model_dir, process, crop_size, batch_size) != Waifu2x::eWaifu2xError_OK) + if (obj->init(1, argv, mode, noise_level, 2.0, model_dir, process, boost::optional(), output_depth, use_tta, crop_size, batch_size) != Waifu2x::eWaifu2xError_OK) { delete obj; return nullptr; @@ -20,14 +20,14 @@ void* Waifu2xInit(const char *mode, const int noise_level, const char *model_dir } __declspec(dllexport) -bool Waifu2xProcess(void *waifu2xObj, int factor, const uint32_t* source, uint32_t* dest, int width, int height) +bool Waifu2xProcess(void *waifu2xObj, double factor, const void* source, void* dest, int width, int height, int in_channel, int in_stride, int out_channel, int out_stride) { if (!waifu2xObj) return false; Waifu2x *obj = (Waifu2x *)waifu2xObj; - return obj->waifu2x(factor, source, dest, width, height) == Waifu2x::eWaifu2xError_OK; + return obj->waifu2x(factor, source, dest, width, height, in_channel, in_stride, out_channel, out_stride) == Waifu2x::eWaifu2xError_OK; } __declspec(dllexport) diff --git a/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj b/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj index 0bc580c..38779da 100644 --- a/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj +++ b/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj @@ -41,15 +41,15 @@ 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;D:\program\boost_1_59_0\stage\x64\lib;$(LibraryPath) + D:\caffe-build\install\include;D:\caffe-build\install\include\boost-1_59;$(CUDA_PATH_V7_5)\include;C:\opencv310\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)stb;$(SolutionDir)include;$(IncludePath) + D:\caffe-build\install\lib;$(CUDA_PATH_V7_5)\lib\$(PlatformName);C:\opencv310\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;D:\program\boost_1_59_0\stage\x64\lib;$(LibraryPath) + D:\caffe-build\install\include;D:\caffe-build\install\include\boost-1_59;$(CUDA_PATH_V7_5)\include;C:\opencv310\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)stb;$(SolutionDir)include;$(IncludePath) + D:\caffe-build\install\lib;$(CUDA_PATH_V7_5)\lib\$(PlatformName);C:\opencv310\build\x64\vc12\lib;D:\program\boost_1_59_0\stage\x64\lib;$(LibraryPath) waifu2x-caffe $(SolutionDir)bin @@ -65,7 +65,7 @@ Console true - cudnn64_70.dll;%(DelayLoadDLLs) + cudnn64_4.dll;%(DelayLoadDLLs) Shlwapi.lib;%(AdditionalDependencies) @@ -85,16 +85,16 @@ true true true - cudnn64_70.dll;%(DelayLoadDLLs) + cudnn64_4.dll;%(DelayLoadDLLs) Shlwapi.lib;%(AdditionalDependencies) + - - + diff --git a/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj.filters b/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj.filters index c230057..86e4980 100644 --- a/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj.filters +++ b/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj.filters @@ -18,12 +18,12 @@ 繧ｽ繝ｼ繧ｹ 繝輔ぃ繧､繝ｫ - + 繧ｽ繝ｼ繧ｹ 繝輔ぃ繧､繝ｫ - + 繝倥ャ繝繝ｼ 繝輔ぃ繧､繝ｫ diff --git a/waifu2x-caffe-dll/waifu2x.cpp b/waifu2x-caffe-dll/waifu2x.cpp deleted file mode 100644 index 3dcb539..0000000 --- a/waifu2x-caffe-dll/waifu2x.cpp +++ /dev/null @@ -1,831 +0,0 @@ -#include "waifu2x.h" -#include -#include -#include -#include -#include -#include -#include -#include - - -#ifdef _DEBUG -#pragma comment(lib, "caffe-d.lib") -#pragma comment(lib, "proto-d.lib") -#pragma comment(lib, "libboost_system-vc120-mt-gd-1_59.lib") -#pragma comment(lib, "libboost_thread-vc120-mt-gd-1_59.lib") -#pragma comment(lib, "libboost_filesystem-vc120-mt-gd-1_59.lib") -#pragma comment(lib, "glogd.lib") -#pragma comment(lib, "gflagsd.lib") -#pragma comment(lib, "libprotobufd.lib") -#pragma comment(lib, "libhdf5_hl_D.lib") -#pragma comment(lib, "libhdf5_D.lib") -#pragma comment(lib, "zlibstaticd.lib") -#pragma comment(lib, "libopenblas.dll.a") -#pragma comment(lib, "cudart.lib") -#pragma comment(lib, "curand.lib") -#pragma comment(lib, "cublas.lib") -#pragma comment(lib, "cudnn.lib") - -#pragma comment(lib, "opencv_calib3d249d.lib") -#pragma comment(lib, "opencv_contrib249d.lib") -#pragma comment(lib, "opencv_core249d.lib") -#pragma comment(lib, "opencv_highgui249d.lib") -#pragma comment(lib, "opencv_imgproc249d.lib") -#else -#pragma comment(lib, "caffe.lib") -#pragma comment(lib, "proto.lib") -#pragma comment(lib, "libboost_system-vc120-mt-1_59.lib") -#pragma comment(lib, "libboost_thread-vc120-mt-1_59.lib") -#pragma comment(lib, "libboost_filesystem-vc120-mt-1_59.lib") -#pragma comment(lib, "glog.lib") -#pragma comment(lib, "gflags.lib") -#pragma comment(lib, "libprotobuf.lib") -#pragma comment(lib, "libhdf5_hl.lib") -#pragma comment(lib, "libhdf5.lib") -#pragma comment(lib, "zlibstatic.lib") -#pragma comment(lib, "libopenblas.dll.a") -#pragma comment(lib, "cudart.lib") -#pragma comment(lib, "curand.lib") -#pragma comment(lib, "cublas.lib") -#pragma comment(lib, "cudnn.lib") - -#pragma comment(lib, "opencv_calib3d249.lib") -#pragma comment(lib, "opencv_contrib249.lib") -#pragma comment(lib, "opencv_core249.lib") -#pragma comment(lib, "opencv_highgui249.lib") -#pragma comment(lib, "opencv_imgproc249.lib") -#endif - -// 入力画像のオフセット -const int offset = 0; -// srcnn.prototxtで定義されたレイヤーの数 -const int layer_num = 7; - -const int ConvertMode = CV_RGB2YUV; -const int ConvertInverseMode = CV_YUV2RGB; - -// 最低限必要なCUDAドライバーのバージョン -const int MinCudaDriverVersion = 6050; - -static std::once_flag waifu2x_once_flag; -static std::once_flag waifu2x_cudnn_once_flag; -static std::once_flag waifu2x_cuda_once_flag; - -#ifndef CUDA_CHECK_WAIFU2X -#define CUDA_CHECK_WAIFU2X(condition) \ - do { \ - cudaError_t error = condition; \ - if(error != cudaSuccess) throw error; \ - } while (0) -#endif - -#define CUDA_HOST_SAFE_FREE(ptr) \ - do { \ - if (ptr) { \ - cudaFreeHost(ptr); \ - ptr = nullptr; \ - } \ - } while (0) - -#define SAFE_DELETE_WAIFU2X(ptr) \ - do { \ - if (ptr) { \ - delete [] ptr; \ - ptr = nullptr; \ - } \ - } while (0) - -namespace -{ - class IgnoreErrorCV - { - private: - static int handleError(int status, const char* func_name, - const char* err_msg, const char* file_name, - int line, void* userdata) - { - return 0; - } - - public: - IgnoreErrorCV() - { - cv::redirectError(handleError); - } - }; - - IgnoreErrorCV g_IgnoreErrorCV; -} - -Waifu2x::Waifu2x() : is_inited(false), isCuda(false), input_block(nullptr), dummy_data(nullptr), output_block(nullptr) -{ -} - -Waifu2x::~Waifu2x() -{ - destroy(); -} - -// 画像を読み込んで値を0.0f〜1.0fの範囲に変換 -Waifu2x::eWaifu2xError Waifu2x::LoadMat(cv::Mat &float_image, const uint32_t* source, int width, int height) -{ - float_image = cv::Mat(cv::Size(width, height), CV_MAKETYPE(CV_8U, 4)); - - const auto LinePixel = float_image.step1() / float_image.channels(); - const auto Channel = float_image.channels(); - const auto Width = float_image.size().width; - const auto Height = float_image.size().height; - - const uint8_t *sptr = (const uint8_t *)source; - auto ptr = float_image.data; - for (int i = 0; i < height; i++) - { - for (int j = 0; j < width; j++) - { - for (int ch = 0; ch < Channel; ch++) - ptr[(i * LinePixel + j) * 4 + ch] = sptr[(i * width + j) * 4 + ch]; - } - } - - // RGBだからBGRに変換 - /* - for (int i = 0; i < height; i++) - { - for (int j = 0; j < width; j++) - std::swap(ptr[(i * LinePixel + j) * 4 + 0], ptr[(i * LinePixel + j) * 4 + 2]); - } - */ - - cv::Mat convert; - float_image.convertTo(convert, CV_32F, 1.0 / 255.0); - float_image.release(); - - { - // アルファチャンネル付きだったらα乗算済みにする - - std::vector planes; - cv::split(convert, planes); - - cv::Mat w = planes[3]; - - planes[0] = planes[0].mul(w); - planes[1] = planes[1].mul(w); - planes[2] = planes[2].mul(w); - - cv::merge(planes, convert); - } - - float_image = convert; - - return eWaifu2xError_OK; -} - -// 入力画像の(Photoshopでいう)キャンバスサイズをoutput_sizeの倍数に変更 -// 画像は左上配置、余白はcv::BORDER_REPLICATEで埋める -Waifu2x::eWaifu2xError Waifu2x::PaddingImage(const cv::Mat &input, cv::Mat &output) -{ - const auto h_blocks = (int)floor(input.size().width / output_size) + (input.size().width % output_size == 0 ? 0 : 1); - const auto w_blocks = (int)floor(input.size().height / output_size) + (input.size().height % output_size == 0 ? 0 : 1); - const auto height = offset + h_blocks * output_size + offset; - const auto width = offset + w_blocks * output_size + offset; - const auto pad_h1 = offset; - const auto pad_w1 = offset; - const auto pad_h2 = (height - offset) - input.size().width; - const auto pad_w2 = (width - offset) - input.size().height; - - cv::copyMakeBorder(input, output, pad_w1, pad_w2, pad_h1, pad_h2, cv::BORDER_REPLICATE); - - return eWaifu2xError_OK; -} - -// 画像をcv::INTER_NEARESTで二倍に拡大して、PaddingImage()でパディングする -Waifu2x::eWaifu2xError Waifu2x::Zoom2xAndPaddingImage(const cv::Mat &input, cv::Mat &output, cv::Size_ &zoom_size) -{ - zoom_size = input.size(); - zoom_size.width *= 2; - zoom_size.height *= 2; - - cv::Mat zoom_image; - cv::resize(input, zoom_image, zoom_size, 0.0, 0.0, cv::INTER_NEAREST); - - return PaddingImage(zoom_image, output); -} - -// 入力画像をzoom_sizeの大きさにcv::INTER_CUBICで拡大し、色情報のみを残す -Waifu2x::eWaifu2xError Waifu2x::CreateZoomColorImage(const cv::Mat &float_image, const cv::Size_ &zoom_size, std::vector &cubic_planes) -{ - cv::Mat zoom_cubic_image; - cv::resize(float_image, zoom_cubic_image, zoom_size, 0.0, 0.0, cv::INTER_CUBIC); - - cv::Mat converted_cubic_image; - cv::cvtColor(zoom_cubic_image, converted_cubic_image, ConvertMode); - zoom_cubic_image.release(); - - cv::split(converted_cubic_image, cubic_planes); - converted_cubic_image.release(); - - // このY成分は使わないので解放 - cubic_planes[0].release(); - - return eWaifu2xError_OK; -} - -// モデルファイルからネットワークを構築 -// 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) -{ - 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); - - caffe::NetParameter param; - if (isModelExist && isModelBinExist && caffe::ReadProtoFromBinaryFile(modelbin_path, ¶m)) - { - const auto ret = SetParameter(param); - if (ret != eWaifu2xError_OK) - return ret; - - net = boost::shared_ptr>(new caffe::Net(param)); - net->CopyTrainedLayersFrom(caffemodel_path); - - input_plane = param.input_dim(1); - } - else - return eWaifu2xError_FailedConstructModel; - - return eWaifu2xError_OK; -} - -Waifu2x::eWaifu2xError Waifu2x::SetParameter(caffe::NetParameter ¶m) const -{ - param.mutable_state()->set_phase(caffe::TEST); - - { - auto mid = param.mutable_input_dim(); - - if (mid->size() != 4) - return eWaifu2xError_FailedParseModelFile; - - *mid->Mutable(0) = batch_size; - *mid->Mutable(2) = input_block_size; - *mid->Mutable(3) = input_block_size; - } - - for (int i = 0; i < param.layer_size(); i++) - { - caffe::LayerParameter *layer_param = param.mutable_layer(i); - const std::string& type = layer_param->type(); - if (type == "Convolution") - { - if (process == "cudnn") - layer_param->mutable_convolution_param()->set_engine(caffe::ConvolutionParameter_Engine_CUDNN); - else - layer_param->mutable_convolution_param()->set_engine(caffe::ConvolutionParameter_Engine_CAFFE); - } - else if (type == "ReLU") - { - if (process == "cudnn") - layer_param->mutable_relu_param()->set_engine(caffe::ReLUParameter_Engine_CUDNN); - else - layer_param->mutable_relu_param()->set_engine(caffe::ReLUParameter_Engine_CAFFE); - } - } - - return eWaifu2xError_OK; -} - -// ネットワークを使って画像を再構築する -Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr> net, cv::Mat &im) -{ - const auto Height = im.size().height; - const auto Width = im.size().width; - const auto Line = im.step1(); - - assert(Width % output_size == 0); - assert(Height % output_size == 0); - - assert(im.channels() == 1 || im.channels() == 3); - - cv::Mat outim(im.rows, im.cols, im.type()); - - // float *imptr = (float *)im.data; - float *imptr = (float *)outim.data; - - try - { - auto input_blobs = net->input_blobs(); - auto input_blob = net->input_blobs()[0]; - - input_blob->Reshape(batch_size, input_plane, input_block_size, input_block_size); - - assert(im.channels() == input_plane); - assert(input_blob->shape(1) == input_plane); - - const int WidthNum = Width / output_size; - const int HeightNum = Height / output_size; - - const int BlockNum = WidthNum * HeightNum; - - const int input_block_plane_size = input_block_size * input_block_size * input_plane; - const int output_block_plane_size = output_block_size * output_block_size * input_plane; - - const int output_padding = inner_padding + outer_padding - layer_num; - - // 画像は(消費メモリの都合上)output_size*output_sizeに分けて再構築する - for (int num = 0; num < BlockNum; num += batch_size) - { - const int processNum = (BlockNum - num) >= batch_size ? batch_size : BlockNum - num; - - if (processNum < batch_size) - input_blob->Reshape(processNum, input_plane, input_block_size, input_block_size); - - for (int n = 0; n < processNum; n++) - { - const int wn = (num + n) % WidthNum; - const int hn = (num + n) / WidthNum; - - const int w = wn * output_size; - const int h = hn * output_size; - - if (w + crop_size <= Width && h + crop_size <= Height) - { - int x, y; - x = w - inner_padding; - y = h - inner_padding; - - int width, height; - - width = crop_size + inner_padding * 2; - height = crop_size + inner_padding * 2; - - int top, bottom, left, right; - - top = outer_padding; - bottom = outer_padding; - left = outer_padding; - right = outer_padding; - - if (x < 0) - { - left += -x; - width -= -x; - x = 0; - } - - if (x + width > Width) - { - right += (x + width) - Width; - width = Width - x; - } - - if (y < 0) - { - top += -y; - height -= -y; - y = 0; - } - - if (y + height > Height) - { - bottom += (y + height) - Height; - height = Height - y; - } - - cv::Mat someimg = im(cv::Rect(x, y, width, height)); - - cv::Mat someborderimg; - // 画像を中央にパディング。余白はcv::BORDER_REPLICATEで埋める - // 実はimで画素が存在する部分は余白と認識されないが、inner_paddingがlayer_numでouter_paddingが1以上ならそこの部分の画素は結果画像として取り出す部分には影響しない - cv::copyMakeBorder(someimg, someborderimg, top, bottom, left, right, cv::BORDER_REPLICATE); - someimg.release(); - - // 画像を直列に変換 - { - float *fptr = input_block + (input_block_plane_size * n); - const float *uptr = (const float *)someborderimg.data; - - const auto Line = someborderimg.step1(); - - if (someborderimg.channels() == 1) - { - if (input_block_size == Line) - memcpy(fptr, uptr, input_block_size * input_block_size * sizeof(float)); - else - { - for (int i = 0; i < input_block_size; i++) - memcpy(fptr + i * input_block_size, uptr + i * Line, input_block_size * sizeof(float)); - } - } - else - { - const auto LinePixel = someborderimg.step1() / someborderimg.channels(); - const auto Channel = someborderimg.channels(); - const auto Width = someborderimg.size().width; - const auto Height = someborderimg.size().height; - - for (int i = 0; i < Height; i++) - { - for (int j = 0; j < LinePixel; j++) - { - for (int ch = 0; ch < Channel; ch++) - fptr[(ch * Height + i) * Width + j] = uptr[(i * LinePixel + j) * Channel + ch]; - } - } - - /* - { - cv::Mat im(someborderimg.size(), CV_32F, fptr, Width * sizeof(float)); - - cv::Mat write_iamge; - im.convertTo(write_iamge, CV_8U, 255.0); - im.release(); - - if (!cv::imwrite("test_in.png", write_iamge)) - return eWaifu2xError_FailedOpenOutputFile; - } - */ - } - } - } - } - - assert(input_blob->count() == input_block_plane_size * processNum); - - // ネットワークに画像を入力 - input_blob->set_cpu_data(input_block); - - // 計算 - auto out = net->ForwardPrefilled(nullptr); - - auto b = out[0]; - - assert(b->count() == output_block_plane_size * processNum); - - const float *ptr = nullptr; - - if (caffe::Caffe::mode() == caffe::Caffe::CPU) - ptr = b->cpu_data(); - else - ptr = b->gpu_data(); - - caffe::caffe_copy(output_block_plane_size * processNum, ptr, output_block); - - for (int n = 0; n < processNum; n++) - { - const int wn = (num + n) % WidthNum; - const int hn = (num + n) / WidthNum; - - const int w = wn * output_size; - const int h = hn * output_size; - - const float *fptr = output_block + (output_block_plane_size * n); - - // 結果を出力画像にコピー - if (outim.channels() == 1) - { - for (int i = 0; i < crop_size; i++) - memcpy(imptr + (h + i) * Line + w, fptr + (i + output_padding) * output_block_size + output_padding, crop_size * sizeof(float)); - } - else - { - const auto LinePixel = outim.step1() / outim.channels(); - const auto Channel = outim.channels(); - - for (int i = 0; i < crop_size; i++) - { - for (int j = 0; j < crop_size; j++) - { - for (int ch = 0; ch < Channel; ch++) - imptr[((h + i) * LinePixel + (w + j)) * Channel + ch] = fptr[(ch * output_block_size + i + output_padding) * output_block_size + j + output_padding]; - } - } - - /* - { - cv::Mat im(someborderimg.size(), CV_32F, fptr, Width * sizeof(float)); - - cv::Mat write_iamge; - im.convertTo(write_iamge, CV_8U, 255.0); - im.release(); - - if (!cv::imwrite("test_in.png", write_iamge)) - return eWaifu2xError_FailedOpenOutputFile; - } - */ - } - } - } - } - catch (...) - { - return eWaifu2xError_FailedProcessCaffe; - } - - im = outim; - - return eWaifu2xError_OK; -} - -Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &Mode, const int NoiseLevel, const std::string &ModelDir, const std::string &Process, - const int CropSize, const int BatchSize) -{ - Waifu2x::eWaifu2xError ret; - - if (is_inited) - return eWaifu2xError_OK; - - try - { - mode = Mode; - noise_level = NoiseLevel; - model_dir = ModelDir; - process = Process; - - crop_size = CropSize; - batch_size = BatchSize; - - inner_padding = layer_num; - outer_padding = 1; - - output_size = crop_size - offset * 2; - input_block_size = crop_size + (inner_padding + outer_padding) * 2; - original_width_height = 128 + layer_num * 2; - - output_block_size = crop_size + (inner_padding + outer_padding - layer_num) * 2; - - std::call_once(waifu2x_once_flag, [argc, argv]() - { - assert(argc >= 1); - - int tmpargc = 1; - char* tmpargvv[] = { argv[0] }; - char** tmpargv = tmpargvv; - // glog等の初期化 - caffe::GlobalInit(&tmpargc, &tmpargv); - }); - - const auto cuDNNCheckStartTime = std::chrono::system_clock::now(); - - if (process == "gpu") - process = "cudnn"; - - const auto cuDNNCheckEndTime = std::chrono::system_clock::now(); - - boost::filesystem::path mode_dir_path(model_dir); - if (!mode_dir_path.is_absolute()) // model_dirが相対パスなら絶対パスに直す - { - // まずはカレントディレクトリ下にあるか探す - mode_dir_path = boost::filesystem::absolute(model_dir); - if (!boost::filesystem::exists(mode_dir_path) && argc >= 1) // 無かったらargv[0]から実行ファイルのあるフォルダを推定し、そのフォルダ下にあるか探す - { - boost::filesystem::path a0(argv[0]); - if (a0.is_absolute()) - mode_dir_path = a0.branch_path() / model_dir; - } - } - - if (!boost::filesystem::exists(mode_dir_path)) - return eWaifu2xError_FailedOpenModelFile; - - if (process == "cpu") - { - caffe::Caffe::set_mode(caffe::Caffe::CPU); - isCuda = false; - } - else - { - caffe::Caffe::set_mode(caffe::Caffe::GPU); - isCuda = true; - } - - if (mode == "noise" || mode == "noise_scale" || mode == "auto_scale") - { - const std::string model_path = (mode_dir_path / "srcnn.prototxt").string(); - const std::string param_path = (mode_dir_path / ("noise" + std::to_string(noise_level) + "_model.json")).string(); - - ret = ConstractNet(net_noise, model_path, param_path, process); - if (ret != eWaifu2xError_OK) - return ret; - } - - if (mode == "scale" || mode == "noise_scale" || mode == "auto_scale") - { - const std::string model_path = (mode_dir_path / "srcnn.prototxt").string(); - const std::string param_path = (mode_dir_path / "scale2.0x_model.json").string(); - - ret = ConstractNet(net_scale, model_path, param_path, process); - if (ret != eWaifu2xError_OK) - return ret; - } - - const int input_block_plane_size = input_block_size * input_block_size * input_plane; - const int output_block_plane_size = output_block_size * output_block_size * input_plane; - - if (isCuda) - { - CUDA_CHECK_WAIFU2X(cudaHostAlloc(&input_block, sizeof(float) * input_block_plane_size * batch_size, cudaHostAllocWriteCombined)); - CUDA_CHECK_WAIFU2X(cudaHostAlloc(&dummy_data, sizeof(float) * input_block_plane_size * batch_size, cudaHostAllocWriteCombined)); - CUDA_CHECK_WAIFU2X(cudaHostAlloc(&output_block, sizeof(float) * output_block_plane_size * batch_size, cudaHostAllocDefault)); - } - else - { - input_block = new float[input_block_plane_size * batch_size]; - dummy_data = new float[input_block_plane_size * batch_size]; - output_block = new float[output_block_plane_size * batch_size]; - } - - for (size_t i = 0; i < input_block_plane_size * batch_size; i++) - dummy_data[i] = 0.0f; - - is_inited = true; - } - catch (...) - { - return eWaifu2xError_InvalidParameter; - } - - return eWaifu2xError_OK; -} - -void Waifu2x::destroy() -{ - net_noise.reset(); - net_scale.reset(); - - if (isCuda) - { - CUDA_HOST_SAFE_FREE(input_block); - CUDA_HOST_SAFE_FREE(dummy_data); - CUDA_HOST_SAFE_FREE(output_block); - } - else - { - SAFE_DELETE_WAIFU2X(input_block); - SAFE_DELETE_WAIFU2X(dummy_data); - SAFE_DELETE_WAIFU2X(output_block); - } - - is_inited = false; -} - -Waifu2x::eWaifu2xError Waifu2x::waifu2x(int factor, const uint32_t* source, uint32_t* dest, int width, int height) -{ - Waifu2x::eWaifu2xError ret; - - if (!is_inited) - return eWaifu2xError_NotInitialized; - - cv::Mat float_image; - ret = LoadMat(float_image, source, width, height); - if (ret != eWaifu2xError_OK) - return ret; - - cv::Mat im; - if (input_plane == 1) - return eWaifu2xError_NotInitialized; - else - { - std::vector planes; - cv::split(float_image, planes); - - if (float_image.channels() == 4) - planes.resize(3); - - // BGRからRGBにする - //std::swap(planes[0], planes[2]); - - cv::merge(planes, im); - } - cv::Size_ image_size = im.size(); - - const bool isReconstructNoise = mode == "noise" || mode == "noise_scale" || mode == "auto_scale"; - const bool isReconstructScale = mode == "scale" || mode == "noise_scale"; - - if (isReconstructNoise) - { - PaddingImage(im, im); - - ret = ReconstructImage(net_noise, im); - if (ret != eWaifu2xError_OK) - return ret; - - // パディングを取り払う - im = im(cv::Rect(offset, offset, image_size.width, image_size.height)); - } - - const int scale2 = ceil(log2((double)factor)); - const double shrinkRatio = (double)factor / std::pow(2.0, (double)scale2); - - if (isReconstructScale) - { - bool isError = false; - for (int i = 0; i < scale2; i++) - { - Zoom2xAndPaddingImage(im, im, image_size); - - ret = ReconstructImage(net_scale, im); - if (ret != eWaifu2xError_OK) - return ret; - - // パディングを取り払う - im = im(cv::Rect(offset, offset, image_size.width, image_size.height)); - } - } - - cv::Mat process_image; - if (input_plane == 1) - { - // 再構築した輝度画像とCreateZoomColorImage()で作成した色情報をマージして通常の画像に変換し、書き込む - - std::vector color_planes; - CreateZoomColorImage(float_image, image_size, color_planes); - - float_image.release(); - - color_planes[0] = im; - im.release(); - - cv::Mat converted_image; - cv::merge(color_planes, converted_image); - color_planes.clear(); - - cv::cvtColor(converted_image, process_image, ConvertInverseMode); - converted_image.release(); - } - else - { - std::vector planes; - cv::split(im, planes); - - // RGBからBGRに直す - //std::swap(planes[0], planes[2]); - - cv::merge(planes, process_image); - } - - cv::Mat alpha; - if (float_image.channels() == 4) - { - std::vector planes; - cv::split(float_image, planes); - alpha = planes[3]; - - cv::resize(alpha, alpha, image_size, 0.0, 0.0, cv::INTER_CUBIC); - } - - // アルファチャンネルがあったら、アルファを付加してカラーからアルファの影響を抜く - if (!alpha.empty()) - { - std::vector planes; - cv::split(process_image, planes); - process_image.release(); - - planes.push_back(alpha); - - cv::Mat w2 = planes[3]; - - planes[0] = (planes[0]).mul(1.0 / w2); - planes[1] = (planes[1]).mul(1.0 / w2); - planes[2] = (planes[2]).mul(1.0 / w2); - - cv::merge(planes, process_image); - } - - const cv::Size_ ns(image_size.width * shrinkRatio, image_size.height * shrinkRatio); - if (image_size.width != ns.width || image_size.height != ns.height) - cv::resize(process_image, process_image, ns, 0.0, 0.0, cv::INTER_LINEAR); - - cv::Mat write_iamge; - process_image.convertTo(write_iamge, CV_8U, 255.0); - process_image.release(); - - /* - ret = WriteMat(write_iamge, output_file); - if (ret != eWaifu2xError_OK) - return ret; - - write_iamge.release(); - */ - - { - const auto width = write_iamge.size().width; - const auto stride = write_iamge.step1(); - for (int i = 0; i < write_iamge.size().height; i++) - memcpy(dest + width * i, write_iamge.data + stride * i, stride); - } - - return eWaifu2xError_OK; -} - -const std::string& Waifu2x::used_process() const -{ - return process; -} diff --git a/waifu2x-caffe-dll/waifu2x.h b/waifu2x-caffe-dll/waifu2x.h deleted file mode 100644 index ebe0810..0000000 --- a/waifu2x-caffe-dll/waifu2x.h +++ /dev/null @@ -1,115 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include -#include -#include - - -namespace caffe -{ - template - class Net; - class NetParameter; -}; - -class Waifu2x -{ -public: - enum eWaifu2xError - { - eWaifu2xError_OK = 0, - eWaifu2xError_Cancel, - eWaifu2xError_NotInitialized, - eWaifu2xError_InvalidParameter, - eWaifu2xError_FailedOpenInputFile, - eWaifu2xError_FailedOpenOutputFile, - eWaifu2xError_FailedOpenModelFile, - eWaifu2xError_FailedParseModelFile, - eWaifu2xError_FailedConstructModel, - eWaifu2xError_FailedProcessCaffe, - eWaifu2xError_FailedCudaCheck, - }; - - enum eWaifu2xCudaError - { - eWaifu2xCudaError_OK = 0, - eWaifu2xCudaError_NotFind, - eWaifu2xCudaError_OldVersion, - }; - - enum eWaifu2xcuDNNError - { - eWaifu2xcuDNNError_OK = 0, - eWaifu2xcuDNNError_NotFind, - eWaifu2xcuDNNError_OldVersion, - eWaifu2xcuDNNError_CannotCreate, - }; - - typedef std::function waifu2xCancelFunc; - -private: - bool is_inited; - - // 一度に処理する画像の幅 - int crop_size; - // 一度に何ブロック分処理するか - int batch_size; - - // ネットに入力する画像のサイズ - int input_block_size; - // ブロック変換後の出力サイズ - int output_size; - // ネットワークに入力する画像のサイズ(出力画像の幅はlayer_num * 2だけ小さくなる) - int block_width_height; - // srcnn.prototxtで定義された入力する画像のサイズ - int original_width_height; - - std::string mode; - int noise_level; - std::string model_dir; - std::string process; - - int inner_padding; - int outer_padding; - - int output_block_size; - - int input_plane; - - bool isCuda; - - boost::shared_ptr> net_noise; - boost::shared_ptr> net_scale; - - float *input_block; - float *dummy_data; - float *output_block; - -private: - eWaifu2xError LoadMat(cv::Mat &float_image, const uint32_t* source, int width, int height); - 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 SetParameter(caffe::NetParameter ¶m) const; - eWaifu2xError ReconstructImage(boost::shared_ptr> net, cv::Mat &im); - -public: - Waifu2x(); - ~Waifu2x(); - - // 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 std::string &model_dir, const std::string &process, - const int crop_size = 128, const int batch_size = 1); - - void destroy(); - - eWaifu2xError waifu2x(int factor, const uint32_t* source, uint32_t* dest, int width, int height); - - const std::string& used_process() const; -};