mirror of
https://github.com/lltcggie/waifu2x-caffe.git
synced 2025-06-26 05:32:47 +00:00
・メモリ使用量を減らすようにした
・拡大後のサイズが3GB超えていたらデータを一時ファイルに書き出してメモリ不足対策をするようにした
This commit is contained in:
lltcggie
parent
77b84972b3
commit
2d395ab3b1
@ -16,6 +16,7 @@
|
||||
#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
|
||||
@ -654,10 +655,9 @@ Waifu2x::eWaifu2xError Waifu2x::Zoom2xAndPaddingImage(const cv::Mat &input, cv::
|
||||
zoom_size.width *= 2;
|
||||
zoom_size.height *= 2;
|
||||
|
||||
cv::Mat zoom_image;
|
||||
cv::resize(input, zoom_image, zoom_size, 0.0, 0.0, cv::INTER_NEAREST);
|
||||
cv::resize(input, output, zoom_size, 0.0, 0.0, cv::INTER_NEAREST);
|
||||
|
||||
return PaddingImage(zoom_image, output);
|
||||
return PaddingImage(output, output);
|
||||
}
|
||||
|
||||
// 入力画像をzoom_sizeの大きさにcv::INTER_CUBICで拡大し、色情報のみを残す
|
||||
@ -1618,7 +1618,7 @@ Waifu2x::eWaifu2xError Waifu2x::Reconstruct(const bool isReconstructNoise, const
|
||||
return eWaifu2xError_OK;
|
||||
}
|
||||
|
||||
Waifu2x::eWaifu2xError Waifu2x::AfterReconstructFloatMatProcess(const bool isReconstructScale, const waifu2xCancelFunc cancel_func, const cv::Mat &floatim, const cv::Mat &in, cv::Mat &out)
|
||||
Waifu2x::eWaifu2xError Waifu2x::AfterReconstructFloatMatProcess(const bool isReconstructScale, const waifu2xCancelFunc cancel_func, const cv::Mat &floatim, cv::Mat &in, cv::Mat &out)
|
||||
{
|
||||
cv::Size_<int> image_size = in.size();
|
||||
|
||||
@ -1631,6 +1631,7 @@ Waifu2x::eWaifu2xError Waifu2x::AfterReconstructFloatMatProcess(const bool isRec
|
||||
CreateZoomColorImage(floatim, image_size, color_planes);
|
||||
|
||||
color_planes[0] = in;
|
||||
in.release();
|
||||
|
||||
cv::Mat converted_image;
|
||||
cv::merge(color_planes, converted_image);
|
||||
@ -1643,6 +1644,7 @@ Waifu2x::eWaifu2xError Waifu2x::AfterReconstructFloatMatProcess(const bool isRec
|
||||
{
|
||||
std::vector<cv::Mat> planes;
|
||||
cv::split(in, planes);
|
||||
in.release();
|
||||
|
||||
// RGBからBGRに直す
|
||||
std::swap(planes[0], planes[2]);
|
||||
@ -1654,30 +1656,6 @@ Waifu2x::eWaifu2xError Waifu2x::AfterReconstructFloatMatProcess(const bool isRec
|
||||
const int scale2 = ceil(log2(ratio));
|
||||
const double shrinkRatio = ratio >= 1.0 ? ratio / std::pow(2.0, (double)scale2) : ratio;
|
||||
|
||||
cv::Mat alpha;
|
||||
if (floatim.channels() == 4)
|
||||
{
|
||||
std::vector<cv::Mat> planes;
|
||||
cv::split(floatim, planes);
|
||||
|
||||
if (isReconstructScale)
|
||||
Reconstruct(false, true, cancel_func, planes[3], alpha);
|
||||
else
|
||||
alpha = planes[3];
|
||||
}
|
||||
|
||||
// アルファチャンネルがあったらアルファを付加する
|
||||
if (!alpha.empty())
|
||||
{
|
||||
std::vector<cv::Mat> planes;
|
||||
cv::split(process_image, planes);
|
||||
process_image.release();
|
||||
|
||||
planes.push_back(alpha);
|
||||
|
||||
cv::merge(planes, process_image);
|
||||
}
|
||||
|
||||
if (isReconstructScale)
|
||||
{
|
||||
const cv::Size_<int> ns(image_size.width * shrinkRatio, image_size.height * shrinkRatio);
|
||||
@ -1691,6 +1669,202 @@ Waifu2x::eWaifu2xError Waifu2x::AfterReconstructFloatMatProcess(const bool isRec
|
||||
}
|
||||
}
|
||||
|
||||
cv::Mat alpha;
|
||||
if (floatim.channels() == 4)
|
||||
{
|
||||
std::vector<cv::Mat> 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_<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);
|
||||
}
|
||||
}
|
||||
|
||||
// アルファチャンネルがあったらアルファを付加する
|
||||
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);
|
||||
}
|
||||
|
||||
// 値を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);
|
||||
|
||||
@ -133,7 +133,7 @@ private:
|
||||
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, 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);
|
||||
|
||||
eWaifu2xError waifu2xConvetedMat(const bool isJpeg, const cv::Mat &inMat, cv::Mat &outMat,
|
||||
const waifu2xCancelFunc cancel_func = nullptr);
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user