diff --git a/common/waifu2x.cpp b/common/waifu2x.cpp index 84a849c..ee0061e 100644 --- a/common/waifu2x.cpp +++ b/common/waifu2x.cpp @@ -242,10 +242,19 @@ static Waifu2x::eWaifu2xError readProtoText(const boost::filesystem::path &path, static Waifu2x::eWaifu2xError readProtoBinary(const boost::filesystem::path &path, ::google::protobuf::Message* proto) { - boost::iostreams::stream is(path, std::ios_base::in | std::ios_base::binary); + boost::iostreams::stream is; + + try + { + is.open(path, std::ios_base::in | std::ios_base::binary); + } + catch (...) + { + return Waifu2x::eWaifu2xError_FailedOpenModelFile; + } if (!is) - return Waifu2x::eWaifu2xError_FailedParseModelFile; + return Waifu2x::eWaifu2xError_FailedOpenModelFile; std::vector tmp; if (!readFile(is, tmp)) @@ -617,7 +626,12 @@ Waifu2x::eWaifu2xError Waifu2x::ConstractNet(boost::shared_ptr modelbin_path += ".protobin"; caffe::NetParameter param; - if (readProtoBinary(modelbin_path, ¶m) == eWaifu2xError_OK) + caffe::NetParameter param_caffemodel; + + const auto retModelBin = readProtoBinary(caffemodel_path, ¶m); + const auto retParamBin = readProtoBinary(modelbin_path, ¶m_caffemodel); + + if (retModelBin == eWaifu2xError_OK && retParamBin == eWaifu2xError_OK) { Waifu2x::eWaifu2xError ret; @@ -625,11 +639,6 @@ Waifu2x::eWaifu2xError Waifu2x::ConstractNet(boost::shared_ptr if (ret != eWaifu2xError_OK) return ret; - caffe::NetParameter param_caffemodel; - ret = readProtoBinary(caffemodel_path, ¶m_caffemodel); - if (ret != eWaifu2xError_OK) - return ret; - if (!caffe::UpgradeNetAsNeeded(caffemodel_path.string(), ¶m_caffemodel)) return Waifu2x::eWaifu2xError_FailedParseModelFile;