コード例 #1
0
ファイル: waifu2x.cpp プロジェクト: nevercannot/waifu2x-caffe
Waifu2x::eWaifu2xError Waifu2x::waifu2x(int factor, const uint32_t* source, uint32_t* dest, int width, int height)
{
	Waifu2x::eWaifu2xError ret;

	if (!is_inited)
		return eWaifu2xError_NotInitialized;

	cv::Mat float_image;
	ret = LoadMat(float_image, source, width, height);
	if (ret != eWaifu2xError_OK)
		return ret;

	cv::Mat im;
	if (input_plane == 1)
		return eWaifu2xError_NotInitialized;
	else
	{
		std::vector<cv::Mat> planes;
		cv::split(float_image, planes);

		if (float_image.channels() == 4)
			planes.resize(3);

		// BGRからRGBにする
		//std::swap(planes[0], planes[2]);

		cv::merge(planes, im);
	}
	cv::Size_<int> image_size = im.size();

	const bool isReconstructNoise = mode == "noise" || mode == "noise_scale" || mode == "auto_scale";
	const bool isReconstructScale = mode == "scale" || mode == "noise_scale";

	if (isReconstructNoise)
	{
		PaddingImage(im, im);

		ret = ReconstructImage(net_noise, im);
		if (ret != eWaifu2xError_OK)
			return ret;

		// パディングを取り払う
		im = im(cv::Rect(offset, offset, image_size.width, image_size.height));
	}

	const int scale2 = ceil(log2((double)factor));
	const double shrinkRatio = (double)factor / std::pow(2.0, (double)scale2);

	if (isReconstructScale)
	{
		bool isError = false;
		for (int i = 0; i < scale2; i++)
		{
			Zoom2xAndPaddingImage(im, im, image_size);

			ret = ReconstructImage(net_scale, im);
			if (ret != eWaifu2xError_OK)
				return ret;

			// パディングを取り払う
			im = im(cv::Rect(offset, offset, image_size.width, image_size.height));
		}
	}

	cv::Mat process_image;
	if (input_plane == 1)
	{
		// 再構築した輝度画像とCreateZoomColorImage()で作成した色情報をマージして通常の画像に変換し、書き込む

		std::vector<cv::Mat> color_planes;
		CreateZoomColorImage(float_image, image_size, color_planes);

		float_image.release();

		color_planes[0] = im;
		im.release();

		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;
		cv::split(im, planes);

		// RGBからBGRに直す
		//std::swap(planes[0], planes[2]);

		cv::merge(planes, process_image);
	}

	cv::Mat alpha;
	if (float_image.channels() == 4)
	{
		std::vector<cv::Mat> planes;
		cv::split(float_image, planes);
		alpha = planes[3];

		cv::resize(alpha, alpha, image_size, 0.0, 0.0, cv::INTER_CUBIC);
	}

	// アルファチャンネルがあったら、アルファを付加してカラーからアルファの影響を抜く
	if (!alpha.empty())
	{
		std::vector<cv::Mat> planes;
		cv::split(process_image, planes);
		process_image.release();

		planes.push_back(alpha);

		cv::Mat w2 = planes[3];

		planes[0] = (planes[0]).mul(1.0 / w2);
		planes[1] = (planes[1]).mul(1.0 / w2);
		planes[2] = (planes[2]).mul(1.0 / w2);

		cv::merge(planes, process_image);
	}

	const cv::Size_<int> ns(image_size.width * shrinkRatio, image_size.height * shrinkRatio);
	if (image_size.width != ns.width || image_size.height != ns.height)
		cv::resize(process_image, process_image, ns, 0.0, 0.0, cv::INTER_LINEAR);

	cv::Mat write_iamge;
	process_image.convertTo(write_iamge, CV_8U, 255.0);
	process_image.release();

	/*
	ret = WriteMat(write_iamge, output_file);
	if (ret != eWaifu2xError_OK)
	return ret;

	write_iamge.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(dest + width * i, write_iamge.data + stride * i, stride);
	}

	return eWaifu2xError_OK;
}
コード例 #2
0
ファイル: waifu2x.cpp プロジェクト: soeaver/waifu2x-caffe
Waifu2x::eWaifu2xError Waifu2x::AfterReconstructFloatMatProcess(const bool isReconstructScale, const waifu2xCancelFunc cancel_func, const cv::Mat &floatim, const cv::Mat &in, cv::Mat &out)
{
	cv::Size_<int> image_size = in.size();

	cv::Mat process_image;
	if (input_plane == 1)
	{
		// 再構築した輝度画像とCreateZoomColorImage()で作成した色情報をマージして通常の画像に変換し、書き込む

		std::vector<cv::Mat> color_planes;
		CreateZoomColorImage(floatim, image_size, color_planes);

		color_planes[0] = in;

		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;
		cv::split(in, planes);

		// RGBからBGRに直す
		std::swap(planes[0], planes[2]);

		cv::merge(planes, process_image);
	}

	const int scale2 = ceil(log2(scale_ratio));
	const double shrinkRatio = scale_ratio / std::pow(2.0, (double)scale2);

	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);
		if (image_size.width != ns.width || image_size.height != ns.height)
			cv::resize(process_image, process_image, ns, 0.0, 0.0, cv::INTER_LINEAR);
	}

	// 値を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;
}