データを直接与えるAPIがうまく動かなかったのを修正

This commit is contained in:
lltcggie 2016-07-08 23:35:51 +09:00
parent 5514995215
commit b880441b02
2 changed files with 27 additions and 1 deletions

View File

@ -380,6 +380,8 @@ Waifu2x::eWaifu2xError stImage::Load(const boost::filesystem::path &input_file)
Waifu2x::eWaifu2xError stImage::Load(const void* source, const int width, const int height, const int channel, const int stride) Waifu2x::eWaifu2xError stImage::Load(const void* source, const int width, const int height, const int channel, const int stride)
{ {
Clear();
cv::Mat original_image(cv::Size(width, height), CV_MAKETYPE(CV_8U, channel), (void *)source, stride); cv::Mat original_image(cv::Size(width, height), CV_MAKETYPE(CV_8U, channel), (void *)source, stride);
if (original_image.channels() >= 3) // RGBなのでBGRにする if (original_image.channels() >= 3) // RGBなのでBGRにする
@ -393,6 +395,10 @@ Waifu2x::eWaifu2xError stImage::Load(const void* source, const int width, const
} }
mOrgFloatImage = original_image; mOrgFloatImage = original_image;
mOrgChannel = original_image.channels();
mOrgSize = original_image.size();
mIsRequestDenoise = false;
return Waifu2x::eWaifu2xError_OK; return Waifu2x::eWaifu2xError_OK;
} }

View File

@ -5,6 +5,7 @@
#include <cudnn.h> #include <cudnn.h>
#include <mutex> #include <mutex>
#include <opencv2/core.hpp> #include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <tclap/CmdLine.h> #include <tclap/CmdLine.h>
#include <boost/filesystem.hpp> #include <boost/filesystem.hpp>
#include <boost/algorithm/string.hpp> #include <boost/algorithm/string.hpp>
@ -790,6 +791,18 @@ Waifu2x::eWaifu2xError Waifu2x::waifu2x(const double factor, const void* source,
if (!mIsInited) if (!mIsInited)
return Waifu2x::eWaifu2xError_NotInitialized; return Waifu2x::eWaifu2xError_NotInitialized;
int cvrSetting = -1;
if (in_channel == 3 && out_channel == 3)
cvrSetting = CV_BGR2RGB;
else if (in_channel == 4 && out_channel == 4)
cvrSetting = CV_BGRA2RGBA;
else if (in_channel == 3 && out_channel == 4)
cvrSetting = CV_BGR2RGBA;
else if (in_channel == 4 && out_channel == 3)
cvrSetting = CV_BGRA2RGB;
else if (!(in_channel == 1 && out_channel == 1))
return Waifu2x::eWaifu2xError_InvalidParameter;
stImage image; stImage image;
ret = image.Load(source, width, height, in_channel, in_stride); ret = image.Load(source, width, height, in_channel, in_stride);
if (ret != Waifu2x::eWaifu2xError_OK) if (ret != Waifu2x::eWaifu2xError_OK)
@ -811,9 +824,16 @@ Waifu2x::eWaifu2xError Waifu2x::waifu2x(const double factor, const void* source,
image.Postprocess(mInputPlane, Factor, 8); image.Postprocess(mInputPlane, Factor, 8);
cv::Mat out_image = image.GetEndImage(); cv::Mat out_bgr_image = image.GetEndImage();
image.Clear(); image.Clear();
cv::Mat out_image;
if (cvrSetting >= 0)
cv::cvtColor(out_bgr_image, out_image, cvrSetting); // BGR©çRGBÉß·
else
out_image = out_bgr_image;
out_bgr_image.release();
// 出力配列へ書き込み // 出力配列へ書き込み
{ {
const auto width = out_image.size().width; const auto width = out_image.size().width;