コード例 #1
0
ファイル: waifu2x.cpp プロジェクト: nevercannot/waifu2x-caffe
// モデルファイルからネットワークを構築
// processでcudnnが指定されなかった場合はcuDNNが呼び出されないように変更する
Waifu2x::eWaifu2xError Waifu2x::ConstractNet(boost::shared_ptr<caffe::Net<float>> &net, const std::string &model_path, const std::string &param_path, const std::string &process)
{
	const std::string caffemodel_path = param_path + ".caffemodel";
	const std::string modelbin_path = model_path + ".protobin";

	FILE *fp = fopen(caffemodel_path.c_str(), "rb");
	const bool isModelExist = fp != nullptr;
	if (fp) fclose(fp);

	fp = fopen(modelbin_path.c_str(), "rb");
	const bool isModelBinExist = fp != nullptr;
	if (fp) fclose(fp);

	caffe::NetParameter param;
	if (isModelExist && isModelBinExist && caffe::ReadProtoFromBinaryFile(modelbin_path, &param))
	{
		const auto ret = SetParameter(param);
		if (ret != eWaifu2xError_OK)
			return ret;

		net = boost::shared_ptr<caffe::Net<float>>(new caffe::Net<float>(param));
		net->CopyTrainedLayersFrom(caffemodel_path);

		input_plane = param.input_dim(1);
	}
	else
		return eWaifu2xError_FailedConstructModel;

	return eWaifu2xError_OK;
}
コード例 #2
0
ファイル: waifu2x.cpp プロジェクト: soeaver/waifu2x-caffe
// モデルファイルからネットワークを構築
// processでcudnnが指定されなかった場合はcuDNNが呼び出されないように変更する
Waifu2x::eWaifu2xError Waifu2x::ConstractNet(boost::shared_ptr<caffe::Net<float>> &net, const boost::filesystem::path &model_path, const boost::filesystem::path &param_path, const std::string &process)
{
	boost::filesystem::path modelbin_path = model_path;
	modelbin_path += ".protobin";
	boost::filesystem::path caffemodel_path = param_path;
	caffemodel_path += ".caffemodel";

	caffe::NetParameter param_model;
	caffe::NetParameter param_caffemodel;

	const auto retModelBin = readProtoBinary(modelbin_path, &param_model);
	const auto retParamBin = readProtoBinary(caffemodel_path, &param_caffemodel);

	if (retModelBin == eWaifu2xError_OK && retParamBin == eWaifu2xError_OK)
	{
		Waifu2x::eWaifu2xError ret;

		ret = SetParameter(param_model, process);
		if (ret != eWaifu2xError_OK)
			return ret;

		if (!caffe::UpgradeNetAsNeeded(caffemodel_path.string(), &param_caffemodel))
			return Waifu2x::eWaifu2xError_FailedParseModelFile;

		net = boost::shared_ptr<caffe::Net<float>>(new caffe::Net<float>(param_model));
		net->CopyTrainedLayersFrom(param_caffemodel);

		input_plane = param_model.input_dim(1);
	}
	else
	{
		const auto ret = LoadParameterFromJson(net, model_path, param_path, modelbin_path, caffemodel_path, process);
		if (ret != eWaifu2xError_OK)
			return ret;
	}

	return eWaifu2xError_OK;
}