From 9d32a969bf1ff371a0482a8c2a8d01a6cc94f262 Mon Sep 17 00:00:00 2001 From: lltcggie Date: Fri, 17 Jul 2015 03:26:35 +0900 Subject: [PATCH] =?UTF-8?q?=E9=81=A9=E5=BD=93=E3=81=AAdll=E3=81=AE?= =?UTF-8?q?=E3=83=97=E3=83=AD=E3=82=B8=E3=82=A7=E3=82=AF=E3=83=88=E3=82=92?= =?UTF-8?q?=E8=BF=BD=E5=8A=A0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- waifu2x-caffe-dll/Source.cpp | 41 + waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj | 96 +++ .../waifu2x-caffe-dll.vcxproj.filters | 30 + waifu2x-caffe-dll/waifu2x.cpp | 794 ++++++++++++++++++ waifu2x-caffe-dll/waifu2x.h | 115 +++ waifu2x-caffe.sln | 6 + 6 files changed, 1082 insertions(+) create mode 100644 waifu2x-caffe-dll/Source.cpp create mode 100644 waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj create mode 100644 waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj.filters create mode 100644 waifu2x-caffe-dll/waifu2x.cpp create mode 100644 waifu2x-caffe-dll/waifu2x.h diff --git a/waifu2x-caffe-dll/Source.cpp b/waifu2x-caffe-dll/Source.cpp new file mode 100644 index 0000000..963606c --- /dev/null +++ b/waifu2x-caffe-dll/Source.cpp @@ -0,0 +1,41 @@ +#include +#include +#include "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) +{ + Waifu2x *obj = new Waifu2x(); + + char *argv[] = { "" }; + + if (obj->init(1, argv, mode, noise_level, model_dir, process, crop_size, batch_size) != Waifu2x::eWaifu2xError_OK) + { + delete obj; + return nullptr; + } + + return obj; +} + +__declspec(dllexport) +bool Waifu2xProcess(void *waifu2xObj, int factor, const uint32_t* source, uint32_t* dest, int width, int height) +{ + if (!waifu2xObj) + return false; + + Waifu2x *obj = (Waifu2x *)waifu2xObj; + + return obj->waifu2x(factor, source, dest, width, height) == Waifu2x::eWaifu2xError_OK; +} + +__declspec(dllexport) +void Waifu2xDestory(void *waifu2xObj) +{ + if (waifu2xObj) + { + Waifu2x *obj = (Waifu2x *)waifu2xObj; + delete obj; + } +} diff --git a/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj b/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj new file mode 100644 index 0000000..530214f --- /dev/null +++ b/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj @@ -0,0 +1,96 @@ + + + + + Debug + x64 + + + Release + x64 + + + + {DFF94FEB-78AB-41B1-9B92-4D8B7D799E04} + Win32Proj + waifu2x-caffe-dll + + + + DynamicLibrary + true + v120 + Unicode + + + DynamicLibrary + false + v120 + true + Unicode + + + + + + + + + + + + + true + $(SolutionDir)caffe\build\include;$(SolutionDir)caffe\3rdparty\include;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v7.0\include;C:\boost_1_56_0;C:\opencv249\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)stb;$(SolutionDir)include;$(IncludePath) + $(SolutionDir)caffe\build\lib;$(SolutionDir)caffe\3rdparty\lib;C:\boost_1_56_0\lib64-msvc-12.0;$(LibraryPath) + waifu2x-caffe + + + false + $(SolutionDir)caffe\build\include;$(SolutionDir)caffe\3rdparty\include;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v7.0\include;C:\boost_1_56_0;C:\opencv249\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)stb;$(SolutionDir)include;$(IncludePath) + $(SolutionDir)caffe\build\lib;$(SolutionDir)caffe\3rdparty\lib;C:\boost_1_56_0\lib64-msvc-12.0;$(LibraryPath) + waifu2x-caffe + + + + + + Level3 + Disabled + WIN32;_DEBUG;_CONSOLE;%(PreprocessorDefinitions) + + + Console + true + cudnn64_65.dll;%(DelayLoadDLLs) + + + + + Level3 + + + MaxSpeed + true + true + WIN32;NDEBUG;_CONSOLE;%(PreprocessorDefinitions) + + + Console + true + true + true + cudnn64_65.dll;%(DelayLoadDLLs) + + + + + + + + + + + + + \ No newline at end of file diff --git a/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj.filters b/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj.filters new file mode 100644 index 0000000..c230057 --- /dev/null +++ b/waifu2x-caffe-dll/waifu2x-caffe-dll.vcxproj.filters @@ -0,0 +1,30 @@ + + + + + {4FC737F1-C7A5-4376-A066-2A32D752A2FF} + cpp;c;cc;cxx;def;odl;idl;hpj;bat;asm;asmx + + + {93995380-89BD-4b04-88EB-625FBE52EBFB} + h;hh;hpp;hxx;hm;inl;inc;xsd + + + {67DA6AB6-F800-4c08-8B7A-83BB121AAD01} + rc;ico;cur;bmp;dlg;rc2;rct;bin;rgs;gif;jpg;jpeg;jpe;resx;tiff;tif;png;wav;mfcribbon-ms + + + + + ソース ファイル + + + ソース ファイル + + + + + ヘッダー ファイル + + + \ No newline at end of file diff --git a/waifu2x-caffe-dll/waifu2x.cpp b/waifu2x-caffe-dll/waifu2x.cpp new file mode 100644 index 0000000..cbf9858 --- /dev/null +++ b/waifu2x-caffe-dll/waifu2x.cpp @@ -0,0 +1,794 @@ +#include "waifu2x.h" +#include +#include +#include +#include +#include +#include +#include +#include + + +#ifdef _MSC_VER +#ifdef _DEBUG +#pragma comment(lib, "libcaffed.lib") +#pragma comment(lib, "libprotobufd.lib") +#else +#pragma comment(lib, "libcaffe.lib") +#pragma comment(lib, "libprotobuf.lib") +#endif +#pragma comment(lib, "libprotoc.lib") +#endif + +// ͉摜̃ItZbg +const int offset = 0; +// srcnn.prototxtŒ`ꂽC[̐ +const int layer_num = 7; + +const int ConvertMode = CV_RGB2YUV; +const int ConvertInverseMode = CV_YUV2RGB; + +// ŒKvCUDAhCo[̃o[W +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(); +} + +// 摜ǂݍŒl0.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]; + } + } + + // RGBBGRɕϊ + /* + 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(); + + { + // At@`lt烿Zς݂ɂ + + 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ł)LoXTCYoutput_size̔{ɕύX +// 摜͍zuA]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œ{Ɋg債āAPaddingImage()ŃpfBO +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_sizȇ傫cv::INTER_CUBICŊg債AF݂̂c +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͎gȂ̂ʼn + cubic_planes[0].release(); + + return eWaifu2xError_OK; +} + +// ft@Clbg[N\z +// processcudnnw肳ȂꍇcuDNNĂяoȂ悤ɕύX +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; +} + +// lbg[Ngĉ摜č\z +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; + + // 摜(̓s)output_size*output_sizeɕčč\z + 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; + // 摜𒆉ɃpfBOB]cv::BORDER_REPLICATEŖ߂ + // imʼnf݂镔͗]ƔFȂAinner_paddinglayer_numouter_padding1ȏȂ炻̉̕f͌ʉ摜ƂĎoɂ͉eȂ + 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); + + // lbg[Nɉ摜 + input_blob->set_cpu_data(input_block); + + // vZ + 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); + + // ʂo͉摜ɃRs[ + 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΃pXȂ΃pXɒ + { + // ܂̓JgfBNgɂ邩T + mode_dir_path = boost::filesystem::absolute(model_dir); + if (!boost::filesystem::exists(mode_dir_path) && argc >= 1) // argv[0]st@ĈtH_𐄒肵ÃtH_ɂ邩T + { + 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); + + // BGRRGBɂ + //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; + + // pfBO蕥 + 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; + + // pfBO蕥 + im = im(cv::Rect(offset, offset, image_size.width, image_size.height)); + } + } + + cv::Mat process_image; + if (input_plane == 1) + { + // č\zPx摜CreateZoomColorImage()ō쐬F}[WĒʏ̉摜ɕϊA + + 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); + + // RGBBGRɒ + //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); + } + + // At@`lAAt@tăJ[At@̉e𔲂 + 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 new file mode 100644 index 0000000..ebe0810 --- /dev/null +++ b/waifu2x-caffe-dll/waifu2x.h @@ -0,0 +1,115 @@ +#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; + + // xɏ摜̕ + int crop_size; + // xɉubN邩 + int batch_size; + + // lbgɓ͂摜̃TCY + int input_block_size; + // ubNϊ̏o̓TCY + int output_size; + // lbg[Nɓ͂摜̃TCY(o͉摜̕layer_num * 2Ȃ) + int block_width_height; + // srcnn.prototxtŒ`ꂽ͂摜̃TCY + 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; +}; diff --git a/waifu2x-caffe.sln b/waifu2x-caffe.sln index 4d49403..d00739a 100644 --- a/waifu2x-caffe.sln +++ b/waifu2x-caffe.sln @@ -7,6 +7,8 @@ Project("{8BC9CEB8-8B4A-11D0-8D11-00A0C91BC942}") = "waifu2x-caffe", "waifu2x-ca EndProject Project("{8BC9CEB8-8B4A-11D0-8D11-00A0C91BC942}") = "waifu2x-caffe-gui", "waifu2x-caffe-gui\waifu2x-caffe-gui.vcxproj", "{63FB3EFC-63B0-401C-BB54-F3A984DC233F}" EndProject +Project("{8BC9CEB8-8B4A-11D0-8D11-00A0C91BC942}") = "waifu2x-caffe-dll", "waifu2x-caffe-dll\waifu2x-caffe-dll.vcxproj", "{DFF94FEB-78AB-41B1-9B92-4D8B7D799E04}" +EndProject Global GlobalSection(SolutionConfigurationPlatforms) = preSolution Debug|x64 = Debug|x64 @@ -21,6 +23,10 @@ Global {63FB3EFC-63B0-401C-BB54-F3A984DC233F}.Debug|x64.Build.0 = Debug|x64 {63FB3EFC-63B0-401C-BB54-F3A984DC233F}.Release|x64.ActiveCfg = Release|x64 {63FB3EFC-63B0-401C-BB54-F3A984DC233F}.Release|x64.Build.0 = Release|x64 + {DFF94FEB-78AB-41B1-9B92-4D8B7D799E04}.Debug|x64.ActiveCfg = Debug|x64 + {DFF94FEB-78AB-41B1-9B92-4D8B7D799E04}.Debug|x64.Build.0 = Debug|x64 + {DFF94FEB-78AB-41B1-9B92-4D8B7D799E04}.Release|x64.ActiveCfg = Release|x64 + {DFF94FEB-78AB-41B1-9B92-4D8B7D799E04}.Release|x64.Build.0 = Release|x64 EndGlobalSection GlobalSection(SolutionProperties) = preSolution HideSolutionNode = FALSE