waifu2x-caffe/common/waifu2x.h

#pragma once

#include <stdint.h>
#include <string>
#include <vector>
#include <utility>
#include <functional>
#include <boost/shared_ptr.hpp>
#include <opencv2/opencv.hpp>


namespace caffe
{
	template <typename Dtype>
	class Net;
	class NetParameter;
};

class Waifu2x
{
public:
	enum eWaifu2xError
	{
		eWaifu2xError_OK = 0,
		eWaifu2xError_Cancel,
		eWaifu2xError_NotInitialized,
		eWaifu2xError_InvalidParameter,
		eWaifu2xError_FailedOpenInputFile,
		eWaifu2xError_FailedOpenOutputFile,
		eWaifu2xError_FailedOpenModelFile,
		eWaifu2xError_FailedParseModelFile,
		eWaifu2xError_FailedConstructModel,
		eWaifu2xError_FailedProcessCaffe,
		eWaifu2xError_FailedCudaCheck,
	};

	enum eWaifu2xCudaError
	{
		eWaifu2xCudaError_OK = 0,
		eWaifu2xCudaError_NotFind,
		eWaifu2xCudaError_OldVersion,
	};

	enum eWaifu2xcuDNNError
	{
		eWaifu2xcuDNNError_OK = 0,
		eWaifu2xcuDNNError_NotFind,
		eWaifu2xcuDNNError_OldVersion,
		eWaifu2xcuDNNError_CannotCreate,
	};

	typedef std::function<bool()> waifu2xCancelFunc;

private:
	bool is_inited;

	// <20><><EFBFBD>x<EFBFBD>ɏ<EFBFBD><C98F><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>摜<EFBFBD>̕<EFBFBD>
	int crop_size;
	// <20><><EFBFBD>x<EFBFBD>ɉ<EFBFBD><C989>u<EFBFBD><75><EFBFBD>b<EFBFBD>N<EFBFBD><4E><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>邩
	int batch_size;

	// <20>l<EFBFBD>b<EFBFBD>g<EFBFBD>ɓ<EFBFBD><C993>͂<EFBFBD><CD82><EFBFBD><EFBFBD>摜<EFBFBD>̃T<CC83>C<EFBFBD>Y
	int input_block_size;
	// <20>u<EFBFBD><75><EFBFBD>b<EFBFBD>N<EFBFBD>ϊ<EFBFBD><CF8A><EFBFBD><EFBFBD>̏o<CC8F>̓T<CD83>C<EFBFBD>Y
	int output_size;
	// <20>l<EFBFBD>b<EFBFBD>g<EFBFBD><67><EFBFBD>[<5B>N<EFBFBD>ɓ<EFBFBD><C993>͂<EFBFBD><CD82><EFBFBD><EFBFBD>摜<EFBFBD>̃T<CC83>C<EFBFBD>Y(<28>o<EFBFBD>͉摜<CD89>̕<EFBFBD><CC95><EFBFBD>layer_num * 2<><32><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ȃ<EFBFBD>)
	int block_width_height;
	// srcnn.prototxt<78>Œ<EFBFBD><C592>`<60><><EFBFBD>ꂽ<EFBFBD><EA82BD><EFBFBD>͂<EFBFBD><CD82><EFBFBD><EFBFBD>摜<EFBFBD>̃T<CC83>C<EFBFBD>Y
	int original_width_height;

	std::string mode;
	int noise_level;
	double scale_ratio;
	std::string model_dir;
	std::string process;

	int inner_padding;
	int outer_padding;

	int output_block_size;

	int input_plane;

	bool isCuda;

	boost::shared_ptr<caffe::Net<float>> net_noise;
	boost::shared_ptr<caffe::Net<float>> net_scale;

	float *input_block;
	float *dummy_data;
	float *output_block;

	bool use_tta;

private:
	static eWaifu2xError LoadMat(cv::Mat &float_image, const std::string &input_file);
	static eWaifu2xError LoadMatBySTBI(cv::Mat &float_image, const std::string &input_file);
	eWaifu2xError CreateBrightnessImage(const cv::Mat &float_image, cv::Mat &im);
	eWaifu2xError PaddingImage(const cv::Mat &input, cv::Mat &output);
	eWaifu2xError Zoom2xAndPaddingImage(const cv::Mat &input, cv::Mat &output, cv::Size_<int> &zoom_size);
	eWaifu2xError CreateZoomColorImage(const cv::Mat &float_image, const cv::Size_<int> &zoom_size, std::vector<cv::Mat> &cubic_planes);
	eWaifu2xError ConstractNet(boost::shared_ptr<caffe::Net<float>> &net, const std::string &model_path, const std::string &param_path, const std::string &process);
	eWaifu2xError LoadParameterFromJson(boost::shared_ptr<caffe::Net<float>> &net, const std::string &model_path, const std::string &param_path);
	eWaifu2xError SetParameter(caffe::NetParameter &param) const;
	eWaifu2xError ReconstructImage(boost::shared_ptr<caffe::Net<float>> net, cv::Mat &im);
	eWaifu2xError WriteMat(const cv::Mat &im, const std::string &output_file);

	eWaifu2xError BeforeReconstructFloatMatProcess(const cv::Mat &in, cv::Mat &out);
	eWaifu2xError ReconstructFloatMat(const bool isJpeg, const waifu2xCancelFunc cancel_func, const cv::Mat &in, cv::Mat &out);
	eWaifu2xError AfterReconstructFloatMatProcess(const cv::Mat &floatim, const cv::Mat &in, cv::Mat &out);

public:
	Waifu2x();
	~Waifu2x();

	static eWaifu2xcuDNNError can_use_cuDNN();
	static eWaifu2xCudaError can_use_CUDA();

	// mode: noise or scale or noise_scale or auto_scale
	// process: cpu or gpu or cudnn
	eWaifu2xError init(int argc, char** argv, const std::string &mode, const int noise_level, const double scale_ratio, const std::string &model_dir, const std::string &process,
		const bool use_tta = false, const int crop_size = 128, const int batch_size = 1);

	void destroy();

	eWaifu2xError waifu2x(const std::string &input_file, const std::string &output_file,
		const waifu2xCancelFunc cancel_func = nullptr);

	const std::string& used_process() const;

	static cv::Mat LoadMat(const std::string &path);
};