About processing pytorch image segmenation in c++

(swp) #1
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”)

#2

I have a similar problem with a fully convolutional network for image segmentation. Python results are as expected but libtorch c++ frontend predictions differ a lot. Did you find a solution for this problem?

(Will Feng) #3

@swp @Kallinski Would you like to post the issue with a minimal reproducible example in https://github.com/pytorch/pytorch/issues?