waifu2x-caffe/common/waifu2x.cpp

2167 lines
56 KiB
C++
Raw Normal View History

2015-05-29 01:47:26 +09:00
#include "waifu2x.h"
#include <caffe/caffe.hpp>
#include <cudnn.h>
#include <mutex>
#include <opencv2/opencv.hpp>
#include <rapidjson/document.h>
#include <tclap/CmdLine.h>
#include <boost/filesystem.hpp>
#include <boost/algorithm/string.hpp>
#include <chrono>
#include <cuda_runtime.h>
2015-05-29 01:47:26 +09:00
#include <boost/iostreams/stream.hpp>
#include <boost/iostreams/device/file_descriptor.hpp>
#include <google/protobuf/io/coded_stream.h>
#include <google/protobuf/io/zero_copy_stream_impl.h>
#include <google/protobuf/text_format.h>
#include <fcntl.h>
#include <zlib.h>
#ifdef _MSC_VER
#include <io.h>
#endif
2015-07-10 04:09:22 +09:00
#define STB_IMAGE_IMPLEMENTATION
#include <stb_image.h>
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include <stb_image_write.h>
2015-05-29 01:47:26 +09:00
#if defined(WIN32) || defined(WIN64)
#include <Windows.h>
#endif
#define CV_VERSION_STR CVAUX_STR(CV_MAJOR_VERSION) CVAUX_STR(CV_MINOR_VERSION) CVAUX_STR(CV_SUBMINOR_VERSION)
// <20>r<EFBFBD><72><EFBFBD>h<EFBFBD><68><EFBFBD>[<5B>h
#ifdef _DEBUG
#define CV_EXT_STR "d.lib"
#else
#define CV_EXT_STR ".lib"
#endif
#ifdef _MSC_VER
2015-05-29 01:47:26 +09:00
#ifdef _DEBUG
#pragma comment(lib, "caffe-d.lib")
#pragma comment(lib, "proto-d.lib")
#pragma comment(lib, "libboost_system-vc120-mt-gd-1_59.lib")
#pragma comment(lib, "libboost_thread-vc120-mt-gd-1_59.lib")
#pragma comment(lib, "libboost_filesystem-vc120-mt-gd-1_59.lib")
#pragma comment(lib, "glogd.lib")
#pragma comment(lib, "gflagsd.lib")
2015-05-29 01:47:26 +09:00
#pragma comment(lib, "libprotobufd.lib")
#pragma comment(lib, "libhdf5_hl_D.lib")
#pragma comment(lib, "libhdf5_D.lib")
#pragma comment(lib, "zlibstaticd.lib")
#pragma comment(lib, "libopenblas.dll.a")
#pragma comment(lib, "cudart.lib")
#pragma comment(lib, "curand.lib")
#pragma comment(lib, "cublas.lib")
#pragma comment(lib, "cudnn.lib")
#pragma comment(lib, "opencv_core" CV_VERSION_STR CV_EXT_STR)
#pragma comment(lib, "opencv_imgcodecs" CV_VERSION_STR CV_EXT_STR)
#pragma comment(lib, "opencv_imgproc" CV_VERSION_STR CV_EXT_STR)
//#pragma comment(lib, "IlmImf" CV_EXT_STR)
//#pragma comment(lib, "libjasper" CV_EXT_STR)
//#pragma comment(lib, "libjpeg" CV_EXT_STR)
//#pragma comment(lib, "libpng" CV_EXT_STR)
//#pragma comment(lib, "libtiff" CV_EXT_STR)
//#pragma comment(lib, "libwebp" CV_EXT_STR)
//#pragma comment(lib, "ippicvmt.lib")
#pragma comment(lib, "libboost_iostreams-vc120-mt-gd-1_59.lib")
2015-05-29 01:47:26 +09:00
#else
#pragma comment(lib, "caffe.lib")
#pragma comment(lib, "proto.lib")
#pragma comment(lib, "libboost_system-vc120-mt-1_59.lib")
#pragma comment(lib, "libboost_thread-vc120-mt-1_59.lib")
#pragma comment(lib, "libboost_filesystem-vc120-mt-1_59.lib")
#pragma comment(lib, "glog.lib")
#pragma comment(lib, "gflags.lib")
2015-05-29 01:47:26 +09:00
#pragma comment(lib, "libprotobuf.lib")
#pragma comment(lib, "libhdf5_hl.lib")
#pragma comment(lib, "libhdf5.lib")
#pragma comment(lib, "zlibstatic.lib")
#pragma comment(lib, "libopenblas.dll.a")
#pragma comment(lib, "cudart.lib")
#pragma comment(lib, "curand.lib")
#pragma comment(lib, "cublas.lib")
#pragma comment(lib, "cudnn.lib")
#pragma comment(lib, "opencv_core" CV_VERSION_STR CV_EXT_STR)
#pragma comment(lib, "opencv_imgcodecs" CV_VERSION_STR CV_EXT_STR)
#pragma comment(lib, "opencv_imgproc" CV_VERSION_STR CV_EXT_STR)
//#pragma comment(lib, "IlmImf" CV_EXT_STR)
//#pragma comment(lib, "libjasper" CV_EXT_STR)
//#pragma comment(lib, "libjpeg" CV_EXT_STR)
//#pragma comment(lib, "libpng" CV_EXT_STR)
//#pragma comment(lib, "libtiff" CV_EXT_STR)
//#pragma comment(lib, "libwebp" CV_EXT_STR)
//#pragma comment(lib, "ippicvmt.lib")
#pragma comment(lib, "libboost_iostreams-vc120-mt-1_59.lib")
#endif
2015-05-29 01:47:26 +09:00
#endif
// <20><><EFBFBD>͉摜<CD89>̃I<CC83>t<EFBFBD>Z<EFBFBD>b<EFBFBD>g
const int offset = 0;
// srcnn.prototxt<78>Œ<EFBFBD><C592>`<60><><EFBFBD><EFBFBD><EA82BD><EFBFBD>C<EFBFBD><43><EFBFBD>[<5B>̐<EFBFBD>
const int layer_num = 7;
2015-05-29 01:47:26 +09:00
const int ConvertMode = CV_RGB2YUV;
const int ConvertInverseMode = CV_YUV2RGB;
// <20>Œ<EFBFBD><C592><EFBFBD><EFBFBD>K<EFBFBD>v<EFBFBD><76>CUDA<44>h<EFBFBD><68><EFBFBD>C<EFBFBD>o<EFBFBD>[<5B>̃o<CC83>[<5B>W<EFBFBD><57><EFBFBD><EFBFBD>
const int MinCudaDriverVersion = 7050;
// float<61>ȉ摜<C889><E6919C>uint8_t<5F>ȉ摜<C889>ɕϊ<C995><CF8A><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ۂ̎l<CC8E>̌ܓ<CC8C><DC93>Ɏg<C98E><67><EFBFBD>l
2015-12-06 14:33:38 +09:00
// 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.
2015-06-03 03:01:56 +09:00
static std::once_flag waifu2x_once_flag;
static std::once_flag waifu2x_cudnn_once_flag;
static std::once_flag waifu2x_cuda_once_flag;
2015-05-29 01:47:26 +09:00
2015-12-27 06:39:13 +09:00
const std::vector<Waifu2x::stOutputExtentionElement> Waifu2x::OutputExtentionList =
{
{ L".png", { 8, 16 }, boost::optional<int>(), boost::optional<int>(), boost::optional<int>(), boost::optional<int>() },
{ L".bmp", { 8 }, boost::optional<int>(), boost::optional<int>(), boost::optional<int>(), boost::optional<int>() },
{ L".jpg", { 8 }, 0, 100, 95, cv::IMWRITE_JPEG_QUALITY },
{ L".jp2", { 8, 16 }, boost::optional<int>(), boost::optional<int>(), boost::optional<int>(), boost::optional<int>() },
{ L".sr", { 8 }, boost::optional<int>(), boost::optional<int>(), boost::optional<int>(), boost::optional<int>() },
{ L".tif", { 8, 16, 32 }, boost::optional<int>(), boost::optional<int>(), boost::optional<int>(), boost::optional<int>() },
{ L".hdr", { 8, 16, 32 }, boost::optional<int>(), boost::optional<int>(), boost::optional<int>(), boost::optional<int>() },
{ L".exr", { 8, 16, 32 }, boost::optional<int>(), boost::optional<int>(), boost::optional<int>(), boost::optional<int>() },
{ L".ppm", { 8, 16 }, boost::optional<int>(), boost::optional<int>(), boost::optional<int>(), boost::optional<int>() },
{ L".webp", { 8 }, 1, 100, 100, cv::IMWRITE_WEBP_QUALITY },
{ L".tga", { 8 }, 0, 1, 0, 0 },
2015-12-27 06:39:13 +09:00
};
#ifndef CUDA_CHECK_WAIFU2X
#define CUDA_CHECK_WAIFU2X(condition) \
do { \
cudaError_t error = condition; \
if(error != cudaSuccess) throw error; \
} while (0)
#endif
2015-05-29 01:47:26 +09:00
#define CUDA_HOST_SAFE_FREE(ptr) \
do { \
if (ptr) { \
cudaFreeHost(ptr); \
ptr = nullptr; \
} \
} while (0)
#define SAFE_DELETE_WAIFU2X(ptr) \
do { \
if (ptr) { \
delete [] ptr; \
ptr = nullptr; \
} \
} while (0)
2015-07-10 04:09:22 +09:00
namespace
{
class IgnoreErrorCV
{
private:
static int handleError(int status, const char* func_name,
const char* err_msg, const char* file_name,
int line, void* userdata)
{
return 0;
}
public:
IgnoreErrorCV()
{
cv::redirectError(handleError);
}
};
IgnoreErrorCV g_IgnoreErrorCV;
}
template<typename BufType>
static bool writeFile(boost::iostreams::stream<boost::iostreams::file_descriptor> &os, const std::vector<BufType> &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<typename BufType>
static bool writeFile(const boost::filesystem::path &path, std::vector<BufType> &buf)
{
boost::iostreams::stream<boost::iostreams::file_descriptor> os;
try
{
os.open(path, std::ios_base::out | std::ios_base::binary | std::ios_base::trunc);
}
catch (...)
{
return false;
}
return writeFile(os, buf);
}
template<typename BufType>
static bool readFile(boost::iostreams::stream<boost::iostreams::file_descriptor_source> &is, std::vector<BufType> &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<typename BufType>
static bool readFile(const boost::filesystem::path &path, std::vector<BufType> &buf)
{
boost::iostreams::stream<boost::iostreams::file_descriptor_source> 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<boost::iostreams::file_descriptor_source> is;
try
{
is.open(path, std::ios_base::in);
}
catch (...)
{
return Waifu2x::eWaifu2xError_FailedOpenModelFile;
}
if (!is)
return Waifu2x::eWaifu2xError_FailedOpenModelFile;
std::vector<char> 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<boost::iostreams::file_descriptor_source> 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<char> 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<boost::iostreams::file_descriptor> 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)
2015-06-03 03:01:56 +09:00
{
}
Waifu2x::~Waifu2x()
{
destroy();
}
2015-05-29 01:47:26 +09:00
// cuDNN<4E><4E><EFBFBD>g<EFBFBD><67><EFBFBD><EFBFBD>`<60>F<EFBFBD>b<EFBFBD>N<EFBFBD>B<EFBFBD><42><EFBFBD><EFBFBD>Windows<77>̂<EFBFBD>
Waifu2x::eWaifu2xcuDNNError Waifu2x::can_use_cuDNN()
2015-05-29 01:47:26 +09:00
{
static eWaifu2xcuDNNError cuDNNFlag = eWaifu2xcuDNNError_NotFind;
2015-05-29 01:47:26 +09:00
std::call_once(waifu2x_cudnn_once_flag, [&]()
{
#if defined(WIN32) || defined(WIN64)
HMODULE hModule = LoadLibrary(TEXT(CUDNN_DLL_NAME));
2015-05-29 01:47:26 +09:00
if (hModule != NULL)
{
typedef cudnnStatus_t(__stdcall * cudnnCreateType)(cudnnHandle_t *);
typedef cudnnStatus_t(__stdcall * cudnnDestroyType)(cudnnHandle_t);
typedef uint64_t(__stdcall * cudnnGetVersionType)();
2015-05-29 01:47:26 +09:00
cudnnCreateType cudnnCreateFunc = (cudnnCreateType)GetProcAddress(hModule, "cudnnCreate");
cudnnDestroyType cudnnDestroyFunc = (cudnnDestroyType)GetProcAddress(hModule, "cudnnDestroy");
cudnnGetVersionType cudnnGetVersionFunc = (cudnnGetVersionType)GetProcAddress(hModule, "cudnnGetVersion");
if (cudnnCreateFunc != nullptr && cudnnDestroyFunc != nullptr && cudnnGetVersionFunc != nullptr)
2015-05-29 01:47:26 +09:00
{
if (cudnnGetVersionFunc() >= 3000)
2015-05-29 01:47:26 +09:00
{
cudnnHandle_t h;
if (cudnnCreateFunc(&h) == CUDNN_STATUS_SUCCESS)
{
if (cudnnDestroyFunc(h) == CUDNN_STATUS_SUCCESS)
cuDNNFlag = eWaifu2xcuDNNError_OK;
else
cuDNNFlag = eWaifu2xcuDNNError_CannotCreate;
}
else
cuDNNFlag = eWaifu2xcuDNNError_CannotCreate;
2015-05-29 01:47:26 +09:00
}
else
cuDNNFlag = eWaifu2xcuDNNError_OldVersion;
2015-05-29 01:47:26 +09:00
}
else
cuDNNFlag = eWaifu2xcuDNNError_NotFind;
2015-05-29 01:47:26 +09:00
FreeLibrary(hModule);
}
#endif
});
return cuDNNFlag;
}
// CUDA<44><41><EFBFBD>g<EFBFBD><67><EFBFBD><EFBFBD>`<60>F<EFBFBD>b<EFBFBD>N
Waifu2x::eWaifu2xCudaError Waifu2x::can_use_CUDA()
{
static eWaifu2xCudaError CudaFlag = eWaifu2xCudaError_NotFind;
std::call_once(waifu2x_cuda_once_flag, [&]()
{
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 CudaFlag;
}
void Waifu2x::init_liblary()
{
}
void Waifu2x::quit_liblary()
{
}
cv::Mat Waifu2x::LoadMat(const boost::filesystem::path &path)
2015-07-10 04:09:22 +09:00
{
cv::Mat mat;
LoadMat(mat, path);
return mat;
}
Waifu2x::eWaifu2xError Waifu2x::AlphaMakeBorder(std::vector<cv::Mat> &planes, const cv::Mat &alpha, const int offset)
{
// <20><><EFBFBD>̃J<CC83>[<5B>l<EFBFBD><6C><EFBFBD>Ɖ摜<C689>̏􍞂݂<F48D9E82><DD82>s<EFBFBD><73><EFBFBD>ƁA(x, y)<29>𒆐S<F0928690>Ƃ<EFBFBD><C682><EFBFBD>3<EFBFBD>~3<>̈<EFBFBD><CC88>̍<EFBFBD><CC8D>v<EFBFBD>l<EFBFBD><6C><EFBFBD><EFBFBD><EFBFBD>܂<EFBFBD>
const static cv::Mat sum2d_kernel = (cv::Mat_<double>(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); // <20>A<EFBFBD><41><EFBFBD>t<EFBFBD>@<40>`<60><><EFBFBD><EFBFBD><EFBFBD>l<EFBFBD><6C><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>l<EFBFBD><6C><EFBFBD><EFBFBD><EFBFBD>ă}<7D>X<EFBFBD>N<EFBFBD>Ƃ<EFBFBD><C682>Ĉ<EFBFBD><C488><EFBFBD>
cv::Mat mask_nega;
cv::threshold(mask, mask_nega, 0.0, 1.0, cv::THRESH_BINARY_INV); // <20><><EFBFBD>]<5D><><EFBFBD><EFBFBD><EFBFBD>}<7D>X<EFBFBD>N<EFBFBD>i<EFBFBD>l<EFBFBD><6C>1<EFBFBD>̉ӏ<CC89><D38F>͊<EFBFBD><CD8A>S<EFBFBD><53><EFBFBD><EFBFBD><EFBFBD>łȂ<C582><C882>L<EFBFBD><4C><EFBFBD>ȉ<EFBFBD><C889>f<EFBFBD>ƂȂ<C682><C882>j
for (auto &p : planes) // <20><><EFBFBD>S<EFBFBD>ɓ<EFBFBD><C993><EFBFBD><EFBFBD>ȃs<C883>N<EFBFBD>Z<EFBFBD><5A><EFBFBD>ɂ<EFBFBD><C982><EFBFBD><EFBFBD>S<EFBFBD>~<7E><><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
{
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); // <20>}<7D>X<EFBFBD>N<EFBFBD><4E>3<EFBFBD>~3<>̈<EFBFBD><CC88>̍<EFBFBD><CC8D>v<EFBFBD>l<EFBFBD><6C><EFBFBD><EFBFBD><EFBFBD>߂<EFBFBD>
cv::Mat mask_nega_u8;
mask_nega.convertTo(mask_nega_u8, CV_8U, 255.0, clip_eps8); // mask_nega<67><61>CV_U8<55>ŁiOpenCV<43><56>API<50><49><EFBFBD>K<EFBFBD>v<EFBFBD>ɂȂ<C982><C882>j
for (auto &p : planes) // 1<>`<60><><EFBFBD><EFBFBD><EFBFBD>l<EFBFBD><6C><EFBFBD><EFBFBD><EFBFBD><EFBFBD><C28F><EFBFBD>
{
// <20>`<60><><EFBFBD><EFBFBD><EFBFBD>l<EFBFBD><6C><EFBFBD><EFBFBD>3<EFBFBD>~3<>̈<EFBFBD><CC88><EFBFBD><EFBFBD>̗L<CC97><4C><EFBFBD><EFBFBD><EFBFBD>f<EFBFBD>̕<EFBFBD><CC95>ϒl<CF92><6C><EFBFBD><EFBFBD><EFBFBD>߂<EFBFBD>
cv::Mat border;
cv::filter2D(p, border, -1, sum2d_kernel, cv::Point(-1, -1), 0, cv::BORDER_DEFAULT);
border /= mask_weight;
// <20>`<60><><EFBFBD><EFBFBD><EFBFBD>l<EFBFBD><6C><EFBFBD>̗L<CC97><4C><EFBFBD>ȉ<EFBFBD><C889>f<EFBFBD>̕<EFBFBD><CC95><EFBFBD><EFBFBD>ɁA<C981>v<EFBFBD>Z<EFBFBD><5A><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ϒl<CF92><6C><EFBFBD>R<EFBFBD>s<EFBFBD>[
border.copyTo(p, mask_nega_u8);
}
// <20>}<7D>X<EFBFBD>N<EFBFBD><4E>1<EFBFBD><31><EFBFBD>c<EFBFBD><63><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>̂<EFBFBD><CC82>V<EFBFBD><56><EFBFBD><EFBFBD><EFBFBD>}<7D>X<EFBFBD>N<EFBFBD>Ƃ<EFBFBD><C682><EFBFBD>(<28>}<7D>X<EFBFBD>N<EFBFBD><4E>3<EFBFBD>~3<>̈<EFBFBD><CC88>̍<EFBFBD><CC8D>v<EFBFBD>l<EFBFBD><6C><EFBFBD><EFBFBD><EFBFBD>߂<EFBFBD><DF82><EFBFBD><EFBFBD>̂̔<CC82>0<EFBFBD>̈<EFBFBD><CC88>́A<CD81>}<7D>X<EFBFBD>N<EFBFBD><4E>1<EFBFBD><31><EFBFBD>c<EFBFBD><63><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>̗̂̈<CC97><CC88>ɓ<EFBFBD><C993><EFBFBD><EFBFBD><EFBFBD>)
cv::threshold(mask_weight, mask, 0.0, 1.0, cv::THRESH_BINARY);
// <20>V<EFBFBD><56><EFBFBD><EFBFBD><EFBFBD>}<7D>X<EFBFBD>N<EFBFBD>̔<EFBFBD><CC94>]<5D><><EFBFBD><EFBFBD><EFBFBD>}<7D>X<EFBFBD>N<EFBFBD><4E><EFBFBD>v<EFBFBD>Z
cv::threshold(mask, mask_nega, 0.0, 1.0, cv::THRESH_BINARY_INV);
}
// <20><><EFBFBD>f<EFBFBD><66>0<EFBFBD><30><EFBFBD><EFBFBD>1<EFBFBD>ɃN<C983><4E><EFBFBD>b<EFBFBD>s<EFBFBD><73><EFBFBD>O
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;
}
2015-05-29 01:47:26 +09:00
// <20><EFBFBD><E6919C><EFBFBD>ǂݍ<C782><DD8D><EFBFBD><EFBFBD>Œl<C592><6C>0.0f<EFBFBD>`1.0f<EFBFBD>͈̔͂ɕϊ<EFBFBD>
Waifu2x::eWaifu2xError Waifu2x::LoadMat(cv::Mat &float_image, const boost::filesystem::path &input_file)
2015-05-29 01:47:26 +09:00
{
cv::Mat original_image;
2015-07-10 04:09:22 +09:00
{
std::vector<char> 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;
}
2015-07-10 04:09:22 +09:00
}
2015-05-29 01:47:26 +09:00
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; // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD>0.0<EFBFBD>`1.0<EFBFBD>̂͂<EFBFBD><EFBFBD>Ȃ̂ŕϊ<EFBFBD><EFBFBD>͕K<EFBFBD>v<EFBFBD>Ȃ<EFBFBD>
break;
}
2015-05-29 01:47:26 +09:00
original_image.release();
if (convert.channels() == 1)
cv::cvtColor(convert, convert, cv::COLOR_GRAY2BGR);
else if (convert.channels() == 4)
{
// <20>A<EFBFBD><41><EFBFBD>t<EFBFBD>@<40>`<60><><EFBFBD><EFBFBD><EFBFBD>l<EFBFBD><6C><EFBFBD>t<EFBFBD><74><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><E793A7><EFBFBD>ȃs<C883>N<EFBFBD>Z<EFBFBD><5A><EFBFBD>̂ƕs<C695><73><EFBFBD><EFBFBD><EFBFBD>ȃs<C883>N<EFBFBD>Z<EFBFBD><5A><EFBFBD>̋<EFBFBD><CC8B>E<EFBFBD><45><EFBFBD><EFBFBD><EFBFBD>̐F<CC90><46><EFBFBD>L<EFBFBD><4C><EFBFBD><EFBFBD>
std::vector<cv::Mat> planes;
cv::split(convert, planes);
cv::Mat alpha = planes[3];
planes.resize(3);
AlphaMakeBorder(planes, alpha, layer_num);
planes.push_back(alpha);
cv::merge(planes, convert);
}
float_image = convert;
2015-05-29 01:47:26 +09:00
return eWaifu2xError_OK;
}
Waifu2x::eWaifu2xError Waifu2x::LoadMatBySTBI(cv::Mat &float_image, const std::vector<char> &img_data)
2015-07-10 04:09:22 +09:00
{
int x, y, comp;
stbi_uc *data = stbi_load_from_memory((const stbi_uc *)img_data.data(), img_data.size(), &x, &y, &comp, 0);
2015-07-10 04:09:22 +09:00
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<47><42><EFBFBD><EFBFBD><EFBFBD><EFBFBD>BGR<47>ɕϊ<C995>
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;
}
2015-05-29 01:47:26 +09:00
// <20><EFBFBD><E6919C><EFBFBD><EFBFBD><EFBFBD>P<EFBFBD>x<EFBFBD>̉摜<CC89><E6919C><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>o<EFBFBD><6F>
2015-06-03 03:01:56 +09:00
Waifu2x::eWaifu2xError Waifu2x::CreateBrightnessImage(const cv::Mat &float_image, cv::Mat &im)
2015-05-29 01:47:26 +09:00
{
if (float_image.channels() > 1)
{
cv::Mat converted_color;
cv::cvtColor(float_image, converted_color, ConvertMode);
2015-05-29 01:47:26 +09:00
std::vector<cv::Mat> planes;
cv::split(converted_color, planes);
2015-05-29 01:47:26 +09:00
im = planes[0];
planes.clear();
}
else
im = float_image;
2015-05-29 01:47:26 +09:00
return eWaifu2xError_OK;
}
// <20><><EFBFBD>͉摜<CD89><E6919C>(Photoshop<6F>ł<EFBFBD><C582><EFBFBD>)<29>L<EFBFBD><4C><EFBFBD><EFBFBD><EFBFBD>o<EFBFBD>X<EFBFBD>T<EFBFBD>C<EFBFBD>Y<EFBFBD><59>output_size<7A>̔{<7B><><EFBFBD>ɕύX
// <20><EFBFBD>͍<EFBFBD><CD8D><EFBFBD><EFBFBD>z<EFBFBD>u<EFBFBD>A<EFBFBD>]<5D><><EFBFBD><EFBFBD>cv::BORDER_REPLICATE<54>Ŗ<EFBFBD><C596>߂<EFBFBD>
2015-06-03 03:01:56 +09:00
Waifu2x::eWaifu2xError Waifu2x::PaddingImage(const cv::Mat &input, cv::Mat &output)
2015-05-29 01:47:26 +09:00
{
const auto h_blocks = (int)floor(input.size().width / output_size) + (input.size().width % output_size == 0 ? 0 : 1);
const auto w_blocks = (int)floor(input.size().height / output_size) + (input.size().height % output_size == 0 ? 0 : 1);
const auto height = offset + h_blocks * output_size + offset;
const auto width = offset + w_blocks * output_size + offset;
const auto pad_h1 = offset;
const auto pad_w1 = offset;
const auto pad_h2 = (height - offset) - input.size().width;
const auto pad_w2 = (width - offset) - input.size().height;
cv::copyMakeBorder(input, output, pad_w1, pad_w2, pad_h1, pad_h2, cv::BORDER_REPLICATE);
return eWaifu2xError_OK;
}
// <20><EFBFBD><E6919C>cv::INTER_NEAREST<53>œ<EFBFBD><C593>{<7B>Ɋg<C98A><EFBFBD>āAPaddingImage()<29>Ńp<C583>f<EFBFBD>B<EFBFBD><42><EFBFBD>O<EFBFBD><4F><EFBFBD><EFBFBD>
2015-06-03 03:01:56 +09:00
Waifu2x::eWaifu2xError Waifu2x::Zoom2xAndPaddingImage(const cv::Mat &input, cv::Mat &output, cv::Size_<int> &zoom_size)
2015-05-29 01:47:26 +09:00
{
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);
2015-05-29 01:47:26 +09:00
return PaddingImage(output, output);
2015-05-29 01:47:26 +09:00
}
// <20><><EFBFBD>͉摜<CD89><E6919C>zoom_size<7A>̑傫<CC91><E582AB><EFBFBD><EFBFBD>cv::INTER_CUBIC<49>Ŋg<C58A><EFBFBD>A<EFBFBD>F<EFBFBD><46><EFBFBD><EFBFBD><EFBFBD>݂̂<CC82><DD82>c<EFBFBD><63>
2015-06-03 03:01:56 +09:00
Waifu2x::eWaifu2xError Waifu2x::CreateZoomColorImage(const cv::Mat &float_image, const cv::Size_<int> &zoom_size, std::vector<cv::Mat> &cubic_planes)
2015-05-29 01:47:26 +09:00
{
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();
// <20><><EFBFBD><EFBFBD>Y<EFBFBD><59><EFBFBD><EFBFBD><EFBFBD>͎g<CD8E><67><EFBFBD>Ȃ<EFBFBD><C882>̂ʼn<CC82><C589><EFBFBD>
cubic_planes[0].release();
return eWaifu2xError_OK;
}
// <20><><EFBFBD>f<EFBFBD><66><EFBFBD>t<EFBFBD>@<40>C<EFBFBD><43><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>l<EFBFBD>b<EFBFBD>g<EFBFBD><67><EFBFBD>[<5B>N<EFBFBD><4E><EFBFBD>\<5C>z
// process<73><73>cudnn<6E><6E><EFBFBD>w<EFBFBD><EFBFBD><E882B3><EFBFBD>Ȃ<EFBFBD><C882><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EA8D87>cuDNN<4E><4E><EFBFBD>Ăяo<D18F><6F><EFBFBD><EFBFBD><EFBFBD>Ȃ<EFBFBD><C882><EFBFBD>ɕύX<CF8D><58><EFBFBD><EFBFBD>
Waifu2x::eWaifu2xError Waifu2x::ConstractNet(boost::shared_ptr<caffe::Net<float>> &net, const boost::filesystem::path &model_path, const boost::filesystem::path &param_path, const std::string &process)
2015-05-29 01:47:26 +09:00
{
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, &param_model);
const auto retParamBin = readProtoBinary(caffemodel_path, &param_caffemodel);
if (retModelBin == eWaifu2xError_OK && retParamBin == eWaifu2xError_OK)
2015-05-29 01:47:26 +09:00
{
Waifu2x::eWaifu2xError ret;
ret = SetParameter(param_model, process);
if (ret != eWaifu2xError_OK)
return ret;
2015-05-29 01:47:26 +09:00
if (!caffe::UpgradeNetAsNeeded(caffemodel_path.string(), &param_caffemodel))
return Waifu2x::eWaifu2xError_FailedParseModelFile;
net = boost::shared_ptr<caffe::Net<float>>(new caffe::Net<float>(param_model));
net->CopyTrainedLayersFrom(param_caffemodel);
2015-05-29 01:47:26 +09:00
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;
}
2015-05-29 01:47:26 +09:00
return eWaifu2xError_OK;
}
2015-05-29 01:47:26 +09:00
Waifu2x::eWaifu2xError Waifu2x::SetParameter(caffe::NetParameter &param, const std::string &process) const
{
param.mutable_state()->set_phase(caffe::TEST);
2015-05-29 01:47:26 +09:00
{
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")
2015-05-29 01:47:26 +09:00
{
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);
}
}
2015-05-29 01:47:26 +09:00
return eWaifu2xError_OK;
}
2015-05-29 01:47:26 +09:00
Waifu2x::eWaifu2xError Waifu2x::LoadParameterFromJson(boost::shared_ptr<caffe::Net<float>> &net, const boost::filesystem::path &model_path, const boost::filesystem::path &param_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, &param);
if (ret != eWaifu2xError_OK)
return ret;
ret = writeProtoBinary(param, modelbin_path);
if (ret != eWaifu2xError_OK)
return ret;
2015-05-29 01:47:26 +09:00
ret = SetParameter(param, process);
if (ret != eWaifu2xError_OK)
return ret;
2015-05-29 01:47:26 +09:00
net = boost::shared_ptr<caffe::Net<float>>(new caffe::Net<float>(param));
2015-05-29 01:47:26 +09:00
rapidjson::Document d;
std::vector<char> jsonBuf;
2015-05-29 01:47:26 +09:00
try
{
2016-03-21 17:16:39 +09:00
boost::iostreams::stream<boost::iostreams::file_descriptor_source> 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);
2015-05-29 01:47:26 +09:00
jsonBuf.resize(size + 1);
is.read(jsonBuf.data(), jsonBuf.size());
jsonBuf[jsonBuf.size() - 1] = '\0';
d.Parse(jsonBuf.data());
}
catch (...)
{
return eWaifu2xError_FailedParseModelFile;
}
2015-05-29 01:47:26 +09:00
if (d.Size() != 7)
return eWaifu2xError_FailedParseModelFile;
2015-05-29 01:47:26 +09:00
int inputPlane = 0;
int outputPlane = 0;
try
{
inputPlane = d[0]["nInputPlane"].GetInt();
outputPlane = d[d.Size() - 1]["nOutputPlane"].GetInt();
}
catch (...)
{
return eWaifu2xError_FailedParseModelFile;
2015-05-29 01:47:26 +09:00
}
if (inputPlane == 0 || outputPlane == 0)
return eWaifu2xError_FailedParseModelFile;
2015-05-29 01:47:26 +09:00
if (inputPlane != outputPlane)
return eWaifu2xError_FailedParseModelFile;
2015-05-29 01:47:26 +09:00
//if (param.layer_size() < 17)
// return eWaifu2xError_FailedParseModelFile;
2015-05-29 01:47:26 +09:00
std::vector<boost::shared_ptr<caffe::Layer<float>>> list;
auto &v = net->layers();
for (auto &l : v)
2015-05-29 01:47:26 +09:00
{
auto lk = l->type();
auto &bv = l->blobs();
if (bv.size() > 0)
list.push_back(l);
}
try
{
std::vector<float> weightList;
std::vector<float> biasList;
int count = 0;
for (auto it = d.Begin(); it != d.End(); ++it)
2015-05-29 01:47:26 +09:00
{
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();
}
2015-05-29 01:47:26 +09:00
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(&param);
ret = writeProtoBinary(param, caffemodel_path);
if (ret != eWaifu2xError_OK)
return ret;
}
catch (...)
{
return eWaifu2xError_FailedConstructModel;
2015-05-29 01:47:26 +09:00
}
input_plane = inputPlane;
2015-05-29 01:47:26 +09:00
return eWaifu2xError_OK;
}
// <20>l<EFBFBD>b<EFBFBD>g<EFBFBD><67><EFBFBD>[<5B>N<EFBFBD><4E><EFBFBD>g<EFBFBD><67><EFBFBD>ĉ摜<C489><E6919C><EFBFBD>č\<5C>z<EFBFBD><7A><EFBFBD><EFBFBD>
2015-06-03 03:01:56 +09:00
Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr<caffe::Net<float>> net, cv::Mat &im)
2015-05-29 01:47:26 +09:00
{
const auto Height = im.size().height;
const auto Width = im.size().width;
const auto Line = im.step1();
assert(Width % output_size == 0);
assert(Height % output_size == 0);
assert(im.channels() == 1 || im.channels() == 3);
2015-05-29 01:47:26 +09:00
cv::Mat outim(im.rows, im.cols, im.type());
// float *imptr = (float *)im.data;
float *imptr = (float *)outim.data;
2015-05-29 01:47:26 +09:00
try
{
auto input_blobs = net->input_blobs();
auto input_blob = net->input_blobs()[0];
2015-05-29 01:47:26 +09:00
input_blob->Reshape(batch_size, input_plane, input_block_size, input_block_size);
2015-05-29 01:47:26 +09:00
assert(im.channels() == input_plane);
assert(input_blob->shape(1) == input_plane);
const int WidthNum = Width / output_size;
const int HeightNum = Height / output_size;
2015-05-29 01:47:26 +09:00
const int BlockNum = WidthNum * HeightNum;
const int input_block_plane_size = input_block_size * input_block_size * input_plane;
const int output_block_plane_size = output_block_size * output_block_size * input_plane;
const int output_padding = inner_padding + outer_padding - layer_num;
2015-05-29 01:47:26 +09:00
// <20><EFBFBD><E6919C>(<28><><EFBFBD><EFBFBD><EF8381><EFBFBD><EFBFBD><EFBFBD>̓s<CC93><73><EFBFBD><EFBFBD>)output_size*output_size<7A>ɕ<EFBFBD><C995><EFBFBD><EFBFBD>čč\<5C>z<EFBFBD><7A><EFBFBD><EFBFBD>
for (int num = 0; num < BlockNum; num += batch_size)
2015-05-29 01:47:26 +09:00
{
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++)
2015-05-29 01:47:26 +09:00
{
const int wn = (num + n) % WidthNum;
const int hn = (num + n) / WidthNum;
const int w = wn * output_size;
const int h = hn * output_size;
if (w + crop_size <= Width && h + crop_size <= Height)
2015-05-29 01:47:26 +09:00
{
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)
2015-05-29 01:47:26 +09:00
{
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;
}
2015-05-29 01:47:26 +09:00
if (y + height > Height)
{
bottom += (y + height) - Height;
height = Height - y;
}
cv::Mat someimg = im(cv::Rect(x, y, width, height));
cv::Mat someborderimg;
// <20><EFBFBD>𒆉<EFBFBD><F0928689>Ƀp<C983>f<EFBFBD>B<EFBFBD><42><EFBFBD>O<EFBFBD>B<EFBFBD>]<5D><><EFBFBD><EFBFBD>cv::BORDER_REPLICATE<54>Ŗ<EFBFBD><C596>߂<EFBFBD>
2015-06-20 04:02:10 +09:00
// <20><><EFBFBD><EFBFBD>im<69>ʼn<EFBFBD><C589>f<EFBFBD><66><EFBFBD><EFBFBD><EFBFBD>݂<EFBFBD><DD82><EFBFBD><E99594><EFBFBD>͗]<5D><><EFBFBD>ƔF<C694><46><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ȃ<EFBFBD><C882><EFBFBD><EFBFBD>Ainner_padding<6E><67>layer_num<75><6D>outer_padding<6E><67>1<EFBFBD>ȏ<EFBFBD><C88F>Ȃ炻<C882><E782BB><EFBFBD>̕<EFBFBD><CC95><EFBFBD><EFBFBD>̉<EFBFBD><CC89>f<EFBFBD>͌<EFBFBD><CD8C>ʉ摜<CA89>Ƃ<EFBFBD><C682>Ď<EFBFBD><C48E><EFBFBD><EFBFBD>o<EFBFBD><6F><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ɂ͉e<CD89><65><EFBFBD><EFBFBD><EFBFBD>Ȃ<EFBFBD>
cv::copyMakeBorder(someimg, someborderimg, top, bottom, left, right, cv::BORDER_REPLICATE);
someimg.release();
// <20><EFBFBD>𒼗<EFBFBD><F092BC97>ɕϊ<C995>
{
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
2015-05-29 01:47:26 +09:00
{
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];
}
}
2015-05-29 01:47:26 +09:00
}
}
}
}
2015-05-29 01:47:26 +09:00
assert(input_blob->count() == input_block_plane_size * processNum);
// <20>l<EFBFBD>b<EFBFBD>g<EFBFBD><67><EFBFBD>[<5B>N<EFBFBD>ɉ摜<C989><E6919C><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
input_blob->set_cpu_data(input_block);
2015-05-29 01:47:26 +09:00
// <20>v<EFBFBD>Z
auto out = net->ForwardPrefilled(nullptr);
2015-05-29 01:47:26 +09:00
auto b = out[0];
2015-05-29 01:47:26 +09:00
assert(b->count() == output_block_plane_size * processNum);
2015-05-29 01:47:26 +09:00
const float *ptr = nullptr;
2015-05-29 01:47:26 +09:00
if (caffe::Caffe::mode() == caffe::Caffe::CPU)
ptr = b->cpu_data();
else
ptr = b->gpu_data();
2015-05-29 01:47:26 +09:00
caffe::caffe_copy(output_block_plane_size * processNum, ptr, output_block);
for (int n = 0; n < processNum; n++)
{
const int wn = (num + n) % WidthNum;
const int hn = (num + n) / WidthNum;
const int w = wn * output_size;
const int h = hn * output_size;
const float *fptr = output_block + (output_block_plane_size * n);
// <20><><EFBFBD>ʂ<EFBFBD><CA82>o<EFBFBD>͉摜<CD89>ɃR<C983>s<EFBFBD>[
if (outim.channels() == 1)
{
for (int i = 0; i < crop_size; i++)
memcpy(imptr + (h + i) * Line + w, fptr + (i + output_padding) * output_block_size + output_padding, crop_size * sizeof(float));
}
else
{
const auto LinePixel = outim.step1() / outim.channels();
const auto Channel = outim.channels();
for (int i = 0; i < crop_size; i++)
{
for (int j = 0; j < crop_size; j++)
{
for (int ch = 0; ch < Channel; ch++)
imptr[((h + i) * LinePixel + (w + j)) * Channel + ch] = fptr[(ch * output_block_size + i + output_padding) * output_block_size + j + output_padding];
}
}
}
2015-05-29 01:47:26 +09:00
}
}
}
catch (...)
{
return eWaifu2xError_FailedProcessCaffe;
}
im = outim;
2015-05-29 01:47:26 +09:00
return eWaifu2xError_OK;
}
Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &Mode, const int NoiseLevel,
const boost::optional<double> ScaleRatio, const boost::optional<int> ScaleWidth, const boost::optional<int> ScaleHeight,
const boost::filesystem::path &ModelDir, const std::string &Process,
2015-12-27 06:39:13 +09:00
const boost::optional<int> OutputQuality, const int OutputDepth, const bool UseTTA, const int CropSize, const int BatchSize)
2015-05-29 01:47:26 +09:00
{
2015-06-03 03:01:56 +09:00
Waifu2x::eWaifu2xError ret;
2015-05-29 01:47:26 +09:00
2015-06-03 03:01:56 +09:00
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.)
2015-06-03 03:01:56 +09:00
return eWaifu2xError_InvalidParameter;
try
{
mode = Mode;
noise_level = NoiseLevel;
scale_ratio = ScaleRatio;
scale_width = ScaleWidth;
scale_height = ScaleHeight;
model_dir = ModelDir;
process = Process;
2015-11-19 01:50:11 +09:00
use_tta = UseTTA;
2015-05-29 01:47:26 +09:00
2015-12-27 06:39:13 +09:00
output_quality = OutputQuality;
output_depth = OutputDepth;
crop_size = CropSize;
batch_size = BatchSize;
inner_padding = layer_num;
outer_padding = 1;
output_size = crop_size - offset * 2;
input_block_size = crop_size + (inner_padding + outer_padding) * 2;
original_width_height = 128 + layer_num * 2;
output_block_size = crop_size + (inner_padding + outer_padding - layer_num) * 2;
std::call_once(waifu2x_once_flag, [argc, argv]()
{
assert(argc >= 1);
2015-05-29 01:47:26 +09:00
int tmpargc = 1;
char* tmpargvv[] = { argv[0] };
char** tmpargv = tmpargvv;
// glog<6F><67><EFBFBD>̏<EFBFBD><CC8F><EFBFBD><EFBFBD><EFBFBD>
caffe::GlobalInit(&tmpargc, &tmpargv);
});
2015-05-29 01:47:26 +09:00
const auto cuDNNCheckStartTime = std::chrono::system_clock::now();
if (process == "gpu")
{
if (can_use_CUDA() != eWaifu2xCudaError_OK)
return eWaifu2xError_FailedCudaCheck;
// cuDNN<4E><4E><EFBFBD>g<EFBFBD><67><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ȃ<EFBFBD>cuDNN<4E><4E><EFBFBD>g<EFBFBD><67>
else if (can_use_cuDNN() == eWaifu2xcuDNNError_OK)
process = "cudnn";
}
2015-05-29 01:47:26 +09:00
const auto cuDNNCheckEndTime = std::chrono::system_clock::now();
boost::filesystem::path mode_dir_path(model_dir);
if (!mode_dir_path.is_absolute()) // model_dir<69><72><EFBFBD><EFBFBD><EFBFBD>΃p<CE83>X<EFBFBD>Ȃ<EFBFBD><C882><EFBFBD><EFBFBD>΃p<CE83>X<EFBFBD>ɒ<EFBFBD><C992><EFBFBD>
2015-05-29 01:47:26 +09:00
{
// <20>܂<EFBFBD><DC82>̓J<CD83><4A><EFBFBD><EFBFBD><EFBFBD>g<EFBFBD>f<EFBFBD>B<EFBFBD><42><EFBFBD>N<EFBFBD>g<EFBFBD><67><EFBFBD><EFBFBD><EFBFBD>ɂ<EFBFBD><C982><EFBFBD>T<EFBFBD><54>
mode_dir_path = boost::filesystem::absolute(model_dir);
if (!boost::filesystem::exists(mode_dir_path) && argc >= 1) // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>argv[0]<5D><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>s<EFBFBD>t<EFBFBD>@<40>C<EFBFBD><43><EFBFBD>̂<EFBFBD><CC82><EFBFBD><EFBFBD>t<EFBFBD>H<EFBFBD><48><EFBFBD>_<EFBFBD>𐄒肵<F0908492>A<EFBFBD><41><EFBFBD>̃t<CC83>H<EFBFBD><48><EFBFBD>_<EFBFBD><5F><EFBFBD>ɂ<EFBFBD><C982><EFBFBD>T<EFBFBD><54>
{
boost::filesystem::path a0(argv[0]);
if (a0.is_absolute())
mode_dir_path = a0.branch_path() / model_dir;
}
2015-05-29 01:47:26 +09:00
}
if (!boost::filesystem::exists(mode_dir_path))
return eWaifu2xError_FailedOpenModelFile;
2015-05-29 01:47:26 +09:00
if (process == "cpu")
{
caffe::Caffe::set_mode(caffe::Caffe::CPU);
isCuda = false;
}
else
{
caffe::Caffe::set_mode(caffe::Caffe::GPU);
isCuda = true;
}
2015-05-29 01:47:26 +09:00
if (mode == "noise" || mode == "noise_scale" || mode == "auto_scale")
{
const boost::filesystem::path model_path = (mode_dir_path / "srcnn.prototxt").string();
const boost::filesystem::path param_path = (mode_dir_path / ("noise" + std::to_string(noise_level) + "_model.json")).string();
2015-05-29 01:47:26 +09:00
ret = ConstractNet(net_noise, model_path, param_path, process);
if (ret != eWaifu2xError_OK)
return ret;
}
2015-05-29 01:47:26 +09:00
if (mode == "scale" || mode == "noise_scale" || mode == "auto_scale")
{
const boost::filesystem::path model_path = (mode_dir_path / "srcnn.prototxt").string();
const boost::filesystem::path param_path = (mode_dir_path / "scale2.0x_model.json").string();
2015-05-29 01:47:26 +09:00
ret = ConstractNet(net_scale, model_path, param_path, process);
if (ret != eWaifu2xError_OK)
return ret;
}
const int input_block_plane_size = input_block_size * input_block_size * input_plane;
const int output_block_plane_size = output_block_size * output_block_size * input_plane;
if (isCuda)
{
CUDA_CHECK_WAIFU2X(cudaHostAlloc(&input_block, sizeof(float) * input_block_plane_size * batch_size, cudaHostAllocWriteCombined));
CUDA_CHECK_WAIFU2X(cudaHostAlloc(&dummy_data, sizeof(float) * input_block_plane_size * batch_size, cudaHostAllocWriteCombined));
CUDA_CHECK_WAIFU2X(cudaHostAlloc(&output_block, sizeof(float) * output_block_plane_size * batch_size, cudaHostAllocDefault));
}
else
{
input_block = new float[input_block_plane_size * batch_size];
dummy_data = new float[input_block_plane_size * batch_size];
output_block = new float[output_block_plane_size * batch_size];
}
2015-05-29 01:47:26 +09:00
for (size_t i = 0; i < input_block_plane_size * batch_size; i++)
dummy_data[i] = 0.0f;
is_inited = true;
}
catch (...)
{
return eWaifu2xError_InvalidParameter;
}
2015-05-29 01:47:26 +09:00
2015-06-03 03:01:56 +09:00
return eWaifu2xError_OK;
}
2015-05-29 01:47:26 +09:00
2015-06-03 03:01:56 +09:00
void Waifu2x::destroy()
{
net_noise.reset();
net_scale.reset();
2015-05-29 01:47:26 +09:00
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);
}
2015-06-03 03:01:56 +09:00
is_inited = false;
}
2015-05-29 01:47:26 +09:00
2015-12-06 21:15:45 +09:00
static void Waifu2x_stbi_write_func(void *context, void *data, int size)
{
boost::iostreams::stream<boost::iostreams::file_descriptor> *osp = (boost::iostreams::stream<boost::iostreams::file_descriptor> *)context;
osp->write((const char *)data, size);
}
2016-02-04 14:14:40 +09:00
Waifu2x::eWaifu2xError Waifu2x::WriteMat(const cv::Mat &im, const boost::filesystem::path &output_file, const boost::optional<int> &output_quality)
2015-07-10 04:09:22 +09:00
{
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<unsigned char> rgbimg;
if (im.channels() >= 3 || im.step1() != im.size().width * im.channels()) // RGB<47>p<EFBFBD>o<EFBFBD>b<EFBFBD>t<EFBFBD>@<40>ɃR<C983>s<EFBFBD>[(<28><><EFBFBD><EFBFBD>̓p<CD83>f<EFBFBD>B<EFBFBD><42><EFBFBD>O<EFBFBD><4F><EFBFBD>Ƃ<EFBFBD>)
{
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<47><52>RGB<47>ɕ<EFBFBD><C995>ёւ<D191>
{
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<boost::iostreams::file_descriptor> 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<4C><45><EFBFBD>k<EFBFBD>̐ݒ<CC90>
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;
}
}
// <20>ݒ肳<DD92><E882B3><EFBFBD>Ȃ<EFBFBD><C882><EFBFBD><EFBFBD><EFBFBD><EFBFBD>̂Ńf<C583>t<EFBFBD>H<EFBFBD><48><EFBFBD>g<EFBFBD>ɂ<EFBFBD><C982><EFBFBD>
if (!isSet)
stbi_write_tga_with_rle = 1;
2015-12-06 21:15:45 +09:00
if (!stbi_write_tga_to_func(Waifu2x_stbi_write_func, &os, im.size().width, im.size().height, im.channels(), data))
2015-07-10 04:09:22 +09:00
return eWaifu2xError_FailedOpenOutputFile;
return eWaifu2xError_OK;
}
try
{
2015-12-27 06:39:13 +09:00
const boost::filesystem::path op(output_file);
const boost::filesystem::path opext(op.extension());
std::vector<int> 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<uchar> buf;
2015-12-27 06:39:13 +09:00
cv::imencode(ext, im, buf, params);
if (writeFile(output_file, buf))
2015-07-10 04:09:22 +09:00
return eWaifu2xError_OK;
}
catch (...)
{
}
return eWaifu2xError_FailedOpenOutputFile;
}
Waifu2x::eWaifu2xError Waifu2x::BeforeReconstructFloatMatProcess(const cv::Mat &in, cv::Mat &out, bool &convertBGRflag)
2015-06-03 03:01:56 +09:00
{
Waifu2x::eWaifu2xError ret;
2015-05-29 01:47:26 +09:00
convertBGRflag = false;
2015-06-03 03:01:56 +09:00
cv::Mat im;
if (input_plane == 1)
2015-11-19 01:50:11 +09:00
CreateBrightnessImage(in, im);
else
{
im = in;
if (in.channels() == 1)
{
cv::cvtColor(in, im, CV_GRAY2BGR);
convertBGRflag = true;
}
std::vector<cv::Mat> planes;
cv::split(im, planes);
if (im.channels() == 4)
planes.resize(3);
// BGR<47><52><EFBFBD><EFBFBD>RGB<47>ɂ<EFBFBD><C982><EFBFBD>
std::swap(planes[0], planes[2]);
cv::merge(planes, im);
}
2015-05-29 01:47:26 +09:00
2015-11-19 01:50:11 +09:00
out = im;
2015-05-29 01:47:26 +09:00
2015-11-19 01:50:11 +09:00
return eWaifu2xError_OK;
}
Waifu2x::eWaifu2xError Waifu2x::ReconstructFloatMat(const bool isReconstructNoise, const bool isReconstructScale, const waifu2xCancelFunc cancel_func, const cv::Mat &in, cv::Mat &out)
2015-11-19 01:50:11 +09:00
{
Waifu2x::eWaifu2xError ret;
cv::Mat im(in);
cv::Size_<int> image_size = im.size();
2015-05-29 01:47:26 +09:00
2015-06-03 03:01:56 +09:00
if (isReconstructNoise)
{
PaddingImage(im, im);
2015-05-29 01:47:26 +09:00
2015-06-03 03:01:56 +09:00
ret = ReconstructImage(net_noise, im);
if (ret != eWaifu2xError_OK)
return ret;
2015-05-29 01:47:26 +09:00
2015-06-03 03:01:56 +09:00
// <20>p<EFBFBD>f<EFBFBD>B<EFBFBD><42><EFBFBD>O<EFBFBD><4F><EFBFBD><EFBFBD><EFBFBD><EFBFBD><E895A5>
im = im(cv::Rect(offset, offset, image_size.width, image_size.height));
// <20>l<EFBFBD><6C>0<EFBFBD>`1<>ɃN<C983><4E><EFBFBD>b<EFBFBD>s<EFBFBD><73><EFBFBD>O
cv::threshold(im, im, 1.0, 1.0, cv::THRESH_TRUNC);
cv::threshold(im, im, 0.0, 0.0, cv::THRESH_TOZERO);
2015-06-03 03:01:56 +09:00
}
2015-05-29 01:47:26 +09:00
2015-06-03 03:01:56 +09:00
if (cancel_func && cancel_func())
return eWaifu2xError_Cancel;
2015-05-29 01:47:26 +09:00
const double ratio = CalcScaleRatio(image_size);
const int scale2 = ceil(log2(ratio));
2015-06-03 03:01:56 +09:00
if (isReconstructScale)
{
bool isError = false;
for (int i = 0; i < scale2; i++)
{
2015-06-03 03:01:56 +09:00
Zoom2xAndPaddingImage(im, im, image_size);
2015-06-03 03:01:56 +09:00
ret = ReconstructImage(net_scale, im);
if (ret != eWaifu2xError_OK)
return ret;
2015-06-03 03:01:56 +09:00
// <20>p<EFBFBD>f<EFBFBD>B<EFBFBD><42><EFBFBD>O<EFBFBD><4F><EFBFBD><EFBFBD><EFBFBD><EFBFBD><E895A5>
im = im(cv::Rect(offset, offset, image_size.width, image_size.height));
// <20>l<EFBFBD><6C>0<EFBFBD>`1<>ɃN<C983><4E><EFBFBD>b<EFBFBD>s<EFBFBD><73><EFBFBD>O
cv::threshold(im, im, 1.0, 1.0, cv::THRESH_TRUNC);
cv::threshold(im, im, 0.0, 0.0, cv::THRESH_TOZERO);
2015-06-03 03:01:56 +09:00
}
}
2015-05-29 01:47:26 +09:00
2015-06-03 03:01:56 +09:00
if (cancel_func && cancel_func())
return eWaifu2xError_Cancel;
2015-05-29 01:47:26 +09:00
2015-11-19 01:50:11 +09:00
out = im;
return eWaifu2xError_OK;
}
Waifu2x::eWaifu2xError Waifu2x::Reconstruct(const bool isReconstructNoise, const bool isReconstructScale, const waifu2xCancelFunc cancel_func, const cv::Mat &in, cv::Mat &out)
2015-12-06 02:13:30 +09:00
{
Waifu2x::eWaifu2xError ret;
bool convertBGRflag = false;
2015-12-06 02:13:30 +09:00
cv::Mat brfm;
ret = BeforeReconstructFloatMatProcess(in, brfm, convertBGRflag);
2015-12-06 02:13:30 +09:00
if (ret != eWaifu2xError_OK)
return ret;
cv::Mat reconstruct_image;
if (!use_tta) // <20><><EFBFBD>ʂɏ<CA82><C98F><EFBFBD>
{
ret = ReconstructFloatMat(isReconstructNoise, isReconstructScale, cancel_func, brfm, reconstruct_image);
2015-12-06 02:13:30 +09:00
if (ret != eWaifu2xError_OK)
return ret;
}
else // Test-Time Augmentation Mode
{
const auto RotateClockwise90 = [](cv::Mat &mat)
{
cv::transpose(mat, mat);
cv::flip(mat, mat, 1);
};
const auto RotateClockwise90N = [RotateClockwise90](cv::Mat &mat, const int rotateNum)
{
for (int i = 0; i < rotateNum; i++)
RotateClockwise90(mat);
};
const auto RotateCounterclockwise90 = [](cv::Mat &mat)
{
cv::transpose(mat, mat);
cv::flip(mat, mat, 0);
};
const auto RotateCounterclockwise90N = [RotateCounterclockwise90](cv::Mat &mat, const int rotateNum)
{
for (int i = 0; i < rotateNum; i++)
RotateCounterclockwise90(mat);
};
cv::Mat ri[8];
for (int i = 0; i < 8; i++)
{
cv::Mat in(brfm.clone());
const int rotateNum = i % 4;
RotateClockwise90N(in, rotateNum);
if (i >= 4)
cv::flip(in, in, 1); // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>]
ret = ReconstructFloatMat(isReconstructNoise, isReconstructScale, cancel_func, in, in);
2015-12-06 02:13:30 +09:00
if (ret != eWaifu2xError_OK)
return ret;
if (i >= 4)
cv::flip(in, in, 1); // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>]
RotateCounterclockwise90N(in, rotateNum);
ri[i] = 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); // <20><><EFBFBD>̒n<CC92>_<EFBFBD>ł͂܂<CD82>RGB<47>Ȃ<EFBFBD><C882>Ƃɒ<C682><C992><EFBFBD>
}
// <20>l<EFBFBD><6C>0<EFBFBD>`1<>ɃN<C983><4E><EFBFBD>b<EFBFBD>s<EFBFBD><73><EFBFBD>O
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);
2015-12-06 02:13:30 +09:00
out = reconstruct_image;
return eWaifu2xError_OK;
}
Waifu2x::eWaifu2xError Waifu2x::AfterReconstructFloatMatProcess(const bool isReconstructScale, const waifu2xCancelFunc cancel_func, const cv::Mat &floatim, cv::Mat &in, cv::Mat &out)
2015-11-19 01:50:11 +09:00
{
cv::Size_<int> image_size = in.size();
cv::Mat process_image;
if (input_plane == 1)
{
// <20>č\<5C>z<EFBFBD><7A><EFBFBD><EFBFBD><EFBFBD>P<EFBFBD>x<EFBFBD><EFBFBD><E6919C>CreateZoomColorImage()<29>ō쐬<C58D><EC90AC><EFBFBD><EFBFBD><EFBFBD>F<EFBFBD><46><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>}<7D>[<5B>W<EFBFBD><57><EFBFBD>Ēʏ<C492><CA8F>̉摜<CC89>ɕϊ<C995><CF8A><EFBFBD><EFBFBD>A<EFBFBD><41><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
std::vector<cv::Mat> color_planes;
2015-11-19 01:50:11 +09:00
CreateZoomColorImage(floatim, image_size, color_planes);
2015-11-19 01:50:11 +09:00
color_planes[0] = in;
in.release();
2015-05-29 01:47:26 +09:00
cv::Mat converted_image;
cv::merge(color_planes, converted_image);
color_planes.clear();
cv::cvtColor(converted_image, process_image, ConvertInverseMode);
converted_image.release();
}
else
{
std::vector<cv::Mat> planes;
2015-11-19 01:50:11 +09:00
cv::split(in, planes);
in.release();
// RGB<47><42><EFBFBD><EFBFBD>BGR<47>ɒ<EFBFBD><C992><EFBFBD>
std::swap(planes[0], planes[2]);
cv::merge(planes, process_image);
}
2015-05-29 01:47:26 +09:00
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_<int> 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);
}
}
2015-06-03 03:01:56 +09:00
cv::Mat alpha;
if (floatim.channels() == 4)
2015-06-03 03:01:56 +09:00
{
std::vector<cv::Mat> planes;
2015-11-19 01:50:11 +09:00
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) // <20>g<EFBFBD><67><EFBFBD><EFBFBD><EFBFBD>̃T<CC83>C<EFBFBD>Y<EFBFBD><59>3GB<47><42><EFBFBD><EFBFBD><EFBFBD>Ă<EFBFBD><C482><EFBFBD><EFBFBD><EFBFBD><EFBFBD>t<EFBFBD>@<40>C<EFBFBD><43><EFBFBD>ɏ<EFBFBD><C98F><EFBFBD><EFBFBD>o<EFBFBD><6F><EFBFBD>ă<EFBFBD><C483><EFBFBD><EFBFBD><EFBFBD><EFBFBD>s<EFBFBD><73><EFBFBD>΍<EFBFBD>
Reconstruct(false, true, cancel_func, alpha, alpha);
else
{
boost::filesystem::path temp = boost::filesystem::unique_path("%%%%-%%%%-%%%%-%%%%.bin");
2015-05-29 01:47:26 +09:00
auto compp = [](const cv::Mat &im, const boost::filesystem::path &temp)
{
static char outbuf[10240000];
FILE *fout;
z_stream z;
int count;
2015-05-29 01:47:26 +09:00
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; // <20><><EFBFBD><EFBFBD>
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; // <20>o<EFBFBD>̓o<CD83>b<EFBFBD>t<EFBFBD>@<40>c<EFBFBD>ʂ<EFBFBD><CA82><EFBFBD><EFBFBD>ɖ߂<C996>
z.avail_out = sizeof(outbuf); // <20>o<EFBFBD>̓|<7C>C<EFBFBD><43><EFBFBD>^<5E><><EFBFBD><EFBFBD><EFBFBD>ɖ߂<C996>
}
}
// <20>c<EFBFBD><63><EFBFBD><EFBFBD><EFBFBD>f<EFBFBD><66><EFBFBD>o<EFBFBD><6F>
if ((count = sizeof(outbuf) - z.avail_out) != 0)
{
if (fwrite(outbuf, 1, count, fout) != count)
{
fclose(fout);
return false;
}
}
// <20><><EFBFBD>n<EFBFBD><6E>
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; // <20>o<EFBFBD>̓|<7C>C<EFBFBD><43><EFBFBD>^
z.avail_out = MaxSize; // <20>o<EFBFBD>̓o<CD83>b<EFBFBD>t<EFBFBD>@<40>c<EFBFBD><63>
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());
}
}
2015-05-29 01:47:26 +09:00
}
if (isReconstructScale)
{
const cv::Size_<int> 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);
}
}
2015-05-29 01:47:26 +09:00
// <20>A<EFBFBD><41><EFBFBD>t<EFBFBD>@<40>`<60><><EFBFBD><EFBFBD><EFBFBD>l<EFBFBD><6C><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>A<EFBFBD><41><EFBFBD>t<EFBFBD>@<40><><EFBFBD>t<EFBFBD><74><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
if (!alpha.empty())
{
std::vector<cv::Mat> planes;
cv::split(process_image, planes);
process_image.release();
planes.push_back(alpha);
alpha.release();
cv::merge(planes, process_image);
}
// <20>l<EFBFBD><6C>0<EFBFBD>`1<>ɃN<C983><4E><EFBFBD>b<EFBFBD>s<EFBFBD><73><EFBFBD>O
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);
2015-11-19 01:50:11 +09:00
out = process_image;
return eWaifu2xError_OK;
}
namespace
{
template<typename T>
void AlphaZeroToZero(std::vector<cv::Mat> &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)
2015-11-19 01:50:11 +09:00
{
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";
2015-11-19 01:50:11 +09:00
2015-12-06 02:13:30 +09:00
cv::Mat reconstruct_image;
ret = Reconstruct(isReconstructNoise, isReconstructScale, cancel_func, inMat, reconstruct_image);
2015-11-19 01:50:11 +09:00
if (ret != eWaifu2xError_OK)
return ret;
cv::Mat process_image;
ret = AfterReconstructFloatMatProcess(isReconstructScale, cancel_func, inMat, reconstruct_image, process_image);
2015-11-19 01:50:11 +09:00
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);
2015-06-03 03:01:56 +09:00
cv::Mat write_iamge;
if (output_depth != 32) // <20>o<EFBFBD>͂<EFBFBD>float<61>`<60><><EFBFBD>Ȃ<EFBFBD><C882>ϊ<EFBFBD><CF8A><EFBFBD><EFBFBD>Ȃ<EFBFBD>
process_image.convertTo(write_iamge, cv_depth, max_val, eps);
else
write_iamge = process_image;
2015-06-03 03:01:56 +09:00
process_image.release();
// <20><><EFBFBD>S<EFBFBD><53><EFBFBD><EFBFBD><EFBFBD>̃s<CC83>N<EFBFBD>Z<EFBFBD><5A><EFBFBD>̐F<CC90><46><EFBFBD><EFBFBD><EFBFBD><EFBFBD>(<28><><EFBFBD><EFBFBD><EFBFBD>̓s<CC93><73><EFBFBD><EFBFBD><EFBFBD>A<EFBFBD><41><EFBFBD>S<EFBFBD><53><EFBFBD><EFBFBD><EFBFBD>̃s<CC83>N<EFBFBD>Z<EFBFBD><5A><EFBFBD>ɂ<EFBFBD><C982>F<EFBFBD><46><EFBFBD>t<EFBFBD><74><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>)
// <20><><EFBFBD>f<EFBFBD><66><EFBFBD>ɂ<EFBFBD><C982><EFBFBD><EFBFBD>Ă͉摜<CD89>S<EFBFBD><53><EFBFBD>̊<EFBFBD><CC8A>S<EFBFBD><53><EFBFBD><EFBFBD><EFBFBD>̏ꏊ<CC8F>ɂ<EFBFBD><C982><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>l<EFBFBD>̃A<CC83><41><EFBFBD>t<EFBFBD>@<40><><EFBFBD>L<EFBFBD><4C><EFBFBD><EFBFBD>Ƃ<EFBFBD><C682><EFBFBD><EFBFBD><EFBFBD><EFBFBD>B<EFBFBD><42><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>߂<EFBFBD>cv_depth<74>֕ϊ<D695><CF8A><EFBFBD><EFBFBD>Ă<EFBFBD><C482><EFBFBD>̏<EFBFBD><CC8F><EFBFBD><EFBFBD><EFBFBD><EFBFBD>s<EFBFBD><73><EFBFBD><EFBFBD><EFBFBD>Ƃɂ<C682><C982><EFBFBD>
// (<28><><EFBFBD><EFBFBD><EFBFBD><EFBFBD>cv_depth<74><68>32<33>̏ꍇ<CC8F><EA8D87><EFBFBD>ƈӖ<C688><D396>͖<EFBFBD><CD96><EFBFBD><EFBFBD><EFBFBD>)
// TODO: <20><><EFBFBD>f<EFBFBD><66>(<28><EFBFBD><E182A6>Photo)<29>ɂ<EFBFBD><C982><EFBFBD><EFBFBD>Ă<EFBFBD>0<EFBFBD><30><EFBFBD><EFBFBD><EFBFBD>Ȃ<EFBFBD><C882><EFBFBD><E6919C><EFBFBD>ϊ<EFBFBD><CF8A><EFBFBD><EFBFBD>Ă<EFBFBD>0.000114856390<EFBFBD>Ƃ<EFBFBD><EFBFBD>ɂȂ<EFBFBD><EFBFBD>̂ŁA<EFBFBD>K<EFBFBD>؂Ȓl<EFBFBD>̃N<EFBFBD><EFBFBD><EFBFBD>b<EFBFBD>s<EFBFBD><EFBFBD><EFBFBD>O<EFBFBD><EFBFBD><EFBFBD>s<EFBFBD><EFBFBD><EFBFBD>H
if (write_iamge.channels() > 3)
{
std::vector<cv::Mat> planes;
cv::split(write_iamge, planes);
write_iamge.release();
const auto depth = planes[0].depth();
switch (depth)
{
case CV_8U:
AlphaZeroToZero<uint8_t>(planes);
break;
case CV_16U:
AlphaZeroToZero<uint16_t>(planes);
break;
case CV_32F:
AlphaZeroToZero<float>(planes);
break;
case CV_64F:
AlphaZeroToZero<double>(planes);
break;
default:
return eWaifu2xError_FailedUnknownType;
}
cv::merge(planes, write_iamge);
}
outMat = write_iamge;
return eWaifu2xError_OK;
}
double Waifu2x::CalcScaleRatio(const cv::Size_<int> &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);
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;
2016-02-04 14:14:40 +09:00
ret = WriteMat(write_iamge, output_file, output_quality);
2015-07-10 04:09:22 +09:00
if (ret != eWaifu2xError_OK)
return ret;
2015-06-03 03:01:56 +09:00
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) // <20>o<EFBFBD>͐[<5B>x<EFBFBD><78>8bit<69><74><EFBFBD><EFBFBD>
return eWaifu2xError_InvalidParameter;
cv::Mat float_image;
// Mat<61>֕ϊ<D695>
{
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; // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD>0.0<EFBFBD>`1.0<EFBFBD>̂͂<EFBFBD><EFBFBD>Ȃ̂ŕϊ<EFBFBD><EFBFBD>͕K<EFBFBD>v<EFBFBD>Ȃ<EFBFBD>
break;
}
original_image.release();
if (convert.channels() == 1)
cv::cvtColor(convert, convert, cv::COLOR_GRAY2BGR);
else if (convert.channels() == 4)
{
// <20>A<EFBFBD><41><EFBFBD>t<EFBFBD>@<40>`<60><><EFBFBD><EFBFBD><EFBFBD>l<EFBFBD><6C><EFBFBD>t<EFBFBD><74><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><E793A7><EFBFBD>ȃs<C883>N<EFBFBD>Z<EFBFBD><5A><EFBFBD>̂ƕs<C695><73><EFBFBD><EFBFBD><EFBFBD>ȃs<C883>N<EFBFBD>Z<EFBFBD><5A><EFBFBD>̋<EFBFBD><CC8B>E<EFBFBD><45><EFBFBD><EFBFBD><EFBFBD>̐F<CC90><46><EFBFBD>L<EFBFBD><4C><EFBFBD><EFBFBD>
std::vector<cv::Mat> planes;
cv::split(convert, planes);
cv::Mat alpha = planes[3];
planes.resize(3);
AlphaMakeBorder(planes, alpha, layer_num);
planes.push_back(alpha);
cv::merge(planes, convert);
}
float_image = convert;
}
const auto 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();
// <20>o<EFBFBD>͔z<CD94><7A><EFBFBD>֏<EFBFBD><D68F><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
{
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);
}
2015-05-29 01:47:26 +09:00
return eWaifu2xError_OK;
}
2015-06-03 03:01:56 +09:00
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;
}
// <20>s<EFBFBD><73><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ǂƂ肠<C682><E882A0><EFBFBD><EFBFBD>CV_8U<38><55><EFBFBD>Ԃ<EFBFBD><D482>Ă<EFBFBD><C482><EFBFBD>
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;
}
// <20>s<EFBFBD><73><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ǂƂ肠<C682><E882A0><EFBFBD><EFBFBD>255.0<EFBFBD><EFBFBD><EFBFBD>Ԃ<EFBFBD><EFBFBD>Ă<EFBFBD><EFBFBD><EFBFBD>
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;
}
// <20>s<EFBFBD><73><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ǂƂ肠<C682><E882A0><EFBFBD><EFBFBD>clip_eps8<73>Ԃ<EFBFBD><D482>Ă<EFBFBD><C482><EFBFBD>
return clip_eps8;
}