Waifu2x::eWaifu2xError Waifu2x::Init(const eWaifu2xModelType mode, const int noise_level,
const boost::filesystem::path &model_dir, const std::string &process, const int GPUNo)
{
Waifu2x::eWaifu2xError ret;
if (mIsInited)
return Waifu2x::eWaifu2xError_OK;
try
{
std::string Process = process;
const auto cuDNNCheckStartTime = std::chrono::system_clock::now();
if (Process == "gpu" || Process == "cudnn")
{
if (can_use_CUDA() != eWaifu2xCudaError_OK)
return Waifu2x::eWaifu2xError_FailedCudaCheck;
}
mMode = mode;
mNoiseLevel = noise_level;
mProcess = Process;
mGPUNo = GPUNo;
const auto cuDNNCheckEndTime = std::chrono::system_clock::now();
if (Process == "cudnn")
{
// exeのディレクトリにcuDNNのアルゴリズムデータ保存
boost::filesystem::path cudnn_data_base_dir_path(ExeDir);
if (cudnn_data_base_dir_path.is_relative())
cudnn_data_base_dir_path = boost::filesystem::system_complete(cudnn_data_base_dir_path);
if (!boost::filesystem::is_directory(cudnn_data_base_dir_path))
cudnn_data_base_dir_path = cudnn_data_base_dir_path.branch_path();
if (!boost::filesystem::exists(cudnn_data_base_dir_path))
{
// exeのディレクトリが取得できなければカレントディレクトリに保存
cudnn_data_base_dir_path = boost::filesystem::current_path();
if (cudnn_data_base_dir_path.is_relative())
cudnn_data_base_dir_path = boost::filesystem::system_complete(cudnn_data_base_dir_path);
if (!boost::filesystem::exists(cudnn_data_base_dir_path))
cudnn_data_base_dir_path = "./";
}
if (boost::filesystem::exists(cudnn_data_base_dir_path))
{
const boost::filesystem::path cudnn_data_dir_path(cudnn_data_base_dir_path / "cudnn_data");
bool isOK = false;
if (boost::filesystem::exists(cudnn_data_dir_path))
isOK = true;
if (!isOK)
{
boost::system::error_code error;
const bool result = boost::filesystem::create_directory(cudnn_data_dir_path, error);
if (result && !error)
isOK = true;
}
if (isOK)
{
cudaDeviceProp prop;
if (cudaGetDeviceProperties(&prop, mGPUNo) == cudaSuccess)
{
std::string conv_filename(prop.name);
conv_filename += " conv ";
std::string deconv_filename(prop.name);
deconv_filename += " deconv ";
const boost::filesystem::path conv_data_path = cudnn_data_dir_path / conv_filename;
const boost::filesystem::path deconv_data_path = cudnn_data_dir_path / deconv_filename;
g_ConvCcuDNNAlgorithm.SetDataPath(conv_data_path.string());
g_DeconvCcuDNNAlgorithm.SetDataPath(deconv_data_path.string());
}
}
}
}
const boost::filesystem::path mode_dir_path(GetModeDirPath(model_dir));
if (!boost::filesystem::exists(mode_dir_path))
return Waifu2x::eWaifu2xError_FailedOpenModelFile;
CudaDeviceSet devset(process, mGPUNo);
if (mProcess == "cpu")
{
caffe::Caffe::set_mode(caffe::Caffe::CPU);
mIsCuda = false;
}
else
{
caffe::Caffe::set_mode(caffe::Caffe::GPU);
mIsCuda = true;
}
caffe::Caffe::SetGetcuDNNAlgorithmFunc(GetcuDNNAlgorithm);
caffe::Caffe::SetSetcuDNNAlgorithmFunc(SetcuDNNAlgorithm);
mInputPlane = 0;
mMaxNetOffset = 0;
const boost::filesystem::path info_path = GetInfoPath(mode_dir_path);
stInfo info;
ret = cNet::GetInfo(info_path, info);
if (ret != Waifu2x::eWaifu2xError_OK)
return ret;
mHasNoiseScale = info.has_noise_scale;
mInputPlane = info.channels;
if (mode == eWaifu2xModelTypeNoise || mode == eWaifu2xModelTypeNoiseScale || mode == eWaifu2xModelTypeAutoScale)
{
std::string base_name;
mNoiseNet.reset(new cNet);
eWaifu2xModelType Mode = mode;
if (info.has_noise_scale) // ノイズ除去と拡大を同時に行う
{
// ノイズ除去拡大ネットの構築はeWaifu2xModelTypeNoiseScaleを指定する必要がある
Mode = eWaifu2xModelTypeNoiseScale;
base_name = "noise" + std::to_string(noise_level) + "_scale2.0x_model";
}
else // ノイズ除去だけ
{
Mode = eWaifu2xModelTypeNoise;
base_name = "noise" + std::to_string(noise_level) + "_model";
}
const boost::filesystem::path model_path = mode_dir_path / (base_name + ".prototxt");
const boost::filesystem::path param_path = mode_dir_path / (base_name + ".json");
ret = mNoiseNet->ConstractNet(Mode, model_path, param_path, info, mProcess);
if (ret != Waifu2x::eWaifu2xError_OK)
return ret;
mMaxNetOffset = mNoiseNet->GetNetOffset();
}
// noise_scaleを持っている場合はαチャンネルの拡大のためにmScaleNetも構築する必要がある
if (info.has_noise_scale || mode == eWaifu2xModelTypeScale || mode == eWaifu2xModelTypeNoiseScale || mode == eWaifu2xModelTypeAutoScale)
{
const std::string base_name = "scale2.0x_model";
const boost::filesystem::path model_path = mode_dir_path / (base_name + ".prototxt");
const boost::filesystem::path param_path = mode_dir_path / (base_name + ".json");
mScaleNet.reset(new cNet);
ret = mScaleNet->ConstractNet(eWaifu2xModelTypeScale, model_path, param_path, info, mProcess);
if (ret != Waifu2x::eWaifu2xError_OK)
return ret;
assert(mInputPlane == 0 || mInputPlane == mScaleNet->GetInputPlane());
mMaxNetOffset = std::max(mScaleNet->GetNetOffset(), mMaxNetOffset);
}
else
{
}
mIsInited = true;
}
catch (...)
{
return Waifu2x::eWaifu2xError_InvalidParameter;
}
return Waifu2x::eWaifu2xError_OK;
}
Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &Mode, const int NoiseLevel, const double ScaleRatio, const boost::filesystem::path &ModelDir, const std::string &Process,
const boost::optional<int> OutputQuality, const int OutputDepth, const bool UseTTA, const int CropSize, const int BatchSize)
{
Waifu2x::eWaifu2xError ret;
if (is_inited)
return eWaifu2xError_OK;
if (ScaleRatio <= 0.0)
return eWaifu2xError_InvalidParameter;
try
{
mode = Mode;
noise_level = NoiseLevel;
scale_ratio = ScaleRatio;
model_dir = ModelDir;
process = Process;
use_tta = UseTTA;
output_quality = OutputQuality;
output_depth = OutputDepth;
crop_size = CropSize;
batch_size = BatchSize;
inner_padding = layer_num;
outer_padding = 1;
output_size = crop_size - offset * 2;
input_block_size = crop_size + (inner_padding + outer_padding) * 2;
original_width_height = 128 + layer_num * 2;
output_block_size = crop_size + (inner_padding + outer_padding - layer_num) * 2;
std::call_once(waifu2x_once_flag, [argc, argv]()
{
assert(argc >= 1);
int tmpargc = 1;
char* tmpargvv[] = { argv[0] };
char** tmpargv = tmpargvv;
// glog等の初期化
caffe::GlobalInit(&tmpargc, &tmpargv);
});
const auto cuDNNCheckStartTime = std::chrono::system_clock::now();
if (process == "gpu")
{
if (can_use_CUDA() != eWaifu2xCudaError_OK)
return eWaifu2xError_FailedCudaCheck;
// cuDNNが使えそうならcuDNNを使う
else if (can_use_cuDNN() == eWaifu2xcuDNNError_OK)
process = "cudnn";
}
const auto cuDNNCheckEndTime = std::chrono::system_clock::now();
boost::filesystem::path mode_dir_path(model_dir);
if (!mode_dir_path.is_absolute()) // model_dirが相対パスなら絶対パスに直す
{
// まずはカレントディレクトリ下にあるか探す
mode_dir_path = boost::filesystem::absolute(model_dir);
if (!boost::filesystem::exists(mode_dir_path) && argc >= 1) // 無かったらargv[0]から実行ファイルのあるフォルダを推定し、そのフォルダ下にあるか探す
{
boost::filesystem::path a0(argv[0]);
if (a0.is_absolute())
mode_dir_path = a0.branch_path() / model_dir;
}
}
if (!boost::filesystem::exists(mode_dir_path))
return eWaifu2xError_FailedOpenModelFile;
if (process == "cpu")
{
caffe::Caffe::set_mode(caffe::Caffe::CPU);
isCuda = false;
}
else
{
caffe::Caffe::set_mode(caffe::Caffe::GPU);
isCuda = true;
}
if (mode == "noise" || mode == "noise_scale" || mode == "auto_scale")
{
const boost::filesystem::path model_path = (mode_dir_path / "srcnn.prototxt").string();
const boost::filesystem::path param_path = (mode_dir_path / ("noise" + std::to_string(noise_level) + "_model.json")).string();
ret = ConstractNet(net_noise, model_path, param_path, process);
if (ret != eWaifu2xError_OK)
return ret;
}
if (mode == "scale" || mode == "noise_scale" || mode == "auto_scale")
{
const boost::filesystem::path model_path = (mode_dir_path / "srcnn.prototxt").string();
const boost::filesystem::path param_path = (mode_dir_path / "scale2.0x_model.json").string();
ret = ConstractNet(net_scale, model_path, param_path, process);
if (ret != eWaifu2xError_OK)
return ret;
}
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;
if (isCuda)
{
CUDA_CHECK_WAIFU2X(cudaHostAlloc(&input_block, sizeof(float) * input_block_plane_size * batch_size, cudaHostAllocWriteCombined));
CUDA_CHECK_WAIFU2X(cudaHostAlloc(&dummy_data, sizeof(float) * input_block_plane_size * batch_size, cudaHostAllocWriteCombined));
CUDA_CHECK_WAIFU2X(cudaHostAlloc(&output_block, sizeof(float) * output_block_plane_size * batch_size, cudaHostAllocDefault));
}
else
{
input_block = new float[input_block_plane_size * batch_size];
dummy_data = new float[input_block_plane_size * batch_size];
output_block = new float[output_block_plane_size * batch_size];
}
for (size_t i = 0; i < input_block_plane_size * batch_size; i++)
dummy_data[i] = 0.0f;
is_inited = true;
}
catch (...)
{
return eWaifu2xError_InvalidParameter;
}
return eWaifu2xError_OK;
}
