diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..1ff0c42 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,63 @@ +############################################################################### +# Set default behavior to automatically normalize line endings. +############################################################################### +* text=auto + +############################################################################### +# Set default behavior for command prompt diff. +# +# This is need for earlier builds of msysgit that does not have it on by +# default for csharp files. +# Note: This is only used by command line +############################################################################### +#*.cs diff=csharp + +############################################################################### +# Set the merge driver for project and solution files +# +# Merging from the command prompt will add diff markers to the files if there +# are conflicts (Merging from VS is not affected by the settings below, in VS +# the diff markers are never inserted). Diff markers may cause the following +# file extensions to fail to load in VS. An alternative would be to treat +# these files as binary and thus will always conflict and require user +# intervention with every merge. To do so, just uncomment the entries below +############################################################################### +#*.sln merge=binary +#*.csproj merge=binary +#*.vbproj merge=binary +#*.vcxproj merge=binary +#*.vcproj merge=binary +#*.dbproj merge=binary +#*.fsproj merge=binary +#*.lsproj merge=binary +#*.wixproj merge=binary +#*.modelproj merge=binary +#*.sqlproj merge=binary +#*.wwaproj merge=binary + +############################################################################### +# behavior for image files +# +# image files are treated as binary by default. +############################################################################### +#*.jpg binary +#*.png binary +#*.gif binary + +############################################################################### +# diff behavior for common document formats +# +# Convert binary document formats to text before diffing them. This feature +# is only available from the command line. Turn it on by uncommenting the +# entries below. +############################################################################### +#*.doc diff=astextplain +#*.DOC diff=astextplain +#*.docx diff=astextplain +#*.DOCX diff=astextplain +#*.dot diff=astextplain +#*.DOT diff=astextplain +#*.pdf diff=astextplain +#*.PDF diff=astextplain +#*.rtf diff=astextplain +#*.RTF diff=astextplain diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3a2238d --- /dev/null +++ b/.gitignore @@ -0,0 +1,245 @@ +## Ignore Visual Studio temporary files, build results, and +## files generated by popular Visual Studio add-ons. + +# User-specific files +*.suo +*.user +*.userosscache +*.sln.docstates + +# User-specific files (MonoDevelop/Xamarin Studio) +*.userprefs + +# Build results +[Dd]ebug/ +[Dd]ebugPublic/ +[Rr]elease/ +[Rr]eleases/ +[Xx]64/ +[Xx]86/ +[Bb]uild/ +bld/ +[Bb]in/ +[Oo]bj/ + +# Visual Studio 2015 cache/options directory +.vs/ +# Uncomment if you have tasks that create the project's static files in wwwroot +#wwwroot/ + +# MSTest test Results +[Tt]est[Rr]esult*/ +[Bb]uild[Ll]og.* + +# NUNIT +*.VisualState.xml +TestResult.xml + +# Build Results of an ATL Project +[Dd]ebugPS/ +[Rr]eleasePS/ +dlldata.c + +# DNX +project.lock.json +artifacts/ + +*_i.c +*_p.c +*_i.h +*.ilk +*.meta +*.obj +*.pch +*.pdb +*.pgc +*.pgd +*.rsp +*.sbr +*.tlb +*.tli +*.tlh +*.tmp +*.tmp_proj +*.log +*.vspscc +*.vssscc +.builds +*.pidb +*.svclog +*.scc + +# Chutzpah Test files +_Chutzpah* + +# Visual C++ cache files +ipch/ +*.aps +*.ncb +*.opendb +*.opensdf +*.sdf +*.cachefile +*.VC.db + +# Visual Studio profiler +*.psess +*.vsp +*.vspx +*.sap + +# TFS 2012 Local Workspace +$tf/ + +# Guidance Automation Toolkit +*.gpState + +# ReSharper is a .NET coding add-in +_ReSharper*/ +*.[Rr]e[Ss]harper +*.DotSettings.user + +# JustCode is a .NET coding add-in +.JustCode + +# TeamCity is a build add-in +_TeamCity* + +# DotCover is a Code Coverage Tool +*.dotCover + +# NCrunch +_NCrunch_* +.*crunch*.local.xml +nCrunchTemp_* + +# MightyMoose +*.mm.* +AutoTest.Net/ + +# Web workbench (sass) +.sass-cache/ + +# Installshield output folder +[Ee]xpress/ + +# DocProject is a documentation generator add-in +DocProject/buildhelp/ +DocProject/Help/*.HxT +DocProject/Help/*.HxC +DocProject/Help/*.hhc +DocProject/Help/*.hhk +DocProject/Help/*.hhp +DocProject/Help/Html2 +DocProject/Help/html + +# Click-Once directory +publish/ + +# Publish Web Output +*.[Pp]ublish.xml +*.azurePubxml + +# TODO: Un-comment the next line if you do not want to checkin +# your web deploy settings because they may include unencrypted +# passwords +#*.pubxml +*.publishproj + +# NuGet Packages +*.nupkg +# The packages folder can be ignored because of Package Restore +**/packages/* +# except build/, which is used as an MSBuild target. +!**/packages/build/ +# Uncomment if necessary however generally it will be regenerated when needed +#!**/packages/repositories.config +# NuGet v3's project.json files produces more ignoreable files +*.nuget.props +*.nuget.targets + +# Microsoft Azure Build Output +csx/ +*.build.csdef + +# Microsoft Azure Emulator +ecf/ +rcf/ + +# Microsoft Azure ApplicationInsights config file +ApplicationInsights.config + +# Windows Store app package directory +AppPackages/ +BundleArtifacts/ + +# Visual Studio cache files +# files ending in .cache can be ignored +*.[Cc]ache +# but keep track of directories ending in .cache +!*.[Cc]ache/ + +# Others +ClientBin/ +[Ss]tyle[Cc]op.* +~$* +*~ +*.dbmdl +*.dbproj.schemaview +*.pfx +*.publishsettings +node_modules/ +orleans.codegen.cs + +# RIA/Silverlight projects +Generated_Code/ + +# Backup & report files from converting an old project file +# to a newer Visual Studio version. Backup files are not needed, +# because we have git ;-) +_UpgradeReport_Files/ +Backup*/ +UpgradeLog*.XML +UpgradeLog*.htm + +# SQL Server files +*.mdf +*.ldf + +# Business Intelligence projects +*.rdl.data +*.bim.layout +*.bim_*.settings + +# Microsoft Fakes +FakesAssemblies/ + +# GhostDoc plugin setting file +*.GhostDoc.xml + +# Node.js Tools for Visual Studio +.ntvs_analysis.dat + +# Visual Studio 6 build log +*.plg + +# Visual Studio 6 workspace options file +*.opt + +# Visual Studio LightSwitch build output +**/*.HTMLClient/GeneratedArtifacts +**/*.DesktopClient/GeneratedArtifacts +**/*.DesktopClient/ModelManifest.xml +**/*.Server/GeneratedArtifacts +**/*.Server/ModelManifest.xml +_Pvt_Extensions + +# LightSwitch generated files +GeneratedArtifacts/ +ModelManifest.xml + +# Paket dependency manager +.paket/paket.exe + +# FAKE - F# Make +.fake/ \ No newline at end of file diff --git a/common/cNet.cpp b/common/cNet.cpp new file mode 100644 index 0000000..8109d8a --- /dev/null +++ b/common/cNet.cpp @@ -0,0 +1,710 @@ +#include "cNet.h" +#include +#include +#include +#include +#include +#include +#include +#include + +const int kProtoReadBytesLimit = INT_MAX; // Max size of 2 GB minus 1 byte. + + +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; + + try + { + is.open(path, std::ios_base::in | std::ios_base::binary); + } + catch (...) + { + return false; + } + + return readFile(is, buf); +} + +static Waifu2x::eWaifu2xError readProtoText(const boost::filesystem::path &path, ::google::protobuf::Message* proto) +{ + boost::iostreams::stream is; + + try + { + is.open(path, std::ios_base::in); + } + catch (...) + { + return Waifu2x::eWaifu2xError_FailedOpenModelFile; + } + + 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 writeProtoBinary(const ::google::protobuf::Message& proto, const boost::filesystem::path &path) +{ + boost::iostreams::stream os; + + try + { + os.open(path, std::ios_base::out | std::ios_base::binary | std::ios_base::trunc); + } + catch (...) + { + return Waifu2x::eWaifu2xError_FailedOpenModelFile; + } + + if (!os) + return Waifu2x::eWaifu2xError_FailedWriteModelFile; + + if (!proto.SerializePartialToOstream(&os)) + return Waifu2x::eWaifu2xError_FailedWriteModelFile; + + return Waifu2x::eWaifu2xError_OK; +} + +static Waifu2x::eWaifu2xError readProtoBinary(const boost::filesystem::path &path, ::google::protobuf::Message* proto) +{ + boost::iostreams::stream is; + + try + { + is.open(path, std::ios_base::in | std::ios_base::binary); + } + catch (...) + { + return Waifu2x::eWaifu2xError_FailedOpenModelFile; + } + + 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()); + + 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; +} + +cNet::cNet() : mModelScale(0), mInnerScale(0), mNetOffset(0), mInputPlane(0) +{} + +cNet::~cNet() +{} + +// モデルファイルからネットワークを構築 +// processでcudnnが指定されなかった場合はcuDNNが呼び出されないように変更する +Waifu2x::eWaifu2xError cNet::ConstractNet(const boost::filesystem::path &model_path, const boost::filesystem::path ¶m_path, const boost::filesystem::path &info_path, const std::string &process) +{ + Waifu2x::eWaifu2xError ret; + + ret = LoadInfoFromJson(info_path); + if (ret != Waifu2x::eWaifu2xError_OK) + return ret; + + boost::filesystem::path modelbin_path = model_path; + modelbin_path += ".protobin"; + boost::filesystem::path caffemodel_path = param_path; + caffemodel_path += ".caffemodel"; + + caffe::NetParameter param_model; + caffe::NetParameter param_caffemodel; + + const auto retModelBin = readProtoBinary(modelbin_path, ¶m_model); + const auto retParamBin = readProtoBinary(caffemodel_path, ¶m_caffemodel); + + if (retModelBin == Waifu2x::eWaifu2xError_OK && retParamBin == Waifu2x::eWaifu2xError_OK) + { + ret = SetParameter(param_model, process); + if (ret != Waifu2x::eWaifu2xError_OK) + return ret; + + if (!caffe::UpgradeNetAsNeeded(caffemodel_path.string(), ¶m_caffemodel)) + return Waifu2x::eWaifu2xError_FailedParseModelFile; + + mNet = boost::shared_ptr>(new caffe::Net(param_model)); + mNet->CopyTrainedLayersFrom(param_caffemodel); + } + else + { + const auto ret = LoadParameterFromJson(model_path, param_path, modelbin_path, caffemodel_path, process); + if (ret != Waifu2x::eWaifu2xError_OK) + return ret; + } + + const auto &inputs = mNet->input_blobs(); + if (inputs.empty()) + return Waifu2x::eWaifu2xError_FailedConstructModel; + + if (mInputPlane != inputs[0]->channels()) + return Waifu2x::eWaifu2xError_FailedConstructModel; + + return Waifu2x::eWaifu2xError_OK; +} + +Waifu2x::eWaifu2xError cNet::LoadInfoFromJson(const boost::filesystem::path &info_path) +{ + rapidjson::Document d; + std::vector jsonBuf; + + try + { + boost::iostreams::stream is; + + try + { + is.open(info_path, std::ios_base::in | std::ios_base::binary); + } + catch (...) + { + return Waifu2x::eWaifu2xError_FailedOpenModelFile; + } + + if (!is) + return Waifu2x::eWaifu2xError_FailedOpenModelFile; + + const size_t size = is.seekg(0, std::ios::end).tellg(); + is.seekg(0, std::ios::beg); + + jsonBuf.resize(size + 1); + is.read(jsonBuf.data(), jsonBuf.size()); + + jsonBuf[jsonBuf.size() - 1] = '\0'; + + d.Parse(jsonBuf.data()); + + const bool resize = d.HasMember("resize") && d["resize"].GetBool() ? true : false; + const auto name = d["name"].GetString(); + const int channels = d["channels"].GetInt(); + const int net_offset = d["offset"].GetInt(); + const int inner_scale = d["scale_factor"].GetInt(); + + mModelScale = 2; // TODO: 動的に設定するようにする + mInnerScale = inner_scale; + mNetOffset = net_offset; + mInputPlane = channels; + } + catch (...) + { + return Waifu2x::eWaifu2xError_FailedParseModelFile; + } + + return Waifu2x::eWaifu2xError_OK; +} + +Waifu2x::eWaifu2xError cNet::SetParameter(caffe::NetParameter ¶m, const std::string &process) const +{ + param.mutable_state()->set_phase(caffe::TEST); + + { + auto input_layer = param.mutable_layer(0); + auto mid = input_layer->mutable_input_param()->mutable_shape(); + if (mid->size() != 1 || mid->Mutable(0)->dim_size() != 4) + return Waifu2x::eWaifu2xError_FailedParseModelFile; + mid->Mutable(0)->set_dim(0, 1); + mid->Mutable(0)->set_dim(2, 142); + mid->Mutable(0)->set_dim(3, 142); + } + + 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 Waifu2x::eWaifu2xError_OK; +} + +Waifu2x::eWaifu2xError cNet::LoadParameterFromJson(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) +{ + Waifu2x::eWaifu2xError ret; + + caffe::NetParameter param; + ret = readProtoText(model_path, ¶m); + if (ret != Waifu2x::eWaifu2xError_OK) + return ret; + + ret = writeProtoBinary(param, modelbin_path); + if (ret != Waifu2x::eWaifu2xError_OK) + return ret; + + ret = SetParameter(param, process); + if (ret != Waifu2x::eWaifu2xError_OK) + return ret; + + mNet = boost::shared_ptr>(new caffe::Net(param)); + + rapidjson::Document d; + std::vector jsonBuf; + + try + { + boost::iostreams::stream is; + + try + { + is.open(param_path, std::ios_base::in | std::ios_base::binary); + } + catch (...) + { + return Waifu2x::eWaifu2xError_FailedOpenModelFile; + } + + if (!is) + return Waifu2x::eWaifu2xError_FailedOpenModelFile; + + const size_t size = is.seekg(0, std::ios::end).tellg(); + is.seekg(0, std::ios::beg); + + jsonBuf.resize(size + 1); + is.read(jsonBuf.data(), jsonBuf.size()); + + jsonBuf[jsonBuf.size() - 1] = '\0'; + + d.Parse(jsonBuf.data()); + } + catch (...) + { + return Waifu2x::eWaifu2xError_FailedParseModelFile; + } + + if (d.Size() != 7) + return Waifu2x::eWaifu2xError_FailedParseModelFile; + + int inputPlane = 0; + int outputPlane = 0; + try + { + inputPlane = d[0]["nInputPlane"].GetInt(); + outputPlane = d[d.Size() - 1]["nOutputPlane"].GetInt(); + } + catch (...) + { + return Waifu2x::eWaifu2xError_FailedParseModelFile; + } + + if (inputPlane == 0 || outputPlane == 0) + return Waifu2x::eWaifu2xError_FailedParseModelFile; + + if (inputPlane != outputPlane) + return Waifu2x::eWaifu2xError_FailedParseModelFile; + + //if (param.layer_size() < 17) + // return Waifu2x::eWaifu2xError_FailedParseModelFile; + + std::vector>> list; + auto &v = mNet->layers(); + for (auto &l : v) + { + auto lk = l->type(); + auto &bv = l->blobs(); + if (bv.size() > 0) + list.push_back(l); + } + + try + { + std::vector weightList; + std::vector biasList; + + int count = 0; + for (auto it = d.Begin(); it != d.End(); ++it) + { + const auto &weight = (*it)["weight"]; + const auto nInputPlane = (*it)["nInputPlane"].GetInt(); + const auto nOutputPlane = (*it)["nOutputPlane"].GetInt(); + const auto kW = (*it)["kW"].GetInt(); + const auto &bias = (*it)["bias"]; + + auto leyer = list[count]; + + auto &b0 = leyer->blobs()[0]; + auto &b1 = leyer->blobs()[1]; + + float *b0Ptr = nullptr; + float *b1Ptr = nullptr; + + if (caffe::Caffe::mode() == caffe::Caffe::CPU) + { + b0Ptr = b0->mutable_cpu_data(); + b1Ptr = b1->mutable_cpu_data(); + } + else + { + b0Ptr = b0->mutable_gpu_data(); + b1Ptr = b1->mutable_gpu_data(); + } + + const auto WeightSize1 = weight.Size(); + const auto WeightSize2 = weight[0].Size(); + const auto KernelHeight = weight[0][0].Size(); + const auto KernelWidth = weight[0][0][0].Size(); + + if (!(b0->count() == WeightSize1 * WeightSize2 * KernelHeight * KernelWidth)) + return Waifu2x::eWaifu2xError_FailedConstructModel; + + if (!(b1->count() == bias.Size())) + return Waifu2x::eWaifu2xError_FailedConstructModel; + + weightList.resize(0); + biasList.resize(0); + + size_t weightCount = 0; + for (auto it2 = weight.Begin(); it2 != weight.End(); ++it2) + { + for (auto it3 = (*it2).Begin(); it3 != (*it2).End(); ++it3) + { + for (auto it4 = (*it3).Begin(); it4 != (*it3).End(); ++it4) + { + for (auto it5 = (*it4).Begin(); it5 != (*it4).End(); ++it5) + weightList.push_back((float)it5->GetDouble()); + } + } + } + + caffe::caffe_copy(b0->count(), weightList.data(), b0Ptr); + + for (auto it2 = bias.Begin(); it2 != bias.End(); ++it2) + biasList.push_back((float)it2->GetDouble()); + + caffe::caffe_copy(b1->count(), biasList.data(), b1Ptr); + + count++; + } + + mNet->ToProto(¶m); + + ret = writeProtoBinary(param, caffemodel_path); + if (ret != Waifu2x::eWaifu2xError_OK) + return ret; + } + catch (...) + { + return Waifu2x::eWaifu2xError_FailedConstructModel; + } + + return Waifu2x::eWaifu2xError_OK; +} + +int cNet::GetInputPlane() const +{ + return mInputPlane; +} + +int cNet::GetInnerScale() const +{ + return mInnerScale; +} + +int cNet::GetNetOffset() const +{ + return mNetOffset; +} + +int cNet::GetScale() const +{ + return mModelScale; +} + +int cNet::GetInputMemorySize(const int crop_w, const int crop_h, const int outer_padding, const int batch_size) const +{ + const int InputPadding = mNetOffset + outer_padding; + const auto input_block_width = crop_w + InputPadding * 2; + const auto input_block_height = crop_h + InputPadding * 2; + + const int input_block_plane_size = input_block_width * input_block_height * mInputPlane; + + return input_block_plane_size * sizeof(float); + +} + +int cNet::GetOutputMemorySize(const int crop_w, const int crop_h, const int outer_padding, const int batch_size) const +{ + const int InputPadding = mNetOffset + outer_padding; + const auto input_block_width = crop_w + InputPadding * 2; + const auto input_block_height = crop_h + InputPadding * 2; + + const auto output_block_width = input_block_width * mInnerScale - mNetOffset * 2; + const auto output_block_height = input_block_height * mInnerScale - mNetOffset * 2; + + const int output_block_plane_size = output_block_width * output_block_height * mInputPlane; + + return output_block_plane_size * sizeof(float); +} + +// ネットワークを使って画像を再構築する +Waifu2x::eWaifu2xError cNet::ReconstructImage(const bool UseTTA, const int crop_w, const int crop_h, const int outer_padding, const int batch_size, float *inputBlockBuf, float *outputBlockBuf, const cv::Mat &inMat, cv::Mat &outMat) +{ + const auto InputHeight = inMat.size().height; + const auto InputWidth = inMat.size().width; + const auto InputLine = inMat.step1(); + + assert(inMat.channels() == 1 || inMat.channels() == 3); + + const int InputPadding = mNetOffset + outer_padding; // 入力パディング + + const auto NoPaddingInputWidth = InputWidth - InputPadding * 2; // パディングを除いた入力画像サイズ(横) + const auto NoPaddingInputHeight = InputHeight - InputPadding * 2; // パディングを除いた入力画像サイズ(縦) + + cv::Mat outim(NoPaddingInputHeight * mInnerScale, NoPaddingInputWidth * mInnerScale, inMat.type()); + + // float *imptr = (float *)im.data; + float *imptr = (float *)outim.data; + + const auto input_block_width = crop_w + InputPadding * 2; // 入力ブロックサイズ(横) + const auto input_block_height = crop_h + InputPadding * 2; // 入力ブロックサイズ(縦) + + const auto output_block_width = input_block_width * mInnerScale - mNetOffset * 2; // 出力ブロックサイズ(横) + const auto output_block_height = input_block_height * mInnerScale - mNetOffset * 2; // 出力ブロックサイズ(縦) + + const auto output_crop_block_width = crop_w * mInnerScale; // クロップ後の出力ブロックサイズ(横) + const auto output_crop_block_height = crop_h * mInnerScale; // クロップ後の出力ブロックサイズ(縦) + + const auto output_crop_w = (output_block_width - crop_w * mInnerScale) / 2; // 出力後のクロップサイズ + const auto output_crop_h = (output_block_height - crop_h * mInnerScale) / 2; // 出力後のクロップサイズ + + assert(NoPaddingInputWidth % crop_w == 0); + assert(NoPaddingInputHeight % crop_h == 0); + + try + { + auto input_blobs = mNet->input_blobs(); + + assert(input_blobs.size() > 0); + + auto input_blob = mNet->input_blobs()[0]; + + input_blob->Reshape(batch_size, mInputPlane, input_block_height, input_block_width); + + assert(inMat.channels() == mInputPlane); + assert(input_blob->shape(1) == mInputPlane); + + const int WidthNum = NoPaddingInputWidth / crop_w; + const int HeightNum = NoPaddingInputHeight / crop_h; + + const int BlockNum = WidthNum * HeightNum; + + const int input_block_plane_size = input_block_width * input_block_height * mInputPlane; + const int output_block_plane_size = output_block_width * output_block_height * mInputPlane; + + // 画像は(消費メモリの都合上)block_size*block_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, mInputPlane, input_block_height, input_block_width); + + for (int n = 0; n < processNum; n++) + { + const int wn = (num + n) % WidthNum; + const int hn = (num + n) / WidthNum; + + const int w = wn * crop_w; + const int h = hn * crop_h; + + assert(w + input_block_width <= InputWidth && h + input_block_height <= InputHeight); + + cv::Mat someimg = inMat(cv::Rect(w, h, input_block_width, input_block_height)); + + // 画像を直列に変換 + { + float *fptr = inputBlockBuf + (input_block_plane_size * n); + const float *uptr = (const float *)someimg.data; + + const auto Line = someimg.step1(); + + if (someimg.channels() == 1) + { + if (input_block_width == Line) + memcpy(fptr, uptr, input_block_width * input_block_height * sizeof(float)); + else + { + for (int i = 0; i < input_block_height; i++) + memcpy(fptr + i * input_block_width, uptr + i * Line, input_block_width * sizeof(float)); + } + } + else + { + const auto LinePixel = someimg.step1() / someimg.channels(); + const auto Channel = someimg.channels(); + const auto Width = someimg.size().width; + const auto Height = someimg.size().height; + + for (int i = 0; i < Height; i++) + { + for (int j = 0; j < Width; j++) + { + for (int ch = 0; ch < Channel; ch++) + { + const size_t IndexSrc = i * someimg.step1() + j * Channel + ch; + const size_t IndexDst = (ch * Height + i) * Width + j; + fptr[IndexDst] = uptr[IndexSrc]; + } + } + } + } + } + } + + assert(input_blob->count() == input_block_plane_size * processNum); + + // ネットワークに画像を入力 + input_blob->set_cpu_data(inputBlockBuf); + + // 計算 + auto out = mNet->Forward(); + + 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, outputBlockBuf); + + for (int n = 0; n < processNum; n++) + { + const int wn = (num + n) % WidthNum; + const int hn = (num + n) / WidthNum; + + const int w = wn * output_crop_block_width; + const int h = hn * output_crop_block_height; + + const float *fptr = outputBlockBuf + (output_block_plane_size * n); + + // 結果を出力画像にコピー + if (outim.channels() == 1) + { + for (int i = 0; i < output_crop_block_height; i++) + memcpy(imptr + (h + i) * InputLine + w, fptr + (i + output_crop_h) * output_block_width + output_crop_w, output_crop_block_width * sizeof(float)); + } + else + { + const auto LinePixel = outim.step1() / outim.channels(); + const auto Channel = outim.channels(); + + //for (int i = 0; i < output_no_padding_block_height; i++) + //{ + // for (int j = 0; j < output_no_padding_block_width; j++) + // { + // for (int ch = 0; ch < Channel; ch++) + // imptr[((h + i) * LinePixel + (w + j)) * Channel + ch] + // = fptr[(ch * output_block_height + i + output_crop_h) * output_block_width + j + output_padding]; + // } + //} + + for (int i = 0; i < output_crop_block_height; i++) + { + for (int j = 0; j < output_crop_block_width; j++) + { + for (int ch = 0; ch < Channel; ch++) + { + const size_t IndexSrc = (ch * output_block_height + i + output_crop_h) * output_block_width + j + output_crop_w; + const size_t IndexDst = ((h + i) * LinePixel + (w + j)) * Channel + ch; + + imptr[IndexDst] = fptr[IndexSrc]; + } + } + } + } + + //{ + // cv::Mat testim(output_block_size, output_block_size, CV_32FC1); + // float *p = (float *)testim.data; + // for (int i = 0; i < output_block_size; i++) + // { + // for (int j = 0; j < output_block_size; j++) + // { + // p[testim.step1() * i + j] = fptr[i * output_block_size + j]; + // } + // } + + // const int cv_depth = DepthBitToCVDepth(8); + // const double max_val = GetValumeMaxFromCVDepth(cv_depth); + // const double eps = GetEPS(cv_depth); + + // cv::Mat write_iamge; + // testim.convertTo(write_iamge, cv_depth, max_val, eps); + + // cv::imwrite("ti.png", write_iamge); + // testim.release(); + //} + } + } + } + catch (...) + { + return Waifu2x::eWaifu2xError_FailedProcessCaffe; + } + + // 値を0〜1にクリッピング + cv::threshold(outim, outim, 1.0, 1.0, cv::THRESH_TRUNC); + cv::threshold(outim, outim, 0.0, 0.0, cv::THRESH_TOZERO); + + outMat = outim; + + return Waifu2x::eWaifu2xError_OK; +} diff --git a/common/cNet.h b/common/cNet.h new file mode 100644 index 0000000..4963acd --- /dev/null +++ b/common/cNet.h @@ -0,0 +1,37 @@ +#pragma once + +#include "waifu2x.h" + + +class cNet +{ +private: + boost::shared_ptr> mNet; + + int mModelScale; // モデルが対象とする拡大率 + int mInnerScale; // ネット内部で拡大される倍率 + int mNetOffset; // ネットに入力するとどれくらい削れるか + int mInputPlane; // ネットへの入力チャンネル数 + +private: + Waifu2x::eWaifu2xError LoadParameterFromJson(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); + Waifu2x::eWaifu2xError LoadInfoFromJson(const boost::filesystem::path &info_path); + Waifu2x::eWaifu2xError SetParameter(caffe::NetParameter ¶m, const std::string &process) const; + +public: + cNet(); + ~cNet(); + + Waifu2x::eWaifu2xError ConstractNet(const boost::filesystem::path &model_path, const boost::filesystem::path ¶m_path, const boost::filesystem::path &info_path, const std::string &process); + + int GetInputPlane() const; + int GetInnerScale() const; + int GetNetOffset() const; + int GetScale() const; + + int GetInputMemorySize(const int crop_w, const int crop_h, const int outer_padding, const int batch_size) const; + int GetOutputMemorySize(const int crop_w, const int crop_h, const int outer_padding, const int batch_size) const; + + Waifu2x::eWaifu2xError ReconstructImage(const bool UseTTA, const int crop_w, const int crop_h, const int outer_padding, const int batch_size, float *inputBlockBuf, float *outputBlockBuf, const cv::Mat &inMat, cv::Mat &outMat); +}; diff --git a/common/stImage.cpp b/common/stImage.cpp new file mode 100644 index 0000000..93a9998 --- /dev/null +++ b/common/stImage.cpp @@ -0,0 +1,873 @@ +#include "stImage.h" +#include +#include +#include +#include +#include +#include + +#define STB_IMAGE_IMPLEMENTATION +#include +#define STB_IMAGE_WRITE_IMPLEMENTATION +#include + +const int YToRGBConvertMode = CV_GRAY2RGB; +const int YToRGBConverInversetMode = CV_RGB2GRAY; +const int BGRToYConvertMode = CV_BGR2YUV; +const int BGRToConvertInverseMode = CV_YUV2BGR; + +// floatな画像をuint8_tな画像に変換する際の四捨五入に使う値 +// 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 double clip_eps16 = (1.0 / 65535.0) * 0.5 - (1.0e-7 * (1.0 / 65535.0) * 0.5); +const double clip_eps32 = 1.0 * 0.5 - (1.0e-7 * 0.5); + +const std::vector stImage::OutputExtentionList = +{ + {L".png",{8, 16}, boost::optional(), boost::optional(), boost::optional(), boost::optional()}, + {L".bmp",{8}, boost::optional(), boost::optional(), boost::optional(), boost::optional()}, + {L".jpg",{8}, 0, 100, 95, cv::IMWRITE_JPEG_QUALITY}, + {L".jp2",{8, 16}, boost::optional(), boost::optional(), boost::optional(), boost::optional()}, + {L".sr",{8}, boost::optional(), boost::optional(), boost::optional(), boost::optional()}, + {L".tif",{8, 16, 32}, boost::optional(), boost::optional(), boost::optional(), boost::optional()}, + {L".hdr",{8, 16, 32}, boost::optional(), boost::optional(), boost::optional(), boost::optional()}, + {L".exr",{8, 16, 32}, boost::optional(), boost::optional(), boost::optional(), boost::optional()}, + {L".ppm",{8, 16}, boost::optional(), boost::optional(), boost::optional(), boost::optional()}, + {L".webp",{8}, 1, 100, 100, cv::IMWRITE_WEBP_QUALITY}, + {L".tga",{8}, 0, 1, 0, 0}, +}; + + +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; + + try + { + is.open(path, std::ios_base::in | std::ios_base::binary); + } + catch (...) + { + return false; + } + + return readFile(is, buf); +} + +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; + + try + { + os.open(path, std::ios_base::out | std::ios_base::binary | std::ios_base::trunc); + } + catch (...) + { + return false; + } + + return writeFile(os, buf); +} + +static void Waifu2x_stbi_write_func(void *context, void *data, int size) +{ + boost::iostreams::stream *osp = (boost::iostreams::stream *)context; + osp->write((const char *)data, size); +} + +int stImage::DepthBitToCVDepth(const int depth_bit) +{ + switch (depth_bit) + { + case 8: + return CV_8U; + + case 16: + return CV_16U; + + case 32: + return CV_32F; + } + + // 不明だけどとりあえずCV_8Uを返しておく + return CV_8U; +} + +double stImage::GetValumeMaxFromCVDepth(const int cv_depth) +{ + switch (cv_depth) + { + case CV_8U: + return 255.0; + + case CV_16U: + return 65535.0; + + case CV_32F: + return 1.0; + } + + // 不明だけどとりあえず255.0を返しておく + return 255.0; +} + +double stImage::GetEPS(const int cv_depth) +{ + switch (cv_depth) + { + case CV_8U: + return clip_eps8; + + case CV_16U: + return clip_eps16; + + case CV_32F: + return clip_eps32; + } + + // 不明だけどとりあえずclip_eps8返しておく + return clip_eps8; +} + + +Waifu2x::eWaifu2xError stImage::AlphaMakeBorder(std::vector &planes, const cv::Mat &alpha, const int offset) +{ + // このカーネルと画像の畳込みを行うと、(x, y)を中心とした3×3領域の合計値が求まる + const static cv::Mat sum2d_kernel = (cv::Mat_(3, 3) << + 1., 1., 1., + 1., 1., 1., + 1., 1., 1.); + + cv::Mat mask; + cv::threshold(alpha, mask, 0.0, 1.0, cv::THRESH_BINARY); // アルファチャンネルを二値化してマスクとして扱う + + cv::Mat mask_nega; + cv::threshold(mask, mask_nega, 0.0, 1.0, cv::THRESH_BINARY_INV); // 反転したマスク（値が1の箇所は完全透明でない有効な画素となる） + + for (auto &p : planes) // 完全に透明なピクセルにあるゴミを取る + { + p = p.mul(mask); + } + + for (int i = 0; i < offset; i++) + { + cv::Mat mask_weight; + cv::filter2D(mask, mask_weight, -1, sum2d_kernel, cv::Point(-1, -1), 0, cv::BORDER_DEFAULT); // マスクの3×3領域の合計値を求める + + cv::Mat mask_nega_u8; + mask_nega.convertTo(mask_nega_u8, CV_8U, 255.0, clip_eps8); // mask_negaのCV_U8版（OpenCVのAPI上必要になる） + + for (auto &p : planes) // 1チャンネルずつ処理 + { + // チャンネルの3×3領域内の有効画素の平均値を求める + cv::Mat border; + cv::filter2D(p, border, -1, sum2d_kernel, cv::Point(-1, -1), 0, cv::BORDER_DEFAULT); + border /= mask_weight; + + // チャンネルの有効な画素の部分に、計算した平均値をコピー + border.copyTo(p, mask_nega_u8); + } + + // マスクを1回膨張させたものを新しいマスクとする(マスクの3×3領域の合計値を求めたものの非0領域は、マスクを1回膨張させたものの領域に等しい) + cv::threshold(mask_weight, mask, 0.0, 1.0, cv::THRESH_BINARY); + // 新しいマスクの反転したマスクを計算 + cv::threshold(mask, mask_nega, 0.0, 1.0, cv::THRESH_BINARY_INV); + } + + // 画素を0から1にクリッピング + for (auto &p : planes) + { + cv::threshold(p, p, 1.0, 1.0, cv::THRESH_TRUNC); + cv::threshold(p, p, 0.0, 0.0, cv::THRESH_TOZERO); + } + + return Waifu2x::eWaifu2xError_OK; +} + +// 画像を読み込んで値を0.0f〜1.0fの範囲に変換 +Waifu2x::eWaifu2xError stImage::LoadMat(cv::Mat &im, const boost::filesystem::path &input_file) +{ + cv::Mat original_image; + + { + 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 Waifu2x::eWaifu2xError ret = LoadMatBySTBI(original_image, img_data); + if (ret != Waifu2x::eWaifu2xError_OK) + return ret; + } + } + + im = original_image; + + return Waifu2x::eWaifu2xError_OK; +} + +Waifu2x::eWaifu2xError stImage::LoadMatBySTBI(cv::Mat &im, const std::vector &img_data) +{ + int x, y, comp; + stbi_uc *data = stbi_load_from_memory((const stbi_uc *)img_data.data(), img_data.size(), &x, &y, &comp, 0); + if (!data) + return Waifu2x::eWaifu2xError_FailedOpenInputFile; + + int type = 0; + switch (comp) + { + case 1: + case 3: + case 4: + type = CV_MAKETYPE(CV_8U, comp); + break; + + default: + return Waifu2x::eWaifu2xError_FailedOpenInputFile; + } + + im = cv::Mat(cv::Size(x, y), type); + + const auto LinePixel = im.step1() / im.channels(); + const auto Channel = im.channels(); + const auto Width = im.size().width; + const auto Height = im.size().height; + + assert(x == Width); + assert(y == Height); + assert(Channel == comp); + + auto ptr = im.data; + for (int i = 0; i < y; i++) + { + for (int j = 0; j < x; j++) + { + for (int ch = 0; ch < Channel; ch++) + ptr[(i * LinePixel + j) * comp + ch] = data[(i * x + j) * comp + ch]; + } + } + + stbi_image_free(data); + + if (comp >= 3) + { + // RGBだからBGRに変換 + for (int i = 0; i < y; i++) + { + for (int j = 0; j < x; j++) + std::swap(ptr[(i * LinePixel + j) * comp + 0], ptr[(i * LinePixel + j) * comp + 2]); + } + } + + return Waifu2x::eWaifu2xError_OK; +} + +cv::Mat stImage::ConvertToFloat(const cv::Mat &im) +{ + cv::Mat convert; + switch (im.depth()) + { + case CV_8U: + im.convertTo(convert, CV_32F, 1.0 / GetValumeMaxFromCVDepth(CV_8U)); + break; + + case CV_16U: + im.convertTo(convert, CV_32F, 1.0 / GetValumeMaxFromCVDepth(CV_16U)); + break; + + case CV_32F: + convert = im; // 元から0.0〜1.0のはずなので変換は必要ない + break; + } + + return convert; +} + + +stImage::stImage() : mIsRequestDenoise(false), pad_w1(0), pad_h1(0), pad_w2(0), pad_h2(0) +{ +} + +stImage::~stImage() +{ +} + +void stImage::Clear() +{ + mOrgFloatImage.release(); + mTmpImageRGB.release(); + mTmpImageA.release(); + mEndImage.release(); +} + +Waifu2x::eWaifu2xError stImage::Load(const boost::filesystem::path &input_file) +{ + Clear(); + + Waifu2x::eWaifu2xError ret; + + cv::Mat im; + ret = LoadMat(im, input_file); + if (ret != Waifu2x::eWaifu2xError_OK) + return ret; + + mOrgFloatImage = im; + mOrgChannel = im.channels(); + mOrgSize = im.size(); + + 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"); + + mIsRequestDenoise = isJpeg; + + return Waifu2x::eWaifu2xError_OK; +} + +Waifu2x::eWaifu2xError stImage::Load(const void* source, const int width, const int height, const int channel, const int stride) +{ + cv::Mat original_image(cv::Size(width, height), CV_MAKETYPE(CV_8U, channel), (void *)source, stride); + + if (original_image.channels() >= 3) // RGBなのでBGRにする + { + std::vector planes; + cv::split(original_image, planes); + + std::swap(planes[0], planes[2]); + + cv::merge(planes, original_image); + } + + mOrgFloatImage = original_image; + + return Waifu2x::eWaifu2xError_OK; +} + +bool stImage::RequestDenoise() const +{ + return mIsRequestDenoise; +} + +void stImage::Preprocess(const int input_plane, const int net_offset) +{ + mOrgFloatImage = ConvertToFloat(mOrgFloatImage); + + ConvertToNetFormat(input_plane, net_offset); +} + +void stImage::ConvertToNetFormat(const int input_plane, const int alpha_offset) +{ + if (input_plane == 1) // Yモデル + { + if (mOrgFloatImage.channels() == 1) // 1chだけなのでそのまま + mTmpImageRGB = mOrgFloatImage; + else // BGRなので変換 + { + mTmpImageRGB = mOrgFloatImage; + + if (mTmpImageRGB.channels() == 4) // BGRAなのでAだけ取り出す + { + std::vector planes; + cv::split(mTmpImageRGB, planes); + + mTmpImageA = planes[3]; + planes.resize(3); + + AlphaMakeBorder(planes, mTmpImageA, alpha_offset); // 透明なピクセルと不透明なピクセルの境界部分の色を広げる + + cv::merge(planes, mTmpImageRGB); + } + + CreateBrightnessImage(mTmpImageRGB, mTmpImageRGB); + } + } + else // RGBモデル + { + if (mOrgFloatImage.channels() == 1) // 1chだけなのでRGBに変換 + { + cv::cvtColor(mOrgFloatImage, mTmpImageRGB, YToRGBConvertMode); + mOrgFloatImage.release(); + } + else // BGRからRGBに変換(AがあったらAも取り出す) + { + std::vector planes; + cv::split(mOrgFloatImage, planes); + mOrgFloatImage.release(); + + if (planes.size() == 4) // BGRAなのでAだけ取り出す + { + mTmpImageA = planes[3]; + planes.resize(3); + + AlphaMakeBorder(planes, mTmpImageA, alpha_offset); // 透明なピクセルと不透明なピクセルの境界部分の色を広げる + + // α拡大用にRGBに変換 + cv::cvtColor(mTmpImageA, mTmpImageA, CV_GRAY2RGB); + } + + // BGRからRGBにする + std::swap(planes[0], planes[2]); + + cv::merge(planes, mTmpImageRGB); + } + + mOrgFloatImage.release(); + } +} + +// 画像から輝度の画像を取り出す +Waifu2x::eWaifu2xError stImage::CreateBrightnessImage(const cv::Mat &float_image, cv::Mat &im) +{ + if (float_image.channels() > 1) + { + cv::Mat converted_color; + cv::cvtColor(float_image, converted_color, BGRToYConvertMode); + + std::vector planes; + cv::split(converted_color, planes); + + im = planes[0]; + planes.clear(); + } + else + im = float_image; + + return Waifu2x::eWaifu2xError_OK; +} + +bool stImage::HasAlpha() const +{ + return !mTmpImageA.empty(); +} + +void stImage::GetScalePaddingedRGB(cv::Mat &im, cv::Size_ &size, const int net_offset, const int outer_padding, + const int crop_w, const int crop_h, const int scale) +{ + GetScalePaddingedImage(mTmpImageRGB, im, size, net_offset, outer_padding, crop_w, crop_h, scale); +} + +void stImage::SetReconstructedRGB(cv::Mat &im, const cv::Size_ &size, const int inner_scale) +{ + SetReconstructedImage(mTmpImageRGB, im, size, inner_scale); +} + +void stImage::GetScalePaddingedA(cv::Mat &im, cv::Size_ &size, const int net_offset, const int outer_padding, + const int crop_w, const int crop_h, const int scale) +{ + GetScalePaddingedImage(mTmpImageA, im, size, net_offset, outer_padding, crop_w, crop_h, scale); +} + +void stImage::SetReconstructedA(cv::Mat &im, const cv::Size_ &size, const int inner_scale) +{ + SetReconstructedImage(mTmpImageA, im, size, inner_scale); +} + +void stImage::GetScalePaddingedImage(cv::Mat &in, cv::Mat &out, cv::Size_ &size, const int net_offset, const int outer_padding, + const int crop_w, const int crop_h, const int scale) +{ + cv::Mat ret; + + if (scale > 1) + { + cv::Size_ zoom_size = in.size(); + zoom_size.width *= scale; + zoom_size.height *= scale; + + cv::resize(in, ret, zoom_size, 0.0, 0.0, cv::INTER_NEAREST); + } + else + ret = in; + + in.release(); + + size = ret.size(); + + PaddingImage(ret, net_offset, outer_padding, crop_w, crop_h, ret); + + out = ret; +} + +// 入力画像の(Photoshopでいう)キャンバスサイズをoutput_sizeの倍数に変更 +// 画像は左上配置、余白はcv::BORDER_REPLICATEで埋める +void stImage::PaddingImage(const cv::Mat &input, const int net_offset, const int outer_padding, + const int crop_w, const int crop_h, cv::Mat &output) +{ + const auto pad_w1 = net_offset + outer_padding; + const auto pad_h1 = net_offset + outer_padding; + const auto pad_w2 = (int)ceil((double)input.size().width / (double)crop_w) * crop_w - input.size().width + net_offset + outer_padding; + const auto pad_h2 = (int)ceil((double)input.size().height / (double)crop_h) * crop_h - input.size().height + net_offset + outer_padding; + + cv::copyMakeBorder(input, output, pad_h1, pad_h2, pad_w1, pad_w2, cv::BORDER_REPLICATE); +} + +// 拡大、パディングされた画像を設定 +void stImage::SetReconstructedImage(cv::Mat &dst, cv::Mat &src, const cv::Size_ &size, const int inner_scale) +{ + const cv::Size_ s(size * inner_scale); + + // ブロックサイズ用のパディングを取り払う(outer_paddingは再構築の過程で取り除かれている) + dst = src(cv::Rect(0, 0, s.width, s.height)); + + src.release(); +} + +void stImage::Postprocess(const int input_plane, const double scale, const int depth) +{ + DeconvertFromNetFormat(input_plane); + ShrinkImage(scale); + + // 値を0〜1にクリッピング + cv::threshold(mEndImage, mEndImage, 1.0, 1.0, cv::THRESH_TRUNC); + cv::threshold(mEndImage, mEndImage, 0.0, 0.0, cv::THRESH_TOZERO); + + mEndImage = DeconvertFromFloat(mEndImage, depth); + + AlphaCleanImage(mEndImage); +} + +void stImage::DeconvertFromNetFormat(const int input_plane) +{ + if (input_plane == 1) // Yモデル + { + if (mOrgChannel == 1) // もともと1chだけなのでそのまま + { + mEndImage = mTmpImageRGB; + mTmpImageRGB.release(); + mOrgFloatImage.release(); + } + else // もともとBGRなので既存アルゴリズムで拡大したUVに拡大したYを合体して戻す + { + std::vector color_planes; + CreateZoomColorImage(mOrgFloatImage, mTmpImageRGB.size(), color_planes); + mOrgFloatImage.release(); + + color_planes[0] = mTmpImageRGB; + mTmpImageRGB.release(); + + cv::Mat converted_image; + cv::merge(color_planes, converted_image); + color_planes.clear(); + + cv::cvtColor(converted_image, mEndImage, BGRToConvertInverseMode); + converted_image.release(); + } + } + else // RGBモデル + { + // ここの地点でmOrgFloatImageは空 + + if (mOrgChannel == 1) // もともと1chだけなので戻す + { + cv::cvtColor(mTmpImageRGB, mEndImage, YToRGBConverInversetMode); + mTmpImageRGB.release(); + } + else // もともとBGRなのでRGBから戻す(AがあったらAも合体して戻す) + { + std::vector planes; + cv::split(mTmpImageRGB, planes); + mTmpImageRGB.release(); + + if (!mTmpImageA.empty()) // Aもあるので合体 + { + // RGBから1chに戻す + cv::cvtColor(mTmpImageA, mTmpImageA, CV_RGB2GRAY); + + planes.push_back(mTmpImageA); + mTmpImageA.release(); + } + + // RGBからBGRにする + std::swap(planes[0], planes[2]); + + cv::merge(planes, mEndImage); + } + } +} + +void stImage::ShrinkImage(const double scale) +{ + // TODO: scale = 1.0 でも悪影響を及ぼさないか調べる + + const int scaleBase = 2; // TODO: モデルの拡大率によって可変できるようにする + + const int scaleNum = ceil(log(scale) / log(scaleBase)); + const double shrinkRatio = scale >= 1.0 ? scale / std::pow(scaleBase, scaleNum) : scale; + + const cv::Size_ ns(mOrgSize.width * scale, mOrgSize.height * scale); + //if (mEndImage.size().width != ns.width || mEndImage.size().height != ns.height) + //{ + // int argo = cv::INTER_CUBIC; + // if (scale < 0.5) + // argo = cv::INTER_AREA; + + // cv::resize(mEndImage, mEndImage, ns, 0.0, 0.0, argo); + //} +} + +cv::Mat stImage::DeconvertFromFloat(const cv::Mat &im, const int depth) +{ + const int cv_depth = DepthBitToCVDepth(depth); + const double max_val = GetValumeMaxFromCVDepth(cv_depth); + const double eps = GetEPS(cv_depth); + + cv::Mat ret; + if (depth == 32) // 出力がfloat形式なら変換しない + ret = im; + else + im.convertTo(ret, cv_depth, max_val, eps); + + return ret; +} + +namespace +{ + template + void AlphaZeroToZero(std::vector &planes) + { + cv::Mat alpha(planes[3]); + + const T *aptr = (const T *)alpha.data; + + T *ptr0 = (T *)planes[0].data; + T *ptr1 = (T *)planes[1].data; + T *ptr2 = (T *)planes[2].data; + + const size_t Line = alpha.step1(); + const size_t Width = alpha.size().width; + const size_t Height = alpha.size().height; + + for (size_t i = 0; i < Height; i++) + { + for (size_t j = 0; j < Width; j++) + { + const size_t pos = Line * i + j; + + if (aptr[pos] == (T)0) + ptr0[pos] = ptr1[pos] = ptr2[pos] = (T)0; + } + } + } +} + +void stImage::AlphaCleanImage(cv::Mat &im) +{ + // 完全透明のピクセルの色を消す(処理の都合上、完全透明のピクセルにも色を付けたから) + // モデルによっては画像全域の完全透明の場所にごく小さい値のアルファが広がることがある。それを消すためにcv_depthへ変換してからこの処理を行うことにした + // (ただしcv_depthが32の場合だと意味は無いが) + // TODO: モデル(例えばPhoto)によっては0しかない画像を変換しても0.000114856390とかになるので、適切な値のクリッピングを行う？ + if (im.channels() > 3) + { + std::vector planes; + cv::split(im, planes); + im.release(); + + const auto depth = planes[0].depth(); + switch (depth) + { + case CV_8U: + AlphaZeroToZero(planes); + break; + + case CV_16U: + AlphaZeroToZero(planes); + break; + + case CV_32F: + AlphaZeroToZero(planes); + break; + + case CV_64F: + AlphaZeroToZero(planes); + break; + } + + cv::merge(planes, im); + } +} + + +// 入力画像をzoom_sizeの大きさにcv::INTER_CUBICで拡大し、色情報のみを残す +Waifu2x::eWaifu2xError stImage::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, BGRToYConvertMode); + zoom_cubic_image.release(); + + cv::split(converted_cubic_image, cubic_planes); + converted_cubic_image.release(); + + // このY成分は使わないので解放 + cubic_planes[0].release(); + + return Waifu2x::eWaifu2xError_OK; +} + +cv::Mat stImage::GetEndImage() const +{ + return mEndImage; +} + +Waifu2x::eWaifu2xError stImage::Save(const boost::filesystem::path &output_file, const boost::optional &output_quality) +{ + return WriteMat(mEndImage, output_file, output_quality); +} + +Waifu2x::eWaifu2xError stImage::WriteMat(const cv::Mat &im, const boost::filesystem::path &output_file, const boost::optional &output_quality) +{ + const boost::filesystem::path ip(output_file); + const std::string ext = ip.extension().string(); + + if (boost::iequals(ext, ".tga")) + { + unsigned char *data = im.data; + + std::vector rgbimg; + if (im.channels() >= 3 || im.step1() != im.size().width * im.channels()) // RGB用バッファにコピー(あるいはパディングをとる) + { + const auto Line = im.step1(); + const auto Channel = im.channels(); + const auto Width = im.size().width; + const auto Height = im.size().height; + + rgbimg.resize(Width * Height * Channel); + + const auto Stride = Width * Channel; + for (int i = 0; i < Height; i++) + memcpy(rgbimg.data() + Stride * i, im.data + Line * i, Stride); + + data = rgbimg.data(); + } + + if (im.channels() >= 3) // BGRをRGBに並び替え + { + const auto Line = im.step1(); + const auto Channel = im.channels(); + const auto Width = im.size().width; + const auto Height = im.size().height; + + auto ptr = rgbimg.data(); + for (int i = 0; i < Height; i++) + { + for (int j = 0; j < Width; j++) + std::swap(ptr[(i * Width + j) * Channel + 0], ptr[(i * Width + j) * Channel + 2]); + } + } + + boost::iostreams::stream os; + + try + { + os.open(output_file, std::ios_base::out | std::ios_base::binary | std::ios_base::trunc); + } + catch (...) + { + return Waifu2x::eWaifu2xError_FailedOpenOutputFile; + } + + if (!os) + return Waifu2x::eWaifu2xError_FailedOpenOutputFile; + + // RLE圧縮の設定 + bool isSet = false; + const auto &OutputExtentionList = stImage::OutputExtentionList; + for (const auto &elm : OutputExtentionList) + { + if (elm.ext == L".tga") + { + if (elm.imageQualitySettingVolume && output_quality) + { + stbi_write_tga_with_rle = *output_quality; + isSet = true; + } + + break; + } + } + + // 設定されなかったのでデフォルトにする + if (!isSet) + stbi_write_tga_with_rle = 1; + + if (!stbi_write_tga_to_func(Waifu2x_stbi_write_func, &os, im.size().width, im.size().height, im.channels(), data)) + return Waifu2x::eWaifu2xError_FailedOpenOutputFile; + + return Waifu2x::eWaifu2xError_OK; + } + + try + { + const boost::filesystem::path op(output_file); + const boost::filesystem::path opext(op.extension()); + + std::vector params; + + const auto &OutputExtentionList = stImage::OutputExtentionList; + for (const auto &elm : OutputExtentionList) + { + if (elm.ext == opext) + { + if (elm.imageQualitySettingVolume && output_quality) + { + params.push_back(*elm.imageQualitySettingVolume); + params.push_back(*output_quality); + } + + break; + } + } + + std::vector buf; + cv::imencode(ext, im, buf, params); + + if (writeFile(output_file, buf)) + return Waifu2x::eWaifu2xError_OK; + + } + catch (...) + { + } + + return Waifu2x::eWaifu2xError_FailedOpenOutputFile; +} diff --git a/common/stImage.h b/common/stImage.h new file mode 100644 index 0000000..e92caab --- /dev/null +++ b/common/stImage.h @@ -0,0 +1,125 @@ +#pragma once + +#include "waifu2x.h" + + +class stImage +{ +private: + cv::Mat mOrgFloatImage; + int mOrgChannel; + cv::Size_ mOrgSize; + + bool mIsRequestDenoise; + + cv::Mat mTmpImageRGB; // RGB(あるいはY) + cv::Mat mTmpImageA; // αチャンネル + + cv::Mat mEndImage; // 完成した画像 + + int pad_w1; + int pad_h1; + int pad_w2; + int pad_h2; + +public: + struct stOutputExtentionElement + { + std::wstring ext; + std::vector depthList; + boost::optional imageQualityStart; + boost::optional imageQualityEnd; + boost::optional imageQualityDefault; + boost::optional imageQualitySettingVolume; + }; + + const static std::vector OutputExtentionList; + +private: + static Waifu2x::eWaifu2xError LoadMatBySTBI(cv::Mat &im, const std::vector &img_data); + + static cv::Mat ConvertToFloat(const cv::Mat &im); + + static Waifu2x::eWaifu2xError AlphaMakeBorder(std::vector &planes, const cv::Mat &alpha, const int offset); + + static cv::Mat DeconvertFromFloat(const cv::Mat &im, const int depth); + static void AlphaCleanImage(cv::Mat &im); + + static Waifu2x::eWaifu2xError WriteMat(const cv::Mat &im, const boost::filesystem::path &output_file, const boost::optional &output_quality); + + void ConvertToNetFormat(const int input_plane, const int alpha_offset); + + Waifu2x::eWaifu2xError CreateBrightnessImage(const cv::Mat &float_image, cv::Mat &im); + void PaddingImage(const cv::Mat &input, const int net_offset, const int outer_padding, + const int crop_w, const int crop_h, cv::Mat &output); + Waifu2x::eWaifu2xError CreateZoomColorImage(const cv::Mat &float_image, const cv::Size_ &zoom_size, std::vector &cubic_planes); + + // 拡大、パディングされた画像を取得 + // この関数を呼び出した後、mTmpImageRGBは空になるので注意 + void GetScalePaddingedImage(cv::Mat &in, cv::Mat &out, cv::Size_ &size, const int net_offset, const int outer_padding, + const int crop_w, const int crop_h, const int scale); + + // 変換された画像を設定 + // size: GetScalePaddingedImage()で取得したsize + void SetReconstructedImage(cv::Mat &dst, cv::Mat &src, const cv::Size_ &size, const int inner_scale); + + void DeconvertFromNetFormat(const int input_plane); + void ShrinkImage(const double scale); + + static int DepthBitToCVDepth(const int depth_bit); + static double GetValumeMaxFromCVDepth(const int cv_depth); + static double GetEPS(const int cv_depth); + +public: + stImage(); + ~stImage(); + + void Clear(); + + static Waifu2x::eWaifu2xError LoadMat(cv::Mat &im, const boost::filesystem::path &input_file); + + Waifu2x::eWaifu2xError Load(const boost::filesystem::path &input_file); + + // source: (4チャンネルの場合は)RGBAな画素配列 + // dest: (4チャンネルの場合は)処理したRGBAな画素配列 + // width: widthの縦幅 + // height: heightの横幅 + // channel: sourceのチャンネル数 + // stride: sourceのストライド(バイト単位) + // sourceはPostprocess()が終わるまで存在している必要がある + Waifu2x::eWaifu2xError Load(const void* source, const int width, const int height, const int channel, const int stride); + + bool RequestDenoise() const; + + // 前処理 + // RGBモデルの場合はこれが終わった時にはmOrgFloatImageが空になっているので注意 + void Preprocess(const int input_plane, const int net_offset); + + bool HasAlpha() const; + + // 拡大、パディングされた画像を取得 + // この関数を呼び出した後、mTmpImageRGBは空になるので注意 + void GetScalePaddingedRGB(cv::Mat &im, cv::Size_ &size, const int net_offset, const int outer_padding, + const int crop_w, const int crop_h, const int scale); + + // 変換された画像を設定 + // この関数を呼び出した後、imは空になるので注意 + // size: GetScalePaddingedImage()で取得したsize + void SetReconstructedRGB(cv::Mat &im, const cv::Size_ &size, const int inner_scale); + + // 拡大、パディングされた画像を取得 + // この関数を呼び出した後、mTmpImageAは空になるので注意 + void GetScalePaddingedA(cv::Mat &im, cv::Size_ &size, const int net_offset, const int outer_padding, + const int crop_w, const int crop_h, const int scale); + + // 拡大された画像を設定 + // この関数を呼び出した後、imは空になるので注意 + // size: GetScalePaddingedImage()で取得したsize + void SetReconstructedA(cv::Mat &im, const cv::Size_ &size, const int inner_scale); + + void Postprocess(const int input_plane, const double scale, const int depth); + + cv::Mat GetEndImage() const; + + Waifu2x::eWaifu2xError Save(const boost::filesystem::path &output_file, const boost::optional &output_quality); +}; diff --git a/common/waifu2x.cpp b/common/waifu2x.cpp index b455114..f778e5e 100644 --- a/common/waifu2x.cpp +++ b/common/waifu2x.cpp @@ -1,11 +1,10 @@ #include "waifu2x.h" +#include "stImage.h" +#include "cNet.h" #include #include #include #include -#include -#include -#include #include #include #include @@ -14,23 +13,16 @@ #include #include -#include -#include -#include + #include #include #ifdef _MSC_VER #include #endif -#define STB_IMAGE_IMPLEMENTATION -#include -#define STB_IMAGE_WRITE_IMPLEMENTATION -#include - -#if defined(WIN32) || defined(WIN64) -#include -#endif +//#if defined(WIN32) || defined(WIN64) +//#include +//#endif #define CV_VERSION_STR CVAUX_STR(CV_MAJOR_VERSION) CVAUX_STR(CV_MINOR_VERSION) CVAUX_STR(CV_SUBMINOR_VERSION) @@ -90,41 +82,20 @@ #endif #endif -// 入力画像のオフセット -const int offset = 0; +const int ScaleBase = 2; // TODO: モデルの拡大率によって可変できるようにする -const int ConvertMode = CV_RGB2YUV; -const int ConvertInverseMode = CV_YUV2RGB; +// 入力画像に追加するパディング +const int OuterPadding = 0; // 最低限必要なCUDAドライバーのバージョン const int MinCudaDriverVersion = 7050; -// floatな画像をuint8_tな画像に変換する際の四捨五入に使う値 -// 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 double clip_eps16 = (1.0 / 65535.0) * 0.5 - (1.0e-7 * (1.0 / 65535.0) * 0.5); -const double clip_eps32 = 1.0 * 0.5 - (1.0e-7 * 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; -const std::vector Waifu2x::OutputExtentionList = -{ - { L".png", { 8, 16 }, boost::optional(), boost::optional(), boost::optional(), boost::optional() }, - { L".bmp", { 8 }, boost::optional(), boost::optional(), boost::optional(), boost::optional() }, - { L".jpg", { 8 }, 0, 100, 95, cv::IMWRITE_JPEG_QUALITY }, - { L".jp2", { 8, 16 }, boost::optional(), boost::optional(), boost::optional(), boost::optional() }, - { L".sr", { 8 }, boost::optional(), boost::optional(), boost::optional(), boost::optional() }, - { L".tif", { 8, 16, 32 }, boost::optional(), boost::optional(), boost::optional(), boost::optional() }, - { L".hdr", { 8, 16, 32 }, boost::optional(), boost::optional(), boost::optional(), boost::optional() }, - { L".exr", { 8, 16, 32 }, boost::optional(), boost::optional(), boost::optional(), boost::optional() }, - { L".ppm", { 8, 16 }, boost::optional(), boost::optional(), boost::optional(), boost::optional() }, - { L".webp", { 8 }, 1, 100, 100, cv::IMWRITE_WEBP_QUALITY }, - { L".tga", { 8 }, 0, 1, 0, 0 }, -}; +std::string Waifu2x::ExeDir; + #ifndef CUDA_CHECK_WAIFU2X #define CUDA_CHECK_WAIFU2X(condition) \ @@ -172,162 +143,43 @@ namespace IgnoreErrorCV g_IgnoreErrorCV; } -template -static bool writeFile(boost::iostreams::stream &os, const std::vector &buf) +// CUDAが使えるかチェック +Waifu2x::eWaifu2xCudaError Waifu2x::can_use_CUDA() { - 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; - - try + static eWaifu2xCudaError CudaFlag = eWaifu2xCudaError_NotFind; + std::call_once(waifu2x_cuda_once_flag, [&]() { - os.open(path, std::ios_base::out | std::ios_base::binary | std::ios_base::trunc); - } - catch (...) - { - return false; - } + int driverVersion = 0; + if (cudaDriverGetVersion(&driverVersion) == cudaSuccess) + { + if (driverVersion > 0) + { + int runtimeVersion; + if (cudaRuntimeGetVersion(&runtimeVersion) == cudaSuccess) + { + if (runtimeVersion >= MinCudaDriverVersion && driverVersion >= runtimeVersion) + { + cudaDeviceProp prop; + cudaGetDeviceProperties(&prop, 0); + if (prop.major >= 2) + CudaFlag = eWaifu2xCudaError_OK; + else + CudaFlag = eWaifu2xCudaError_OldDevice; + } + else + CudaFlag = eWaifu2xCudaError_OldVersion; + } + else + CudaFlag = eWaifu2xCudaError_NotFind; + } + else + CudaFlag = eWaifu2xCudaError_NotFind; + } + else + CudaFlag = eWaifu2xCudaError_NotFind; + }); - 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; - - try - { - is.open(path, std::ios_base::in | std::ios_base::binary); - } - catch (...) - { - return false; - } - - return readFile(is, buf); -} - -static Waifu2x::eWaifu2xError readProtoText(const boost::filesystem::path &path, ::google::protobuf::Message* proto) -{ - boost::iostreams::stream is; - - try - { - is.open(path, std::ios_base::in); - } - catch (...) - { - return Waifu2x::eWaifu2xError_FailedOpenModelFile; - } - - 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; - - try - { - is.open(path, std::ios_base::in | std::ios_base::binary); - } - catch (...) - { - return Waifu2x::eWaifu2xError_FailedOpenModelFile; - } - - 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()); - - 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; - - try - { - os.open(path, std::ios_base::out | std::ios_base::binary | std::ios_base::trunc); - } - catch (...) - { - return Waifu2x::eWaifu2xError_FailedOpenModelFile; - } - - 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), model_scale(2), net_offset(0), inner_scale(1) -{ -} - -Waifu2x::~Waifu2x() -{ - destroy(); + return CudaFlag; } // cuDNNが使えるかチェック。現状Windowsのみ @@ -376,907 +228,61 @@ Waifu2x::eWaifu2xcuDNNError Waifu2x::can_use_cuDNN() return cuDNNFlag; } -// CUDAが使えるかチェック -Waifu2x::eWaifu2xCudaError Waifu2x::can_use_CUDA() +void Waifu2x::init_liblary(int argc, char** argv) { - static eWaifu2xCudaError CudaFlag = eWaifu2xCudaError_NotFind; - std::call_once(waifu2x_cuda_once_flag, [&]() + if (argc > 0) + ExeDir = argv[0]; + + std::call_once(waifu2x_once_flag, [argc, argv]() { - int driverVersion = 0; - if (cudaDriverGetVersion(&driverVersion) == cudaSuccess) - { - if (driverVersion > 0) - { - int runtimeVersion; - if (cudaRuntimeGetVersion(&runtimeVersion) == cudaSuccess) - { - if (runtimeVersion >= MinCudaDriverVersion && driverVersion >= runtimeVersion) - { - cudaDeviceProp prop; - cudaGetDeviceProperties(&prop, 0); - if (prop.major >= 2) - CudaFlag = eWaifu2xCudaError_OK; - else - CudaFlag = eWaifu2xCudaError_OldDevice; - } - else - CudaFlag = eWaifu2xCudaError_OldVersion; - } - else - CudaFlag = eWaifu2xCudaError_NotFind; - } - else - CudaFlag = eWaifu2xCudaError_NotFind; - } - else - CudaFlag = eWaifu2xCudaError_NotFind; + assert(argc >= 1); + + int tmpargc = 1; + char* tmpargvv[] = {argv[0]}; + char** tmpargv = tmpargvv; + // glog等の初期化 + caffe::GlobalInit(&tmpargc, &tmpargv); }); - - return CudaFlag; -} - -void Waifu2x::init_liblary() -{ } void Waifu2x::quit_liblary() +{} + + +Waifu2x::Waifu2x() : mIsInited(false), mNoiseLevel(0), mIsCuda(false), mInputBlock(nullptr), mInputBlockSize(0), mOutputBlock(nullptr), mOutputBlockSize(0) +{} + +Waifu2x::~Waifu2x() { + Destroy(); } -cv::Mat Waifu2x::LoadMat(const boost::filesystem::path &path) -{ - cv::Mat mat; - LoadMat(mat, path, 0); - - return mat; -} - -Waifu2x::eWaifu2xError Waifu2x::AlphaMakeBorder(std::vector &planes, const cv::Mat &alpha, const int offset) -{ - // このカーネルと画像の畳込みを行うと、(x, y)を中心とした3×3領域の合計値が求まる - const static cv::Mat sum2d_kernel = (cv::Mat_(3, 3) << - 1., 1., 1., - 1., 1., 1., - 1., 1., 1.); - - cv::Mat mask; - cv::threshold(alpha, mask, 0.0, 1.0, cv::THRESH_BINARY); // アルファチャンネルを二値化してマスクとして扱う - - cv::Mat mask_nega; - cv::threshold(mask, mask_nega, 0.0, 1.0, cv::THRESH_BINARY_INV); // 反転したマスク（値が1の箇所は完全透明でない有効な画素となる） - - for (auto &p : planes) // 完全に透明なピクセルにあるゴミを取る - { - p = p.mul(mask); - } - - for (int i = 0; i < offset; i++) - { - cv::Mat mask_weight; - cv::filter2D(mask, mask_weight, -1, sum2d_kernel, cv::Point(-1, -1), 0, cv::BORDER_DEFAULT); // マスクの3×3領域の合計値を求める - - cv::Mat mask_nega_u8; - mask_nega.convertTo(mask_nega_u8, CV_8U, 255.0, clip_eps8); // mask_negaのCV_U8版（OpenCVのAPI上必要になる） - - for (auto &p : planes) // 1チャンネルずつ処理 - { - // チャンネルの3×3領域内の有効画素の平均値を求める - cv::Mat border; - cv::filter2D(p, border, -1, sum2d_kernel, cv::Point(-1, -1), 0, cv::BORDER_DEFAULT); - border /= mask_weight; - - // チャンネルの有効な画素の部分に、計算した平均値をコピー - border.copyTo(p, mask_nega_u8); - } - - // マスクを1回膨張させたものを新しいマスクとする(マスクの3×3領域の合計値を求めたものの非0領域は、マスクを1回膨張させたものの領域に等しい) - cv::threshold(mask_weight, mask, 0.0, 1.0, cv::THRESH_BINARY); - // 新しいマスクの反転したマスクを計算 - cv::threshold(mask, mask_nega, 0.0, 1.0, cv::THRESH_BINARY_INV); - } - - // 画素を0から1にクリッピング - for (auto &p : planes) - { - cv::threshold(p, p, 1.0, 1.0, cv::THRESH_TRUNC); - cv::threshold(p, p, 0.0, 0.0, cv::THRESH_TOZERO); - } - - return eWaifu2xError_OK; -} - -// 画像を読み込んで値を0.0f〜1.0fの範囲に変換 -Waifu2x::eWaifu2xError Waifu2x::LoadMat(cv::Mat &float_image, const boost::filesystem::path &input_file, const int alpha_offset) -{ - cv::Mat original_image; - - { - 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; - 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, alpha_offset); - - planes.push_back(alpha); - cv::merge(planes, convert); - } - - float_image = convert; - - return eWaifu2xError_OK; -} - -Waifu2x::eWaifu2xError Waifu2x::LoadMatBySTBI(cv::Mat &float_image, const std::vector &img_data) -{ - int x, y, comp; - stbi_uc *data = stbi_load_from_memory((const stbi_uc *)img_data.data(), img_data.size(), &x, &y, &comp, 0); - if (!data) - return eWaifu2xError_FailedOpenInputFile; - - int type = 0; - switch (comp) - { - case 1: - case 3: - case 4: - type = CV_MAKETYPE(CV_8U, comp); - break; - - default: - return eWaifu2xError_FailedOpenInputFile; - } - - float_image = cv::Mat(cv::Size(x, y), type); - - 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; - - assert(x == Width); - assert(y == Height); - assert(Channel == comp); - - auto ptr = float_image.data; - for (int i = 0; i < y; i++) - { - for (int j = 0; j < x; j++) - { - for (int ch = 0; ch < Channel; ch++) - ptr[(i * LinePixel + j) * comp + ch] = data[(i * x + j) * comp + ch]; - } - } - - stbi_image_free(data); - - if (comp >= 3) - { - // RGBだからBGRに変換 - for (int i = 0; i < y; i++) - { - for (int j = 0; j < x; j++) - std::swap(ptr[(i * LinePixel + j) * comp + 0], ptr[(i * LinePixel + j) * comp + 2]); - } - } - - return eWaifu2xError_OK; -} - -// 画像から輝度の画像を取り出す -Waifu2x::eWaifu2xError Waifu2x::CreateBrightnessImage(const cv::Mat &float_image, cv::Mat &im) -{ - if (float_image.channels() > 1) - { - cv::Mat converted_color; - cv::cvtColor(float_image, converted_color, ConvertMode); - - std::vector planes; - cv::split(converted_color, planes); - - im = planes[0]; - planes.clear(); - } - else - im = float_image; - - 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 / block_size) + (input.size().width % block_size == 0 ? 0 : 1); - const auto w_blocks = (int)floor(input.size().height / block_size) + (input.size().height % block_size == 0 ? 0 : 1); - const auto height = offset + h_blocks * block_size + block_size; - const auto width = offset + w_blocks * block_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::resize(input, output, zoom_size, 0.0, 0.0, cv::INTER_NEAREST); - - return PaddingImage(output, 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 boost::filesystem::path &model_path, const boost::filesystem::path ¶m_path, const boost::filesystem::path &info_path, const std::string &process) +Waifu2x::eWaifu2xError Waifu2x::Init(const std::string &mode, const int noise_level, + const boost::filesystem::path &model_dir, const std::string &process) { Waifu2x::eWaifu2xError ret; - ret = LoadInfoFromJson(info_path); - if (ret != eWaifu2xError_OK) - return ret; - - boost::filesystem::path modelbin_path = model_path; - modelbin_path += ".protobin"; - boost::filesystem::path caffemodel_path = param_path; - caffemodel_path += ".caffemodel"; - - caffe::NetParameter param_model; - caffe::NetParameter param_caffemodel; - - const auto retModelBin = readProtoBinary(modelbin_path, ¶m_model); - const auto retParamBin = readProtoBinary(caffemodel_path, ¶m_caffemodel); - - if (retModelBin == eWaifu2xError_OK && retParamBin == eWaifu2xError_OK) - { - ret = SetParameter(param_model, process); - 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_model)); - net->CopyTrainedLayersFrom(param_caffemodel); - - input_plane = param_model.layer(0).input_param().shape().Get(0).dim(1); - } - else - { - const auto ret = LoadParameterFromJson(net, model_path, param_path, modelbin_path, caffemodel_path, process); - if (ret != eWaifu2xError_OK) - return ret; - } - - return eWaifu2xError_OK; -} - -Waifu2x::eWaifu2xError Waifu2x::SetParameter(caffe::NetParameter ¶m, const std::string &process) const -{ - param.mutable_state()->set_phase(caffe::TEST); - - { - auto input_layer = param.mutable_layer(0); - auto mid = input_layer->mutable_input_param()->mutable_shape(); - if (mid->size() != 1 || mid->Mutable(0)->dim_size() != 4) - return eWaifu2xError_FailedParseModelFile; - mid->Mutable(0)->set_dim(0, batch_size); - mid->Mutable(0)->set_dim(2, input_block_size); - mid->Mutable(0)->set_dim(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::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) -{ - Waifu2x::eWaifu2xError ret; - - caffe::NetParameter param; - ret = readProtoText(model_path, ¶m); - if (ret != eWaifu2xError_OK) - return ret; - - ret = writeProtoBinary(param, modelbin_path); - if (ret != eWaifu2xError_OK) - return ret; - - ret = SetParameter(param, process); - if (ret != eWaifu2xError_OK) - return ret; - - net = boost::shared_ptr>(new caffe::Net(param)); - - rapidjson::Document d; - std::vector jsonBuf; + if (mIsInited) + return Waifu2x::eWaifu2xError_OK; try { - boost::iostreams::stream is; - - try - { - is.open(param_path, std::ios_base::in | std::ios_base::binary); - } - catch (...) - { - return Waifu2x::eWaifu2xError_FailedOpenModelFile; - } - - if(!is) - return eWaifu2xError_FailedOpenModelFile; - - const size_t size = is.seekg(0, std::ios::end).tellg(); - is.seekg(0, std::ios::beg); - - jsonBuf.resize(size + 1); - is.read(jsonBuf.data(), jsonBuf.size()); - - jsonBuf[jsonBuf.size() - 1] = '\0'; - - d.Parse(jsonBuf.data()); - } - catch (...) - { - return eWaifu2xError_FailedParseModelFile; - } - - if (d.Size() != 7) - return eWaifu2xError_FailedParseModelFile; - - int inputPlane = 0; - int outputPlane = 0; - try - { - inputPlane = d[0]["nInputPlane"].GetInt(); - outputPlane = d[d.Size() - 1]["nOutputPlane"].GetInt(); - } - catch (...) - { - return eWaifu2xError_FailedParseModelFile; - } - - if (inputPlane == 0 || outputPlane == 0) - return eWaifu2xError_FailedParseModelFile; - - if (inputPlane != outputPlane) - return eWaifu2xError_FailedParseModelFile; - - //if (param.layer_size() < 17) - // return eWaifu2xError_FailedParseModelFile; - - std::vector>> list; - auto &v = net->layers(); - for (auto &l : v) - { - auto lk = l->type(); - auto &bv = l->blobs(); - if (bv.size() > 0) - list.push_back(l); - } - - try - { - std::vector weightList; - std::vector biasList; - - int count = 0; - for (auto it = d.Begin(); it != d.End(); ++it) - { - const auto &weight = (*it)["weight"]; - const auto nInputPlane = (*it)["nInputPlane"].GetInt(); - const auto nOutputPlane = (*it)["nOutputPlane"].GetInt(); - const auto kW = (*it)["kW"].GetInt(); - const auto &bias = (*it)["bias"]; - - auto leyer = list[count]; - - auto &b0 = leyer->blobs()[0]; - auto &b1 = leyer->blobs()[1]; - - float *b0Ptr = nullptr; - float *b1Ptr = nullptr; - - if (caffe::Caffe::mode() == caffe::Caffe::CPU) - { - b0Ptr = b0->mutable_cpu_data(); - b1Ptr = b1->mutable_cpu_data(); - } - else - { - b0Ptr = b0->mutable_gpu_data(); - b1Ptr = b1->mutable_gpu_data(); - } - - const auto WeightSize1 = weight.Size(); - const auto WeightSize2 = weight[0].Size(); - const auto KernelHeight = weight[0][0].Size(); - const auto KernelWidth = weight[0][0][0].Size(); - - if (!(b0->count() == WeightSize1 * WeightSize2 * KernelHeight * KernelWidth)) - return eWaifu2xError_FailedConstructModel; - - if (!(b1->count() == bias.Size())) - return eWaifu2xError_FailedConstructModel; - - weightList.resize(0); - biasList.resize(0); - - size_t weightCount = 0; - for (auto it2 = weight.Begin(); it2 != weight.End(); ++it2) - { - for (auto it3 = (*it2).Begin(); it3 != (*it2).End(); ++it3) - { - for (auto it4 = (*it3).Begin(); it4 != (*it3).End(); ++it4) - { - for (auto it5 = (*it4).Begin(); it5 != (*it4).End(); ++it5) - weightList.push_back((float)it5->GetDouble()); - } - } - } - - caffe::caffe_copy(b0->count(), weightList.data(), b0Ptr); - - for (auto it2 = bias.Begin(); it2 != bias.End(); ++it2) - biasList.push_back((float)it2->GetDouble()); - - caffe::caffe_copy(b1->count(), biasList.data(), b1Ptr); - - count++; - } - - net->ToProto(¶m); - - ret = writeProtoBinary(param, caffemodel_path); - if (ret != eWaifu2xError_OK) - return ret; - } - catch (...) - { - return eWaifu2xError_FailedConstructModel; - } - - input_plane = inputPlane; - - return eWaifu2xError_OK; -} - -Waifu2x::eWaifu2xError Waifu2x::LoadInfoFromJson(const boost::filesystem::path &info_path) -{ - rapidjson::Document d; - std::vector jsonBuf; - - try - { - boost::iostreams::stream is; - - try - { - is.open(info_path, std::ios_base::in | std::ios_base::binary); - } - catch (...) - { - return Waifu2x::eWaifu2xError_FailedOpenModelFile; - } - - if (!is) - return eWaifu2xError_FailedOpenModelFile; - - const size_t size = is.seekg(0, std::ios::end).tellg(); - is.seekg(0, std::ios::beg); - - jsonBuf.resize(size + 1); - is.read(jsonBuf.data(), jsonBuf.size()); - - jsonBuf[jsonBuf.size() - 1] = '\0'; - - d.Parse(jsonBuf.data()); - - const bool resize = d.HasMember("resize") && d["resize"].GetBool() ? true : false; - const auto name = d["name"].GetString(); - const int channels = d["channels"].GetInt(); - net_offset = d["offset"].GetInt(); - inner_scale = d["scale_factor"].GetInt(); - - inner_padding = net_offset; - input_block_size = crop_size + (inner_padding + outer_padding) * 2; - output_block_size = (crop_size + (inner_padding + outer_padding) * 2) * inner_scale - net_offset * 2; - original_width_height = 128 + net_offset * 2; - } - catch (...) - { - return eWaifu2xError_FailedParseModelFile; - } - - return eWaifu2xError_OK; -} - -// ネットワークを使って画像を再構築する -Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr> net, const int reconstructed_scale, cv::Mat &im) -{ - const auto Height = im.size().height; - const auto Width = im.size().width; - const auto Line = im.step1(); - - assert(Width % block_size == 0); - assert(Height % block_size == 0); - - assert(im.channels() == 1 || im.channels() == 3); - - cv::Mat outim(im.rows * reconstructed_scale, im.cols * reconstructed_scale, 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 / block_size; - const int HeightNum = Height / block_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) * reconstructed_scale - net_offset; - const int output_no_padding_block_size = block_size * reconstructed_scale; - - // 画像は(消費メモリの都合上)block_size*block_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 * block_size; - const int h = hn * block_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がnet_offsetで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]; - } - } - } - } - } - } - - 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_no_padding_block_size; - const int h = hn * output_no_padding_block_size; - - const float *fptr = output_block + (output_block_plane_size * n); - - // 結果を出力画像にコピー - if (outim.channels() == 1) - { - for (int i = 0; i < output_no_padding_block_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 < output_no_padding_block_size; i++) - { - for (int j = 0; j < output_no_padding_block_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 testim(output_block_size, output_block_size, CV_32FC1); - // float *p = (float *)testim.data; - // for (int i = 0; i < output_block_size; i++) - // { - // for (int j = 0; j < output_block_size; j++) - // { - // p[testim.step1() * i + j] = fptr[i * output_block_size + j]; - // } - // } - - // const int cv_depth = DepthBitToCVDepth(8); - // const double max_val = GetValumeMaxFromCVDepth(cv_depth); - // const double eps = GetEPS(cv_depth); - - // cv::Mat write_iamge; - // testim.convertTo(write_iamge, cv_depth, max_val, eps); - - // cv::imwrite("ti.png", write_iamge); - // testim.release(); - //} - } - } - } - 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 boost::optional ScaleRatio, const boost::optional ScaleWidth, const boost::optional ScaleHeight, - const boost::filesystem::path &ModelDir, const std::string &Process, - const boost::optional OutputQuality, const int OutputDepth, const bool UseTTA, const int CropSize, const int BatchSize) -{ - Waifu2x::eWaifu2xError ret; - - if (is_inited) - return eWaifu2xError_OK; - - int valid_num = 0; - if (ScaleRatio) - valid_num++; - if (ScaleWidth) - valid_num++; - if (ScaleHeight) - valid_num++; - - if (valid_num != 1) - return eWaifu2xError_InvalidParameter; - - if (ScaleRatio && *ScaleRatio <= 0.0) - return eWaifu2xError_InvalidParameter; - if (ScaleWidth && *ScaleWidth <= 0) - return eWaifu2xError_InvalidParameter; - if (ScaleHeight && *ScaleHeight <= 0.) - return eWaifu2xError_InvalidParameter; - - try - { - mode = Mode; - noise_level = NoiseLevel; - scale_ratio = ScaleRatio; - scale_width = ScaleWidth; - scale_height = ScaleHeight; - model_dir = ModelDir; - process = Process; - use_tta = UseTTA; - - output_quality = OutputQuality; - output_depth = OutputDepth; - - crop_size = CropSize; - batch_size = BatchSize; - - outer_padding = 1; - - block_size = crop_size - offset * 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); - }); - + std::string Process = process; const auto cuDNNCheckStartTime = std::chrono::system_clock::now(); - if (process == "gpu") + if (Process == "gpu") { if (can_use_CUDA() != eWaifu2xCudaError_OK) - return eWaifu2xError_FailedCudaCheck; + return Waifu2x::eWaifu2xError_FailedCudaCheck; // cuDNNが使えそうならcuDNNを使う else if (can_use_cuDNN() == eWaifu2xcuDNNError_OK) - process = "cudnn"; + Process = "cudnn"; } + mMode = mode; + mNoiseLevel = noise_level; + mProcess = Process; + const auto cuDNNCheckEndTime = std::chrono::system_clock::now(); boost::filesystem::path mode_dir_path(model_dir); @@ -1284,28 +290,33 @@ Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &M { // まずはカレントディレクトリ下にあるか探す mode_dir_path = boost::filesystem::absolute(model_dir); - if (!boost::filesystem::exists(mode_dir_path) && argc >= 1) // 無かったらargv[0]から実行ファイルのあるフォルダを推定し、そのフォルダ下にあるか探す + if (!boost::filesystem::exists(mode_dir_path) && !ExeDir.empty()) // 無かったらargv[0]から実行ファイルのあるフォルダを推定し、そのフォルダ下にあるか探す { - boost::filesystem::path a0(argv[0]); + boost::filesystem::path a0(ExeDir); if (a0.is_absolute()) mode_dir_path = a0.branch_path() / model_dir; } } if (!boost::filesystem::exists(mode_dir_path)) - return eWaifu2xError_FailedOpenModelFile; + return Waifu2x::eWaifu2xError_FailedOpenModelFile; - if (process == "cpu") + if (mProcess == "cpu") { caffe::Caffe::set_mode(caffe::Caffe::CPU); - isCuda = false; + mIsCuda = false; } else { caffe::Caffe::set_mode(caffe::Caffe::GPU); - isCuda = true; + mIsCuda = true; } + mInputPlane = 0; + mMaxNetOffset = 0; + + // TODO: ノイズ除去と拡大を同時に行うネットワークへの対処を考える + if (mode == "noise" || mode == "noise_scale" || mode == "auto_scale") { const std::string base_name = "noise" + std::to_string(noise_level) + "_model"; @@ -1314,9 +325,14 @@ Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &M const boost::filesystem::path param_path = mode_dir_path / (base_name + ".json"); const boost::filesystem::path info_path = mode_dir_path / "info.json"; - ret = ConstractNet(net_noise, model_path, param_path, info_path, process); - if (ret != eWaifu2xError_OK) + mNoiseNet.reset(new cNet); + + ret = mNoiseNet->ConstractNet(model_path, param_path, info_path, mProcess); + if (ret != Waifu2x::eWaifu2xError_OK) return ret; + + mInputPlane = mNoiseNet->GetInputPlane(); + mMaxNetOffset = mNoiseNet->GetNetOffset(); } if (mode == "scale" || mode == "noise_scale" || mode == "auto_scale") @@ -1327,305 +343,184 @@ Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &M const boost::filesystem::path param_path = mode_dir_path / (base_name + ".json"); const boost::filesystem::path info_path = mode_dir_path / "info.json"; - ret = ConstractNet(net_scale, model_path, param_path, info_path, process); - if (ret != eWaifu2xError_OK) + mScaleNet.reset(new cNet); + + ret = mScaleNet->ConstractNet(model_path, param_path, info_path, mProcess); + if (ret != Waifu2x::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; + assert(mInputPlane == 0 || mInputPlane == mScaleNet->GetInputPlane()); - 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)); + mInputPlane = mScaleNet->GetInputPlane(); + mMaxNetOffset = std::max(mScaleNet->GetNetOffset(), mMaxNetOffset); } 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; + mIsInited = true; } catch (...) { - return eWaifu2xError_InvalidParameter; + return Waifu2x::eWaifu2xError_InvalidParameter; } - return eWaifu2xError_OK; + return Waifu2x::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; -} - -static void Waifu2x_stbi_write_func(void *context, void *data, int size) -{ - boost::iostreams::stream *osp = (boost::iostreams::stream *)context; - osp->write((const char *)data, size); -} - -Waifu2x::eWaifu2xError Waifu2x::WriteMat(const cv::Mat &im, const boost::filesystem::path &output_file, const boost::optional &output_quality) -{ - const boost::filesystem::path ip(output_file); - const std::string ext = ip.extension().string(); - - const bool isJpeg = boost::iequals(ext, ".jpg") || boost::iequals(ext, ".jpeg"); - - if (boost::iequals(ext, ".tga")) - { - unsigned char *data = im.data; - - std::vector rgbimg; - if (im.channels() >= 3 || im.step1() != im.size().width * im.channels()) // RGB用バッファにコピー(あるいはパディングをとる) - { - const auto Line = im.step1(); - const auto Channel = im.channels(); - const auto Width = im.size().width; - const auto Height = im.size().height; - - rgbimg.resize(Width * Height * Channel); - - const auto Stride = Width * Channel; - for (int i = 0; i < Height; i++) - memcpy(rgbimg.data() + Stride * i, im.data + Line * i, Stride); - - data = rgbimg.data(); - } - - if (im.channels() >= 3) // BGRをRGBに並び替え - { - const auto Line = im.step1(); - const auto Channel = im.channels(); - const auto Width = im.size().width; - const auto Height = im.size().height; - - auto ptr = rgbimg.data(); - for (int i = 0; i < Height; i++) - { - for (int j = 0; j < Width; j++) - std::swap(ptr[(i * Width + j) * Channel + 0], ptr[(i * Width + j) * Channel + 2]); - } - } - - boost::iostreams::stream os; - - try - { - os.open(output_file, std::ios_base::out | std::ios_base::binary | std::ios_base::trunc); - } - catch (...) - { - return Waifu2x::eWaifu2xError_FailedOpenOutputFile; - } - - if(!os) - return eWaifu2xError_FailedOpenOutputFile; - - // RLE圧縮の設定 - bool isSet = false; - const auto &OutputExtentionList = Waifu2x::OutputExtentionList; - for (const auto &elm : OutputExtentionList) - { - if (elm.ext == L".tga") - { - if (elm.imageQualitySettingVolume && output_quality) - { - stbi_write_tga_with_rle = *output_quality; - isSet = true; - } - - break; - } - } - - // 設定されなかったのでデフォルトにする - if (!isSet) - stbi_write_tga_with_rle = 1; - - if (!stbi_write_tga_to_func(Waifu2x_stbi_write_func, &os, im.size().width, im.size().height, im.channels(), data)) - return eWaifu2xError_FailedOpenOutputFile; - - return eWaifu2xError_OK; - } - - try - { - const boost::filesystem::path op(output_file); - const boost::filesystem::path opext(op.extension()); - - std::vector params; - - const auto &OutputExtentionList = Waifu2x::OutputExtentionList; - for (const auto &elm : OutputExtentionList) - { - if (elm.ext == opext) - { - if (elm.imageQualitySettingVolume && output_quality) - { - params.push_back(*elm.imageQualitySettingVolume); - params.push_back(*output_quality); - } - - break; - } - } - - std::vector buf; - cv::imencode(ext, im, buf, params); - - if (writeFile(output_file, buf)) - return eWaifu2xError_OK; - - } - catch (...) - { - } - - return eWaifu2xError_FailedOpenOutputFile; -} - -Waifu2x::eWaifu2xError Waifu2x::BeforeReconstructFloatMatProcess(const cv::Mat &in, cv::Mat &out, bool &convertBGRflag) +Waifu2x::eWaifu2xError Waifu2x::waifu2x(const boost::filesystem::path &input_file, const boost::filesystem::path &output_file, + const double factor, const waifu2xCancelFunc cancel_func, const int crop_w, const int crop_h, + const boost::optional output_quality, const int output_depth, const bool use_tta, + const int batch_size) { Waifu2x::eWaifu2xError ret; - convertBGRflag = false; + if (!mIsInited) + return Waifu2x::eWaifu2xError_NotInitialized; - cv::Mat im; - if (input_plane == 1) - CreateBrightnessImage(in, im); - else - { - im = in; - if (in.channels() == 1) - { - cv::cvtColor(in, im, CV_GRAY2BGR); - convertBGRflag = true; - } + stImage image; + ret = image.Load(input_file); + if (ret != Waifu2x::eWaifu2xError_OK) + return ret; - std::vector planes; - cv::split(im, planes); + image.Preprocess(mInputPlane, mMaxNetOffset); - if (im.channels() == 4) - planes.resize(3); + const bool isReconstructNoise = mMode == "noise" || mMode == "noise_scale" || (mMode == "auto_scale" && image.RequestDenoise()); + const bool isReconstructScale = mMode == "scale" || mMode == "noise_scale" || mMode == "auto_scale"; - // BGRからRGBにする - std::swap(planes[0], planes[2]); + cv::Mat reconstruct_image; + ret = ReconstructImage(factor, crop_w, crop_h, use_tta, batch_size, isReconstructNoise, isReconstructScale, cancel_func, image); + if (ret != Waifu2x::eWaifu2xError_OK) + return ret; - cv::merge(planes, im); - } + image.Postprocess(mInputPlane, factor, output_depth); - out = im; + ret = image.Save(output_file, output_quality); + if (ret != Waifu2x::eWaifu2xError_OK) + return ret; - return eWaifu2xError_OK; + return Waifu2x::eWaifu2xError_OK; } -Waifu2x::eWaifu2xError Waifu2x::ReconstructFloatMat(const bool isReconstructNoise, const bool isReconstructScale, const waifu2xCancelFunc cancel_func, const cv::Mat &in, cv::Mat &out) +Waifu2x::eWaifu2xError Waifu2x::waifu2x(const double factor, const void* source, void* dest, const int width, const int height, + const int in_channel, const int in_stride, const int out_channel, const int out_stride, + const int crop_w, const int crop_h, const bool use_tta, const int batch_size) { Waifu2x::eWaifu2xError ret; - cv::Mat im(in); - cv::Size_ image_size = im.size(); + if (!mIsInited) + return Waifu2x::eWaifu2xError_NotInitialized; + stImage image; + ret = image.Load(source, width, height, in_channel, in_stride); + if (ret != Waifu2x::eWaifu2xError_OK) + return ret; + + image.Preprocess(mInputPlane, mMaxNetOffset); + + const bool isReconstructNoise = mMode == "noise" || mMode == "noise_scale"; + const bool isReconstructScale = mMode == "scale" || mMode == "noise_scale" || mMode == "auto_scale"; + + cv::Mat reconstruct_image; + ret = ReconstructImage(factor, crop_w, crop_h, use_tta, batch_size, isReconstructNoise, isReconstructScale, nullptr, image); + if (ret != Waifu2x::eWaifu2xError_OK) + return ret; + + image.Postprocess(mInputPlane, factor, 8); + + cv::Mat out_image = image.GetEndImage(); + image.Clear(); + + // 出力配列へ書き込み + { + const auto width = out_image.size().width; + const auto stride = out_image.step1(); + for (int i = 0; i < out_image.size().height; i++) + memcpy((uint8_t *)dest + out_stride * i, out_image.data + stride * i, stride); + } + + return Waifu2x::eWaifu2xError_OK; +} + +Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(const double factor, const int crop_w, const int crop_h, const bool use_tta, const int batch_size, + const bool isReconstructNoise, const bool isReconstructScale, const Waifu2x::waifu2xCancelFunc cancel_func, stImage &image) +{ + Waifu2x::eWaifu2xError ret; + + // TODO: ノイズ除去と拡大を同時に行うネットワークへの対処を考える if (isReconstructNoise) { - PaddingImage(im, im); + cv::Mat im; + cv::Size_ size; + image.GetScalePaddingedRGB(im, size, mNoiseNet->GetNetOffset(), OuterPadding, crop_w, crop_h, 1); - ret = ReconstructImage(net_noise, 1, im); - if (ret != eWaifu2xError_OK) + ret = ProcessNet(mNoiseNet, crop_w, crop_h, use_tta, batch_size, im); + if (ret != Waifu2x::eWaifu2xError_OK) return ret; - // パディングを取り払う - im = im(cv::Rect(offset, offset, image_size.width, image_size.height)); - - // 値を0〜1にクリッピング - cv::threshold(im, im, 1.0, 1.0, cv::THRESH_TRUNC); - cv::threshold(im, im, 0.0, 0.0, cv::THRESH_TOZERO); + image.SetReconstructedRGB(im, size, 1); } if (cancel_func && cancel_func()) - return eWaifu2xError_Cancel; + return Waifu2x::eWaifu2xError_Cancel; - image_size = im.size() * inner_scale; - const double ratio = CalcScaleRatio(image_size); - const int scale2 = ceil(log(ratio) / log(model_scale)); + const int scaleNum = ceil(log(factor) / log(ScaleBase)); if (isReconstructScale) { - const int prior_scale = model_scale / inner_scale; - bool isError = false; - for (int i = 0; i < scale2; i++) + for (int i = 0; i < scaleNum; i++) { - if (prior_scale > 1) - { - // TODO: 今は2倍拡大しか考慮していないので、それ以外に対応 - Zoom2xAndPaddingImage(im, im, image_size); - } - else - { - PaddingImage(im, im); - } - - ret = ReconstructImage(net_scale, inner_scale, im); - if (ret != eWaifu2xError_OK) + ret = ReconstructScale(crop_w, crop_h, use_tta, batch_size, cancel_func, image); + if (ret != Waifu2x::eWaifu2xError_OK) return ret; - - // パディングを取り払う - im = im(cv::Rect(offset, offset, image_size.width, image_size.height)); - - // 値を0〜1にクリッピング - cv::threshold(im, im, 1.0, 1.0, cv::THRESH_TRUNC); - cv::threshold(im, im, 0.0, 0.0, cv::THRESH_TOZERO); } } - if (cancel_func && cancel_func()) - return eWaifu2xError_Cancel; - - out = im; - - return eWaifu2xError_OK; + return Waifu2x::eWaifu2xError_OK; } -Waifu2x::eWaifu2xError Waifu2x::Reconstruct(const bool isReconstructNoise, const bool isReconstructScale, const waifu2xCancelFunc cancel_func, const cv::Mat &in, cv::Mat &out) +Waifu2x::eWaifu2xError Waifu2x::ReconstructScale(const int crop_w, const int crop_h, const bool use_tta, const int batch_size, + const Waifu2x::waifu2xCancelFunc cancel_func, stImage &image) { Waifu2x::eWaifu2xError ret; - bool convertBGRflag = false; - cv::Mat brfm; - ret = BeforeReconstructFloatMatProcess(in, brfm, convertBGRflag); - if (ret != eWaifu2xError_OK) + if (image.HasAlpha()) + { + cv::Mat im; + cv::Size_ size; + image.GetScalePaddingedA(im, size, mScaleNet->GetNetOffset(), OuterPadding, crop_w, crop_h, mScaleNet->GetScale() / mScaleNet->GetInnerScale()); + + ret = ReconstructByNet(mScaleNet, crop_w, crop_h, use_tta, batch_size, cancel_func, im); + if (ret != Waifu2x::eWaifu2xError_OK) + return ret; + + image.SetReconstructedA(im, size, mScaleNet->GetInnerScale()); + } + + cv::Mat im; + cv::Size_ size; + image.GetScalePaddingedRGB(im, size, mScaleNet->GetNetOffset(), OuterPadding, crop_w, crop_h, mScaleNet->GetScale() / mScaleNet->GetInnerScale()); + + ret = ReconstructByNet(mScaleNet, crop_w, crop_h, use_tta, batch_size, cancel_func, im); + if (ret != Waifu2x::eWaifu2xError_OK) return ret; - cv::Mat reconstruct_image; + image.SetReconstructedRGB(im, size, mScaleNet->GetInnerScale()); + + return Waifu2x::eWaifu2xError_OK; +} + +Waifu2x::eWaifu2xError Waifu2x::ReconstructByNet(std::shared_ptr net, const int crop_w, const int crop_h, const bool use_tta, const int batch_size, + const Waifu2x::waifu2xCancelFunc cancel_func, cv::Mat &im) +{ + Waifu2x::eWaifu2xError ret; + if (!use_tta) // 普通に処理 { - ret = ReconstructFloatMat(isReconstructNoise, isReconstructScale, cancel_func, brfm, reconstruct_image); - if (ret != eWaifu2xError_OK) + ret = ProcessNet(net, crop_w, crop_h, use_tta, batch_size, im); + if (ret != Waifu2x::eWaifu2xError_OK) return ret; } else // Test-Time Augmentation Mode @@ -1654,10 +549,10 @@ Waifu2x::eWaifu2xError Waifu2x::Reconstruct(const bool isReconstructNoise, const RotateCounterclockwise90(mat); }; - cv::Mat ri[8]; + cv::Mat reconstruct_image; for (int i = 0; i < 8; i++) { - cv::Mat in(brfm.clone()); + cv::Mat in(im); const int rotateNum = i % 4; RotateClockwise90N(in, rotateNum); @@ -1665,8 +560,8 @@ Waifu2x::eWaifu2xError Waifu2x::Reconstruct(const bool isReconstructNoise, const if (i >= 4) cv::flip(in, in, 1); // 垂直軸反転 - ret = ReconstructFloatMat(isReconstructNoise, isReconstructScale, cancel_func, in, in); - if (ret != eWaifu2xError_OK) + ret = ProcessNet(net, crop_w, crop_h, use_tta, batch_size, im); + if (ret != Waifu2x::eWaifu2xError_OK) return ret; if (i >= 4) @@ -1674,567 +569,86 @@ Waifu2x::eWaifu2xError Waifu2x::Reconstruct(const bool isReconstructNoise, const RotateCounterclockwise90N(in, rotateNum); - ri[i] = in; + if (i == 0) + reconstruct_image = in; + else + reconstruct_image += in; } - reconstruct_image = ri[0]; - for (int i = 1; i < 8; i++) - reconstruct_image += ri[i]; - reconstruct_image /= 8.0; } - if (convertBGRflag) - { - cv::cvtColor(reconstruct_image, reconstruct_image, CV_RGB2GRAY); // この地点ではまだRGBなことに注意 - } - - // 値を0〜1にクリッピング - cv::threshold(reconstruct_image, reconstruct_image, 1.0, 1.0, cv::THRESH_TRUNC); - cv::threshold(reconstruct_image, reconstruct_image, 0.0, 0.0, cv::THRESH_TOZERO); - - out = reconstruct_image; - - return eWaifu2xError_OK; + return Waifu2x::eWaifu2xError_OK; } -Waifu2x::eWaifu2xError Waifu2x::AfterReconstructFloatMatProcess(const bool isReconstructScale, const waifu2xCancelFunc cancel_func, const cv::Mat &floatim, cv::Mat &in, cv::Mat &out) +Waifu2x::eWaifu2xError Waifu2x::ProcessNet(std::shared_ptr net, const int crop_w, const int crop_h, const bool use_tta, const int batch_size, cv::Mat &im) { - cv::Size_ image_size = in.size(); + Waifu2x::eWaifu2xError ret; - cv::Mat process_image; - if (input_plane == 1) + const auto InputMemorySize = net->GetInputMemorySize(crop_w, crop_h, OuterPadding, batch_size); + if (InputMemorySize > mInputBlockSize) { - // 再構築した輝度画像とCreateZoomColorImage()で作成した色情報をマージして通常の画像に変換し、書き込む + if (mIsCuda) + CUDA_HOST_SAFE_FREE(mInputBlock); + else + SAFE_DELETE_WAIFU2X(mInputBlock); - std::vector color_planes; - CreateZoomColorImage(floatim, image_size, color_planes); + CUDA_CHECK_WAIFU2X(cudaHostAlloc(&mInputBlock, InputMemorySize, cudaHostAllocWriteCombined)); - color_planes[0] = in; - in.release(); + mInputBlockSize = InputMemorySize; + } - cv::Mat converted_image; - cv::merge(color_planes, converted_image); - color_planes.clear(); + const auto OutputMemorySize = net->GetOutputMemorySize(crop_w, crop_h, OuterPadding, batch_size); + if (OutputMemorySize > mOutputBlockSize) + { + if (mIsCuda) + CUDA_HOST_SAFE_FREE(mOutputBlock); + else + SAFE_DELETE_WAIFU2X(mOutputBlock); - cv::cvtColor(converted_image, process_image, ConvertInverseMode); - converted_image.release(); + CUDA_CHECK_WAIFU2X(cudaHostAlloc(&mOutputBlock, OutputMemorySize, cudaHostAllocDefault)); + + mInputBlockSize = OutputMemorySize; + } + + ret = net->ReconstructImage(use_tta, crop_w, crop_h, OuterPadding, batch_size, mInputBlock, mOutputBlock, im, im); + if (ret != Waifu2x::eWaifu2xError_OK) + return ret; + + return Waifu2x::eWaifu2xError_OK; +} + +//double Waifu2x::CalcScaleRatio(const cv::Size_ &size) const +//{ +// if (scale_ratio) +// return *scale_ratio; +// +// if (scale_width) +// return (double)*scale_width / (double)size.width; +// +// return (double)*scale_height / (double)size.height; +//} + +void Waifu2x::Destroy() +{ + mNoiseNet.reset(); + mScaleNet.reset(); + + if (mIsCuda) + { + CUDA_HOST_SAFE_FREE(mInputBlock); + CUDA_HOST_SAFE_FREE(mOutputBlock); } else { - std::vector planes; - cv::split(in, planes); - in.release(); - - // RGBからBGRに直す - std::swap(planes[0], planes[2]); - - cv::merge(planes, process_image); + SAFE_DELETE_WAIFU2X(mInputBlock); + SAFE_DELETE_WAIFU2X(mOutputBlock); } - const double ratio = CalcScaleRatio(image_size); - const int scale2 = ceil(log2(ratio)); - const double shrinkRatio = ratio >= 1.0 ? ratio / std::pow(2.0, (double)scale2) : ratio; - - if (isReconstructScale) - { - const cv::Size_ ns(image_size.width * shrinkRatio, image_size.height * shrinkRatio); - if (image_size.width != ns.width || image_size.height != ns.height) - { - int argo = cv::INTER_CUBIC; - if (ratio < 0.5) - argo = cv::INTER_AREA; - - cv::resize(process_image, process_image, ns, 0.0, 0.0, argo); - } - } - - cv::Mat alpha; - if (floatim.channels() == 4) - { - std::vector planes; - cv::split(floatim, planes); - - alpha = planes[3]; - planes.clear(); - - if (isReconstructScale) - { - const auto memSize = process_image.step1() * process_image.elemSize1() * process_image.size().height; - - if (memSize < 3ULL * 1000ULL * 1000ULL * 1000ULL) // 拡大後のサイズが3GB超えていたらファイルに書き出してメモリ不足対策 - Reconstruct(false, true, cancel_func, alpha, alpha); - else - { - boost::filesystem::path temp = boost::filesystem::unique_path("%%%%-%%%%-%%%%-%%%%.bin"); - - auto compp = [](const cv::Mat &im, const boost::filesystem::path &temp) - { - static char outbuf[10240000]; - FILE *fout; - z_stream z; - int count; - - if (!(fout = fopen(temp.string().c_str(), "wb"))) - return false; - - z.zalloc = Z_NULL; - z.zfree = Z_NULL; - z.opaque = Z_NULL; - - if (deflateInit(&z, Z_DEFAULT_COMPRESSION) != Z_OK) - { - fclose(fout); - return false; - } - - z.next_in = (z_const Bytef *)im.data; - z.avail_in = im.step1() * im.elemSize1() * im.size().height; - z.next_out = (Bytef *)outbuf; - z.avail_out = sizeof(outbuf); - - for(;;) - { - const int status = deflate(&z, Z_FINISH); - if (status == Z_STREAM_END) - break; // 完了 - - if (status != Z_OK) - { - fclose(fout); - return false; - } - - if (z.avail_out == 0) - { - if (fwrite(outbuf, 1, sizeof(outbuf), fout) != sizeof(outbuf)) - { - fclose(fout); - return false; - } - z.next_out = (Bytef *)outbuf; // 出力バッファ残量を元に戻す - z.avail_out = sizeof(outbuf); // 出力ポインタを元に戻す - } - } - - // 残りを吐き出す - if ((count = sizeof(outbuf) - z.avail_out) != 0) - { - if (fwrite(outbuf, 1, count, fout) != count) - { - fclose(fout); - return false; - } - } - - // 後始末 - if (deflateEnd(&z) != Z_OK) - { - fclose(fout); - return false; - } - - fclose(fout); - - return true; - }; - - auto decompp = [](const cv::Size &size, const int type, cv::Mat &out, const boost::filesystem::path &temp) - { - static char inbuf[102400]; - FILE *fin; - z_stream z; - - if (!(fin = fopen(temp.string().c_str(), "rb"))) - return false; - - z.zalloc = Z_NULL; - z.zfree = Z_NULL; - z.opaque = Z_NULL; - - z.next_in = Z_NULL; - z.avail_in = 0; - if (inflateInit(&z) != Z_OK) - { - fclose(fin); - return false; - } - - out = cv::Mat(size, type); - - const int MaxSize = out.step1() * out.elemSize1() * out.size().height; - z.next_out = (Bytef *)out.data; // 出力ポインタ - z.avail_out = MaxSize; // 出力バッファ残量 - - for(;;) - { - if (z.avail_in == 0) - { - z.next_in = (Bytef *)inbuf; - z.avail_in = fread(inbuf, 1, sizeof(inbuf), fin); - } - - const int status = inflate(&z, Z_NO_FLUSH); - if (status == Z_STREAM_END) - break; - - if (status != Z_OK) - { - fclose(fin); - return false; - } - - if (z.avail_out == 0) - { - fclose(fin); - return false; - } - } - - if (inflateEnd(&z) != Z_OK) - { - fclose(fin); - return false; - } - - fclose(fin); - - return true; - }; - - const auto step1Old = process_image.step1(); - const auto size = process_image.size(); - const auto type = process_image.type(); - compp(process_image, temp); - process_image.release(); - - Reconstruct(false, true, cancel_func, alpha, alpha); - - decompp(size, type, process_image, temp); - boost::filesystem::remove(temp); - - assert(step1Old == process_image.step1()); - } - } - } - - if (isReconstructScale) - { - const cv::Size_ ns(image_size.width * shrinkRatio, image_size.height * shrinkRatio); - if (image_size.width != ns.width || image_size.height != ns.height) - { - int argo = cv::INTER_CUBIC; - if (ratio < 0.5) - argo = cv::INTER_AREA; - - if (!alpha.empty()) - cv::resize(alpha, alpha, ns, 0.0, 0.0, argo); - } - } - - // アルファチャンネルがあったらアルファを付加する - if (!alpha.empty()) - { - std::vector planes; - cv::split(process_image, planes); - process_image.release(); - - planes.push_back(alpha); - alpha.release(); - - cv::merge(planes, process_image); - } - - // 値を0〜1にクリッピング - cv::threshold(process_image, process_image, 1.0, 1.0, cv::THRESH_TRUNC); - cv::threshold(process_image, process_image, 0.0, 0.0, cv::THRESH_TOZERO); - - out = process_image; - - return eWaifu2xError_OK; -} - -namespace -{ - template - void AlphaZeroToZero(std::vector &planes) - { - cv::Mat alpha(planes[3]); - - const T *aptr = (const T *)alpha.data; - - T *ptr0 = (T *)planes[0].data; - T *ptr1 = (T *)planes[1].data; - T *ptr2 = (T *)planes[2].data; - - const size_t Line = alpha.step1(); - const size_t Width = alpha.size().width; - const size_t Height = alpha.size().height; - - for (size_t i = 0; i < Height; i++) - { - for (size_t j = 0; j < Width; j++) - { - const size_t pos = Line * i + j; - - if (aptr[pos] == (T)0) - ptr0[pos] = ptr1[pos] = ptr2[pos] = (T)0; - } - } - } -} - -Waifu2x::eWaifu2xError Waifu2x::waifu2xConvetedMat(const bool isJpeg, const cv::Mat &inMat, cv::Mat &outMat, const waifu2xCancelFunc cancel_func) -{ - Waifu2x::eWaifu2xError ret; - - 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 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, inMat, reconstruct_image, process_image); - if (ret != eWaifu2xError_OK) - return ret; - - const int cv_depth = DepthBitToCVDepth(output_depth); - const double max_val = GetValumeMaxFromCVDepth(cv_depth); - const double eps = GetEPS(cv_depth); - - cv::Mat write_iamge; - if (output_depth != 32) // 出力がfloat形式なら変換しない - process_image.convertTo(write_iamge, cv_depth, max_val, eps); - else - write_iamge = process_image; - - process_image.release(); - - // 完全透明のピクセルの色を消す(処理の都合上、完全透明のピクセルにも色を付けたから) - // モデルによっては画像全域の完全透明の場所にごく小さい値のアルファが広がることがある。それを消すためにcv_depthへ変換してからこの処理を行うことにした - // (ただしcv_depthが32の場合だと意味は無いが) - // TODO: モデル(例えばPhoto)によっては0しかない画像を変換しても0.000114856390とかになるので、適切な値のクリッピングを行う？ - if (write_iamge.channels() > 3) - { - std::vector planes; - cv::split(write_iamge, planes); - write_iamge.release(); - - const auto depth = planes[0].depth(); - switch (depth) - { - case CV_8U: - AlphaZeroToZero(planes); - break; - - case CV_16U: - AlphaZeroToZero(planes); - break; - - case CV_32F: - AlphaZeroToZero(planes); - break; - - case CV_64F: - AlphaZeroToZero(planes); - break; - - default: - return eWaifu2xError_FailedUnknownType; - } - - cv::merge(planes, write_iamge); - } - - outMat = write_iamge; - - return eWaifu2xError_OK; -} - -double Waifu2x::CalcScaleRatio(const cv::Size_ &size) const -{ - if (scale_ratio) - return *scale_ratio; - - if (scale_width) - return (double)*scale_width / (double)size.width; - - return (double)*scale_height / (double)size.height; -} - -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, net_offset); - 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; - - 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, net_offset); - - planes.push_back(alpha); - cv::merge(planes, convert); - } - - float_image = convert; - } - - const auto oldScaleRatio = scale_ratio; - const auto oldScaleWidth = scale_width; - const auto oldScaleHeight = scale_height; - - scale_ratio = factor; - scale_width.reset(); - scale_height.reset(); - - cv::Mat write_iamge; - ret = waifu2xConvetedMat(false, float_image, write_iamge); - - scale_ratio = oldScaleRatio; - scale_width = oldScaleWidth; - scale_height = oldScaleHeight; - - 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; + mIsInited = false; } const std::string& Waifu2x::used_process() const { - return process; -} - -int Waifu2x::DepthBitToCVDepth(const int depth_bit) -{ - switch (depth_bit) - { - case 8: - return CV_8U; - - case 16: - return CV_16U; - - case 32: - return CV_32F; - } - - // 不明だけどとりあえずCV_8Uを返しておく - return CV_8U; -} - -double Waifu2x::GetValumeMaxFromCVDepth(const int cv_depth) -{ - switch (cv_depth) - { - case CV_8U: - return 255.0; - - case CV_16U: - return 65535.0; - - case CV_32F: - return 1.0; - } - - // 不明だけどとりあえず255.0を返しておく - return 255.0; -} - -double Waifu2x::GetEPS(const int cv_depth) -{ - switch (cv_depth) - { - case CV_8U: - return clip_eps8; - - case CV_16U: - return clip_eps16; - - case CV_32F: - return clip_eps32; - } - - // 不明だけどとりあえずclip_eps8返しておく - return clip_eps8; + return mProcess; } diff --git a/common/waifu2x.h b/common/waifu2x.h index 0f20855..e2e22e7 100644 --- a/common/waifu2x.h +++ b/common/waifu2x.h @@ -21,6 +21,10 @@ namespace caffe class NetParameter; }; +class cNet; +class stImage; + + class Waifu2x { public: @@ -59,130 +63,72 @@ public: typedef std::function waifu2xCancelFunc; - struct stOutputExtentionElement - { - std::wstring ext; - std::vector depthList; - boost::optional imageQualityStart; - boost::optional imageQualityEnd; - boost::optional imageQualityDefault; - boost::optional imageQualitySettingVolume; - }; - - const static std::vector OutputExtentionList; + static std::string ExeDir; private: - bool is_inited; + bool mIsInited; - // 一度に処理する画像の幅 - int crop_size; - // 一度に何ブロック分処理するか - int batch_size; + std::string mMode; + int mNoiseLevel; + std::string mProcess; - // ネットに入力する画像のサイズ - int input_block_size; - // ネットに入力するブロックのサイズ(パディングなし) - int block_size; - // ネットワークに入力する画像のサイズ(出力画像の幅はlayer_num * 2だけ小さくなる) - int block_width_height; - // srcnn.prototxtで定義された入力する画像のサイズ - int original_width_height; + bool mIsCuda; - std::string mode; - int noise_level; - boost::optional scale_ratio; - boost::optional scale_width; - boost::optional scale_height; - boost::filesystem::path model_dir; - std::string process; + std::shared_ptr mNoiseNet; + std::shared_ptr mScaleNet; - int inner_padding; - int outer_padding; + int mInputPlane; // ネットへの入力チャンネル数 + int mMaxNetOffset; // ネットに入力するとどれくらい削れるか - int output_block_size; + float *mInputBlock; + size_t mInputBlockSize; - int input_plane; - - bool isCuda; - - boost::shared_ptr> net_noise; - boost::shared_ptr> net_scale; - - float *input_block; - float *dummy_data; - float *output_block; - - bool use_tta; - - boost::optional output_quality; - int output_depth; - - int model_scale; // モデルが対象とする拡大率 - - int net_offset; // ネットに入力するとどれくらい削れるか - int inner_scale; // ネットで拡大される倍率 + float *mOutputBlock; + size_t mOutputBlockSize; private: - static eWaifu2xError LoadMat(cv::Mat &float_image, const boost::filesystem::path &input_file, const int alpha_offset); - 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 boost::filesystem::path &model_path, const boost::filesystem::path ¶m_path, const boost::filesystem::path &info_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 LoadInfoFromJson(const boost::filesystem::path &info_path); - eWaifu2xError SetParameter(caffe::NetParameter ¶m, const std::string &process) const; eWaifu2xError ReconstructImage(boost::shared_ptr> net, const int reconstructed_scale, cv::Mat &im); - static eWaifu2xError WriteMat(const cv::Mat &im, const boost::filesystem::path &output_file, const boost::optional &output_quality); - 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); - 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, cv::Mat &in, cv::Mat &out); + Waifu2x::eWaifu2xError ReconstructImage(const double factor, const int crop_w, const int crop_h, const bool use_tta, const int batch_size, + const bool isReconstructNoise, const bool isReconstructScale, const Waifu2x::waifu2xCancelFunc cancel_func, stImage &image); + Waifu2x::eWaifu2xError ReconstructScale(const int crop_w, const int crop_h, const bool use_tta, const int batch_size, + const Waifu2x::waifu2xCancelFunc cancel_func, stImage &image); + Waifu2x::eWaifu2xError ReconstructByNet(std::shared_ptr net, const int crop_w, const int crop_h, const bool use_tta, const int batch_size, + const Waifu2x::waifu2xCancelFunc cancel_func, cv::Mat &im); + Waifu2x::eWaifu2xError ProcessNet(std::shared_ptr net, const int crop_w, const int crop_h, const bool use_tta, const int batch_size, cv::Mat &im); - eWaifu2xError waifu2xConvetedMat(const bool isJpeg, const cv::Mat &inMat, cv::Mat &outMat, - const waifu2xCancelFunc cancel_func = nullptr); - - double CalcScaleRatio(const cv::Size_ &size) const; - - static int DepthBitToCVDepth(const int depth_bit); - static double GetValumeMaxFromCVDepth(const int cv_depth); - static double GetEPS(const int cv_depth); + // double CalcScaleRatio(const cv::Size_ &size) const; public: Waifu2x(); ~Waifu2x(); - static eWaifu2xcuDNNError can_use_cuDNN(); static eWaifu2xCudaError can_use_CUDA(); + static eWaifu2xcuDNNError can_use_cuDNN(); - static void init_liblary(); + static void init_liblary(int argc, char** argv); static void quit_liblary(); // 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 boost::optional scale_ratio, const boost::optional scale_width, const boost::optional scale_height, - const boost::filesystem::path &model_dir, const std::string &process, - const boost::optional output_quality = boost::optional(), const int output_depth = 8, const bool use_tta = false, - const int crop_size = 128, const int batch_size = 1); - - void destroy(); + eWaifu2xError Init(const std::string &mode, const int noise_level, + const boost::filesystem::path &model_dir, const std::string &process); eWaifu2xError waifu2x(const boost::filesystem::path &input_file, const boost::filesystem::path &output_file, - const waifu2xCancelFunc cancel_func = nullptr); + const double factor, const waifu2xCancelFunc cancel_func = nullptr, const int crop_w = 128, const int crop_h = 128, + const boost::optional output_quality = boost::optional(), const int output_depth = 8, const bool use_tta = false, + const int batch_size = 1); // factor: 倍率 // source: (4チャンネルの場合は)RGBAな画素配列 // dest: (4チャンネルの場合は)処理したRGBAな画素配列 // 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); + eWaifu2xError waifu2x(const double factor, const void* source, void* dest, const int width, const int height, + const int in_channel, const int in_stride, const int out_channel, const int out_stride, + const int crop_w = 128, const int crop_h = 128, const bool use_tta = false, const int batch_size = 1); + + void Destroy(); const std::string& used_process() const; - - static cv::Mat LoadMat(const boost::filesystem::path &path); }; diff --git a/waifu2x-caffe-gui/MainDialog.cpp b/waifu2x-caffe-gui/MainDialog.cpp index 6d0617b..007d23d 100644 --- a/waifu2x-caffe-gui/MainDialog.cpp +++ b/waifu2x-caffe-gui/MainDialog.cpp @@ -14,6 +14,7 @@ #include #include #include "../common/waifu2x.h" +#include "../common/stImage.h" #include "CDialog.h" #include "CControl.h" //#include @@ -295,7 +296,7 @@ bool DialogEvent::SyncMember(const bool NotSyncCropSize, const bool silent) } { - const auto &OutputExtentionList = Waifu2x::OutputExtentionList; + const auto &OutputExtentionList = stImage::OutputExtentionList; const int cur = SendMessage(GetDlgItem(dh, IDC_COMBO_OUT_EXT), CB_GETCURSEL, 0, 0); if (cur < 0 || cur >= OutputExtentionList.size()) @@ -387,8 +388,9 @@ void DialogEvent::SetCropSizeList(const boost::filesystem::path & input_path) int gcd = 1; if (boost::filesystem::exists(input_path) && !boost::filesystem::is_directory(input_path)) { - auto mat = Waifu2x::LoadMat(input_path.string()); - if (mat.empty()) + cv::Mat mat; + const auto ret = stImage::LoadMat(mat, input_path.string()); + if (ret != Waifu2x::eWaifu2xError_OK) return; auto size = mat.size(); @@ -674,25 +676,8 @@ void DialogEvent::ProcessWaifu2x() Waifu2x::eWaifu2xError ret; - boost::optional ScaleRatio; - boost::optional ScaleWidth; - boost::optional ScaleHeight; - - switch (scaleType) - { - case eScaleTypeRatio: - ScaleRatio = scale_ratio; - break; - case eScaleTypeWidth: - ScaleWidth = scale_width; - break; - default: - ScaleHeight = scale_height; - break; - } - Waifu2x w; - ret = w.init(__argc, __argv, mode, noise_level, ScaleRatio, ScaleWidth, ScaleHeight, model_dir, process, output_quality, output_depth, use_tta, crop_size, batch_size); + ret = w.Init(mode, noise_level, model_dir, process); if (ret != Waifu2x::eWaifu2xError_OK) SendMessage(dh, WM_ON_WAIFU2X_ERROR, (WPARAM)&ret, 0); else @@ -723,10 +708,27 @@ void DialogEvent::ProcessWaifu2x() continue; } - ret = w.waifu2x(p.first, p.second, [this]() + double factor; + + //switch (scaleType) + //{ + //case eScaleTypeRatio: + // ScaleRatio = scale_ratio; + // break; + //case eScaleTypeWidth: + // ScaleWidth = scale_width; + // break; + //default: + // ScaleHeight = scale_height; + // break; + //} + + factor = scale_ratio; + + ret = w.waifu2x(p.first, p.second, factor, [this]() { return cancelFlag; - }); + }, crop_size, crop_size, output_quality, output_depth, use_tta, batch_size); num++; ProgessFunc(maxFile, num); @@ -1441,7 +1443,7 @@ void DialogEvent::SetDepthAndQuality(const bool SetDefaultQuality) if (cur < 0) return; - const auto &OutputExtentionList = Waifu2x::OutputExtentionList; + const auto &OutputExtentionList = stImage::OutputExtentionList; if (cur >= OutputExtentionList.size()) return; @@ -1583,7 +1585,7 @@ void DialogEvent::Create(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpData) { HWND houtext = GetDlgItem(dh, IDC_COMBO_OUT_EXT); - const auto &OutputExtentionList = Waifu2x::OutputExtentionList; + const auto &OutputExtentionList = stImage::OutputExtentionList; for (const auto &elm : OutputExtentionList) { SendMessageW(houtext, CB_ADDSTRING, 0, (LPARAM)elm.ext.c_str()); @@ -1845,7 +1847,7 @@ void DialogEvent::Create(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpData) HWND houtext = GetDlgItem(dh, IDC_COMBO_OUT_EXT); size_t defaultIndex = 0; - const auto &OutputExtentionList = Waifu2x::OutputExtentionList; + const auto &OutputExtentionList = stImage::OutputExtentionList; for (size_t i = 0; i < OutputExtentionList.size(); i++) { const auto &elm = OutputExtentionList[i]; diff --git a/waifu2x-caffe-gui/Source.cpp b/waifu2x-caffe-gui/Source.cpp index babf999..412a590 100644 --- a/waifu2x-caffe-gui/Source.cpp +++ b/waifu2x-caffe-gui/Source.cpp @@ -13,7 +13,7 @@ int WINAPI WinMain(HINSTANCE hInstance, LPSTR lpCmdLine, int nCmdShow) { - Waifu2x::init_liblary(); + Waifu2x::init_liblary(__argc, __argv); // 管理者権限で起動してもファイルのドロップを受け付けるようにする ChangeWindowMessageFilter(WM_DROPFILES, MSGFLT_ADD); diff --git a/waifu2x-caffe-gui/waifu2x-caffe-gui.vcxproj b/waifu2x-caffe-gui/waifu2x-caffe-gui.vcxproj index 03f8981..82712cd 100644 --- a/waifu2x-caffe-gui/waifu2x-caffe-gui.vcxproj +++ b/waifu2x-caffe-gui/waifu2x-caffe-gui.vcxproj @@ -95,6 +95,8 @@ + + @@ -103,6 +105,8 @@ + + diff --git a/waifu2x-caffe-gui/waifu2x-caffe-gui.vcxproj.filters b/waifu2x-caffe-gui/waifu2x-caffe-gui.vcxproj.filters index bc5a3cc..866b002 100644 --- a/waifu2x-caffe-gui/waifu2x-caffe-gui.vcxproj.filters +++ b/waifu2x-caffe-gui/waifu2x-caffe-gui.vcxproj.filters @@ -13,14 +13,14 @@ {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 + + {8f85ea70-f834-4cfc-8869-629e95423258} + 繧ｽ繝ｼ繧ｹ 繝輔ぃ繧､繝ｫ - - 繧ｽ繝ｼ繧ｹ 繝輔ぃ繧､繝ｫ - 繧ｽ繝ｼ繧ｹ 繝輔ぃ繧､繝ｫ @@ -33,11 +33,17 @@ 繧ｽ繝ｼ繧ｹ 繝輔ぃ繧､繝ｫ + + common + + + common + + + common + - - 繝倥ャ繝繝ｼ 繝輔ぃ繧､繝ｫ - 繝倥ャ繝繝ｼ 繝輔ぃ繧､繝ｫ @@ -71,6 +77,15 @@ 繝倥ャ繝繝ｼ 繝輔ぃ繧､繝ｫ + + common + + + common + + + common +