// ネットワークを使って画像を再構築する
Waifu2x::eWaifu2xError cNet::ReconstructImage(const bool UseTTA, const int crop_w, const int crop_h, const int outer_padding, const int batch_size, float *outputBlockBuf, const cv::Mat &inMat, cv::Mat &outMat)
{
	const auto InputHeight = inMat.size().height;
	const auto InputWidth = inMat.size().width;
	const auto InputLine = inMat.step1();

	assert(inMat.channels() == 1 || inMat.channels() == 3);

	const int InputPadding = mNetOffset + outer_padding; // 入力パディング

	const auto NoPaddingInputWidth = InputWidth - InputPadding * 2; // パディングを除いた入力画像サイズ(横)
	const auto NoPaddingInputHeight = InputHeight - InputPadding * 2; // パディングを除いた入力画像サイズ(縦)

	cv::Mat outim(NoPaddingInputHeight * mInnerScale, NoPaddingInputWidth * mInnerScale, inMat.type());

	// float *imptr = (float *)im.data;
	float *imptr = (float *)outim.data;

	const auto input_block_width = crop_w + InputPadding * 2; // 入力ブロックサイズ(横)
	const auto input_block_height = crop_h + InputPadding * 2; // 入力ブロックサイズ(縦)

	const auto output_block_width = input_block_width * mInnerScale - mNetOffset * 2; // 出力ブロックサイズ(横)
	const auto output_block_height = input_block_height * mInnerScale - mNetOffset * 2; // 出力ブロックサイズ(縦)

	const auto output_crop_block_width = crop_w * mInnerScale; // クロップ後の出力ブロックサイズ(横)
	const auto output_crop_block_height = crop_h * mInnerScale; // クロップ後の出力ブロックサイズ(縦)

	const auto output_crop_w = (output_block_width - crop_w * mInnerScale) / 2; // 出力後のクロップサイズ
	const auto output_crop_h = (output_block_height - crop_h * mInnerScale) / 2; // 出力後のクロップサイズ

	assert(NoPaddingInputWidth % crop_w == 0);
	assert(NoPaddingInputHeight % crop_h == 0);

	try
	{
		auto input_blobs = mNet->input_blobs();

		assert(input_blobs.size() > 0);

		auto input_blob = mNet->input_blobs()[0];

		input_blob->Reshape(batch_size, mInputPlane, input_block_height, input_block_width);

		assert(inMat.channels() == mInputPlane);
		assert(input_blob->shape(1) == mInputPlane);

		const int WidthNum = NoPaddingInputWidth / crop_w;
		const int HeightNum = NoPaddingInputHeight / crop_h;

		const int BlockNum = WidthNum * HeightNum;

		const int input_block_plane_size = input_block_width * input_block_height * mInputPlane;
		const int output_block_plane_size = output_block_width * output_block_height * mInputPlane;

		// 画像は(消費メモリの都合上)block_size*block_sizeに分けて再構築する
		for (int num = 0; num < BlockNum; num += batch_size)
		{
			const int processNum = (BlockNum - num) >= batch_size ? batch_size : BlockNum - num;

			if (processNum < batch_size)
				input_blob->Reshape(processNum, mInputPlane, input_block_height, input_block_width);

			for (int n = 0; n < processNum; n++)
			{
				const int wn = (num + n) % WidthNum;
				const int hn = (num + n) / WidthNum;

				const int w = wn * crop_w;
				const int h = hn * crop_h;

				assert(w + input_block_width <= InputWidth && h + input_block_height <= InputHeight);

				cv::Mat someimg = inMat(cv::Rect(w, h, input_block_width, input_block_height));

				// 画像を直列に変換
				{
					float *fptr = input_blob->mutable_cpu_data() + (input_block_plane_size * n);
					const float *uptr = (const float *)someimg.data;

					const auto Line = someimg.step1();

					if (someimg.channels() == 1)
					{
						if (input_block_width == Line)
							memcpy(fptr, uptr, input_block_width * input_block_height * sizeof(float));
						else
						{
							for (int i = 0; i < input_block_height; i++)
								memcpy(fptr + i * input_block_width, uptr + i * Line, input_block_width * sizeof(float));
						}
					}
					else
					{
						const auto LinePixel = someimg.step1() / someimg.channels();
						const auto Channel = someimg.channels();
						const auto Width = someimg.size().width;
						const auto Height = someimg.size().height;

						for (int i = 0; i < Height; i++)
						{
							for (int j = 0; j < Width; j++)
							{
								for (int ch = 0; ch < Channel; ch++)
								{
									const size_t IndexSrc = i * someimg.step1() + j * Channel + ch;
									const size_t IndexDst = (ch * Height + i) * Width + j;
									fptr[IndexDst] = uptr[IndexSrc];
								}
							}
						}
					}
				}
			}

			assert(input_blob->count() == input_block_plane_size * processNum);

			// 計算
			auto out = mNet->Forward();

			auto b = out[0];

			assert(b->count() == output_block_plane_size * processNum);

			const float *ptr = nullptr;

			if (caffe::Caffe::mode() == caffe::Caffe::CPU)
				ptr = b->cpu_data();
			else
				ptr = b->gpu_data();

			caffe::caffe_copy(output_block_plane_size * processNum, ptr, outputBlockBuf);

			for (int n = 0; n < processNum; n++)
			{
				const int wn = (num + n) % WidthNum;
				const int hn = (num + n) / WidthNum;

				const int w = wn * output_crop_block_width;
				const int h = hn * output_crop_block_height;

				const float *fptr = outputBlockBuf + (output_block_plane_size * n);

				const auto Line = outim.step1();

				// 結果を出力画像にコピー
				if (outim.channels() == 1)
				{
					for (int i = 0; i < output_crop_block_height; i++)
						memcpy(imptr + (h + i) * Line + w, fptr + (i + output_crop_h) * output_block_width + output_crop_w, output_crop_block_width * sizeof(float));
				}
				else
				{
					const auto LinePixel = Line / outim.channels();
					const auto Channel = outim.channels();

					for (int i = 0; i < output_crop_block_height; i++)
					{
						for (int j = 0; j < output_crop_block_width; j++)
						{
							for (int ch = 0; ch < Channel; ch++)
							{
								const size_t IndexSrc = (ch * output_block_height + i + output_crop_h) * output_block_width + j + output_crop_w;
								const size_t IndexDst = ((h + i) * LinePixel + (w + j)) * Channel + ch;

								imptr[IndexDst] = fptr[IndexSrc];
							}
						}
					}
				}

				//{
				//	cv::Mat testim(output_block_size, output_block_size, CV_32FC1);
				//	float *p = (float *)testim.data;
				//	for (int i = 0; i < output_block_size; i++)
				//	{
				//		for (int j = 0; j < output_block_size; j++)
				//		{
				//			p[testim.step1() * i + j] = fptr[i * output_block_size + j];
				//		}
				//	}

				//	const int cv_depth = DepthBitToCVDepth(8);
				//	const double max_val = GetValumeMaxFromCVDepth(cv_depth);
				//	const double eps = GetEPS(cv_depth);

				//	cv::Mat write_iamge;
				//	testim.convertTo(write_iamge, cv_depth, max_val, eps);

				//	cv::imwrite("ti.png", write_iamge);
				//	testim.release();
				//}
			}
		}
	}
	catch (...)
	{
		return Waifu2x::eWaifu2xError_FailedProcessCaffe;
	}

	// 値を0~1にクリッピング
	cv::threshold(outim, outim, 1.0, 1.0, cv::THRESH_TRUNC);
	cv::threshold(outim, outim, 0.0, 0.0, cv::THRESH_TOZERO);

	outMat = outim;

	return Waifu2x::eWaifu2xError_OK;
}
Ejemplo n.º 2
0
// ネットワークを使って画像を再構築する
Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr<caffe::Net<float>> net, cv::Mat &im)
{
	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);

	cv::Mat outim(im.rows, im.cols, im.type());

	// float *imptr = (float *)im.data;
	float *imptr = (float *)outim.data;

	try
	{
		auto input_blobs = net->input_blobs();
		auto input_blob = net->input_blobs()[0];

		input_blob->Reshape(batch_size, input_plane, input_block_size, input_block_size);

		assert(im.channels() == input_plane);
		assert(input_blob->shape(1) == input_plane);

		const int WidthNum = Width / output_size;
		const int HeightNum = Height / output_size;

		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;

		// 画像は(消費メモリの都合上)output_size*output_sizeに分けて再構築する
		for (int num = 0; num < BlockNum; num += batch_size)
		{
			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++)
			{
				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)
				{
					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)
					{
						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;
					}

					if (y + height > Height)
					{
						bottom += (y + height) - Height;
						height = Height - y;
					}

					cv::Mat someimg = im(cv::Rect(x, y, width, height));

					cv::Mat someborderimg;
					// 画像を中央にパディング。余白はcv::BORDER_REPLICATEで埋める
					// 実はimで画素が存在する部分は余白と認識されないが、inner_paddingがlayer_numでouter_paddingが1以上ならそこの部分の画素は結果画像として取り出す部分には影響しない
					cv::copyMakeBorder(someimg, someborderimg, top, bottom, left, right, cv::BORDER_REPLICATE);
					someimg.release();

					// 画像を直列に変換
					{
						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
						{
							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];
								}
							}

							/*
							{
								cv::Mat im(someborderimg.size(), CV_32F, fptr, Width * sizeof(float));

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

								if (!cv::imwrite("test_in.png", write_iamge))
									return eWaifu2xError_FailedOpenOutputFile;
							}
							*/
						}
					}
				}
			}

			assert(input_blob->count() == input_block_plane_size * processNum);

			// ネットワークに画像を入力
			input_blob->set_cpu_data(input_block);

			// 計算
			auto out = net->ForwardPrefilled(nullptr);

			auto b = out[0];

			assert(b->count() == output_block_plane_size * processNum);

			const float *ptr = nullptr;

			if (caffe::Caffe::mode() == caffe::Caffe::CPU)
				ptr = b->cpu_data();
			else
				ptr = b->gpu_data();

			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);

				// 結果を出力画像にコピー
				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];
						}
					}

					/*
					{
						cv::Mat im(someborderimg.size(), CV_32F, fptr, Width * sizeof(float));

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

						if (!cv::imwrite("test_in.png", write_iamge))
							return eWaifu2xError_FailedOpenOutputFile;
					}
					*/
				}
			}
		}
	}
	catch (...)
	{
		return eWaifu2xError_FailedProcessCaffe;
	}

	im = outim;

	return eWaifu2xError_OK;
}