diff --git a/common/stImage.cpp b/common/stImage.cpp index 530b301..8cc55b5 100644 --- a/common/stImage.cpp +++ b/common/stImage.cpp @@ -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) { + Clear(); + cv::Mat original_image(cv::Size(width, height), CV_MAKETYPE(CV_8U, channel), (void *)source, stride); 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; + mOrgChannel = original_image.channels(); + mOrgSize = original_image.size(); + + mIsRequestDenoise = false; return Waifu2x::eWaifu2xError_OK; } diff --git a/common/waifu2x.cpp b/common/waifu2x.cpp index 6325e5d..6e22a1c 100644 --- a/common/waifu2x.cpp +++ b/common/waifu2x.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include #include @@ -790,6 +791,18 @@ Waifu2x::eWaifu2xError Waifu2x::waifu2x(const double factor, const void* source, if (!mIsInited) 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; ret = image.Load(source, width, height, in_channel, in_stride); if (ret != Waifu2x::eWaifu2xError_OK) @@ -811,9 +824,16 @@ Waifu2x::eWaifu2xError Waifu2x::waifu2x(const double factor, const void* source, image.Postprocess(mInputPlane, Factor, 8); - cv::Mat out_image = image.GetEndImage(); + cv::Mat out_bgr_image = image.GetEndImage(); 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;