string model_path = "Deep0.93.pt";
std::ifstream is(model_path, std::ifstream::binary);
std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(is);
module->to(at::kCUDA);
assert(module != nullptr);
std::cout << "ok\n";
img_color = imread("psw.jpg", IMREAD_COLOR);
imshow("abc", img_color);
cvtColor(img_color, img_color, COLOR_BGR2RGB);
if (img_color.empty()) {
//continue;
return -1;
}
std::vector<int64_t> sizes = { 1, 3, img_color.rows, img_color.cols };
at::TensorOptions options(at::ScalarType::Byte);
at::Tensor tensor_image = torch::from_blob(img_color.data, at::IntList(sizes), options);
tensor_image = tensor_image.toType(at::kFloat);
std::vector<torch::jit::IValue> inputs;
inputs.emplace_back(tensor_image);
at::Tensor results = module->forward(inputs).toTensor().to(at::kCUDA);
auto output = results.cpu().detach();
output = output.squeeze();
Mat result(output.size(0), output.size(1), CV_32F, output.data_ptr());
Mat temp = result.clone();
for (int i = 0; i < temp.rows; i++)
{
for (int j = 0; j < temp.cols; j++)
{
if (temp.at<float>(i, j) > 0.5f)
temp.at<float>(i, j) = 255;
else
temp.at<float>(i, j) = 0;
}
}
temp.convertTo(temp, CV_8U);
cv::resize(temp, temp, Size(224, 224), 1);
imshow("result", temp);
waitKey(0);
This is my code…
but the result is too weird.
Do I have any mistakes?
The model is not a problem
My guess is that the mat type changes strangely in the tensor conversion process.
refer… save pt file
net = net()
net.to(device)
net.load_state_dict(torch.load(…))
…
example = torch.read(1,3,224,224)
example = example.to(device)
traced_script_module = torch.jit.trace(net, example)
output = traced_script_module(example)
traced_script_module.save(“xxx.pt”)