Cannot load squeezeNet 1.0 or 1.1 by opencv3.2 - c++

I pull the source codes from opencv and opencv_contrib today(2017/06/28), trying to load the squeezeNet models list at here as following example
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
using namespace cv;
using namespace cv::dnn;
#include <fstream>
#include <iostream>
#include <cstdlib>
using namespace std;
/* Find best class for the blob (i. e. class with maximal probability) */
static void getMaxClass(const Mat &probBlob, int *classId, double *classProb)
{
Mat probMat = probBlob.reshape(1, 1); //reshape the blob to 1x1000 matrix
Point classNumber;
minMaxLoc(probMat, NULL, classProb, NULL, &classNumber);
*classId = classNumber.x;
}
static std::vector<String> readClassNames(const char *filename = "synset_words.txt")
{
std::vector<String> classNames;
std::ifstream fp(filename);
if (!fp.is_open())
{
std::cerr << "File with classes labels not found: " << filename << std::endl;
exit(-1);
}
std::string name;
while (!fp.eof())
{
std::getline(fp, name);
if (name.length())
classNames.push_back( name.substr(name.find(' ')+1) );
}
fp.close();
return classNames;
}
int main(int argc, char **argv)
{
cv::dnn::initModule(); //Required if OpenCV is built as static libs
String modelTxt = "train_val.prototxt";
String modelBin = "squeezenet_v1.1.caffemodel";
String imageFile = (argc > 1) ? argv[1] : "space_shuttle.jpg";
Net net = dnn::readNetFromCaffe(modelTxt, modelBin);
if (net.empty())
{
std::cerr << "Can't load network by using the following files: " << std::endl;
std::cerr << "prototxt: " << modelTxt << std::endl;
std::cerr << "caffemodel: " << modelBin << std::endl;
std::cerr << "bvlc_googlenet.caffemodel can be downloaded here:" << std::endl;
std::cerr << "http://dl.caffe.berkeleyvision.org/bvlc_googlenet.caffemodel" << std::endl;
exit(-1);
}
Mat img = imread(imageFile);
if (img.empty())
{
std::cerr << "Can't read image from the file: " << imageFile << std::endl;
exit(-1);
}
//GoogLeNet accepts only 224x224 RGB-images
Mat inputBlob = blobFromImage(img, 1, Size(224, 224),
Scalar(104, 117, 123)); //Convert Mat to batch of images
net.setInput(inputBlob, "data"); //set the network input
Mat prob = net.forward("prob"); //compute output
int classId;
double classProb;
getMaxClass(prob, &classId, &classProb);//find the best class
std::vector<String> classNames = readClassNames();
std::cout << "Best class: #" << classId << " '" << classNames.at(classId) << "'" << std::endl;
std::cout << "Probability: " << classProb * 100 << "%" << std::endl;
return 0;
} //main
But it give me following error(works well with google net)
OpenCV Error: Bad argument (Duplicate blobs produced by multiple
sources) in addOutput, file
/home/ramsus/Qt/3rdLibs/opencv/modules/dnn/src/caffe/caffe_importer.cpp,
line 327
/home/ramsus/Qt/3rdLibs/opencv/modules/dnn/src/caffe/caffe_importer.cpp:327:
error: (-5) Duplicate blobs produced by multiple sources in function
addOutput
How could I import squeezeNet_v1.1 by opencv?Thanks

Mat prob = net.forward("prob");
change to
Mat prob = net.forward("softmax");

Related

Cannot load traced_torch_model_pt in C++

I tried to load pre-built traced_torch_model.pt to C++; however, the error message showed "example-app.exe" existed with code -1
Please see below my code snippets and error message screenshot.
I appreciate your input.
using namespace std;
#include <torch/script.h>
#include <iostream>
#include <memory>
int main(int argc, const char* argv[]) {
if (argc != 2) {
std::cerr << "usage: example-app <C:\\Users\\C at work\\0813_2021\\traced_torch_model.pt\\>n";
return -1;
}
torch::jit::script::Module module;
try {
// Deserialize the ScriptModule from a file using torch::jit::load().
module = torch::jit::load(argv[1]);
}
catch (const c10::Error& e) {
std::cerr << "error loading the model\n";
return -1;
}
std::cout << "Model " << argv[1] << " loaded fine\n";
// Create a vector of inputs.
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::randn({ 1, 1, 64, 101 }));
// Execute the model and turn its output into a tensor.
at::Tensor output = module.forward(inputs).toTensor();
std::cout << output << "\n";
int y_hat = output.argmax(1).item().toInt();
std::cout << "Predicted class: " << y_hat << "\n";
}

translate a boost ASIO buffer data to an opencv matrix of JPEG format

I've built a C++ program to send a single jpg frame (along with a metadata header beforehand) over a TCP socket.. my trouble is.. I am getting the bytes successfully.. parsing my header and seeing my dimenions.. the number of bytes... but my question is.. how do I take this data now.. and recreate my JPEG?
I keep trying to cast the buffer but getting stopped out
std::cout << "[IMAGE_WIDTH] = " << std::to_string(IMAGE_WIDTH) << std::endl;
std::cout << "[IMAGE_HEIGHT] = " << std::to_string(IMAGE_HEIGHT) << std::endl;
std::cout << "image bytes = " << std::to_string(header_data.image_size_bytes) << std::endl;
Mat img = Mat::zeros(IMAGE_WIDTH, IMAGE_HEIGHT, CV_8UC3);
const unsigned char* image_data_matrix = boost::asio::buffer_cast<const unsigned char*>(second_receive_buffer.data());
img.data = image_data_matrix;
imshow("server", img);
waitKey(100);
error I get from gcc
[ 50%] Building CXX object CMakeFiles/sockets-client.dir/src/client.cpp.o
/bootstrap-cpp-sockets/client/src/client.cpp:81:28: error: assigning to 'uchar *' (aka 'unsigned char *') from 'const unsigned char *' discards qualifiers
img.data = image_data_matrix;
^~~~~~~~~~~~~~~~~
1 error generated.
make[2]: *** [CMakeFiles/sockets-client.dir/src/client.cpp.o] Error 1
make[1]: *** [CMakeFiles/sockets-client.dir/all] Error 2
make: *** [all] Error 2
Full Client.cpp
int main()
{
try{
boost::asio::io_service io_service;
tcp::endpoint end_point(boost::asio::ip::address::from_string("127.0.0.1"), 3200);
tcp::socket socket(io_service);
socket.connect(end_point);
boost::system::error_code ignored_error;
boost::asio::streambuf receive_buffer;
// Now we retrieve the message header of 64 bytes
size_t header_size = 64;
boost::asio::read(socket, receive_buffer, boost::asio::transfer_exactly(header_size), ignored_error);
if ( ignored_error && ignored_error != boost::asio::error::eof ) {
cout << "first receive failed: " << ignored_error.message() << endl;
} else {
image_metadata_t header_data = parse_header(receive_buffer);
const int IMAGE_WIDTH = header_data.width;
const int IMAGE_HEIGHT = header_data.height;
// Now we retrieve the frame itself
boost::asio::streambuf second_receive_buffer;
boost::asio::read(socket, second_receive_buffer, boost::asio::transfer_exactly(header_data.image_size_bytes), ignored_error);
if( ignored_error && ignored_error != boost::asio::error::eof ) {
cout << "SECOND receive failed: " << ignored_error.message() << endl;
}
else {
std::cout << "[IMAGE_WIDTH] = " << std::to_string(IMAGE_WIDTH) << std::endl;
std::cout << "[IMAGE_HEIGHT] = " << std::to_string(IMAGE_HEIGHT) << std::endl;
std::cout << "image bytes = " << std::to_string(header_data.image_size_bytes) << std::endl;
Mat img = Mat::zeros(IMAGE_WIDTH, IMAGE_HEIGHT, CV_8UC3);
const unsigned char* image_data_matrix = boost::asio::buffer_cast<const unsigned char*>(second_receive_buffer.data());
img.data = image_data_matrix;
imshow("server", img);
waitKey(100);
}
}
} catch (std::exception& e) {
std::cerr << e.what() << std::endl;
}
return 0;
}
the eventual finished program will keep reading the socket, receiving a header with the metadata.. and showing the next frame.. hopefully 30 Frames per second will be hit but.. first thing first I just want to prove i can do this with one frame.. and build up from there.
Thanks !
after some clarification on my part with #Dan in the chat.. I modified both the server side and client side.. I append them both here now for anyone else trying to send JPG data over the wire
#include <iostream>
#include <boost/array.hpp>
#include <boost/asio.hpp>
#include <boost/thread/thread.hpp>
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/imgproc/imgproc.hpp"
#include <opencv2/core/core.hpp>
#include <string>
#include <vector>
using boost::asio::ip::tcp;
using namespace std;
using namespace cv;
// source https://www.programmerall.com/article/7721373696/
// source https://cppsecrets.com/users/14041151035752494957504952535764103109971051084699111109/Programming-in-C00-using-boostasio.php
bool flag = false;
void servershow()
{
while (true)
{
if (flag)
{
//imshow("server",img);
waitKey(1);
}
}
}
cv::Mat retrieve_data(){
std::string image_path = "/YOUR_IMAGE.jpg";
cv::Mat image;
image = imread(image_path, cv::IMREAD_COLOR);
if(! image.data ) {
std::cout << "Could not open or find the image" << std::endl;
}
return image;
}
int main()
{
boost::thread thrd(&servershow);
try
{
boost::asio::io_service io_service;
tcp::acceptor acceptor(io_service, tcp::endpoint(tcp::v4(), 3200));
for (;;) {
tcp::socket socket(io_service);
acceptor.accept(socket);
boost::system::error_code ignored_error;
//retrieve the frame to be sent
cv::Mat frame = retrieve_data();
std::vector<std::uint8_t> buff;
cv::imencode(".jpg", frame, buff, param);
// now we send the header message
std:: string image_dimensions = "6016Wx3384H";
std:: string image_buff_bytes = std::to_string(buff.size());
std::string message_header = image_dimensions + "," + image_buff_bytes;
std::cout << "sending measage header of " << std::to_string(message_header.length()) << " bytes...." << std::endl;
message_header.append(63 - message_header.length(), ' ');
message_header = message_header + '\0';
socket.write_some(boost::asio::buffer(message_header), ignored_error);
socket.write_some(boost::asio::buffer(buff), ignored_error);
flag = true;
}
}
catch (std::exception& e)
{
std::cerr << e.what() << std::endl;
}
thrd.join();
return 0;
}
Client retrieval
#include <iostream>
#include <boost/asio.hpp>
#include <boost/array.hpp>
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/imgproc/imgproc.hpp"
#include <opencv2/objdetect/objdetect.hpp>
#include <string>
#include <vector>
using boost::asio::ip::tcp;
using namespace std;
using namespace cv;
// source https://www.programmerall.com/article/7721373696/
// source https://cppsecrets.com/users/14041151035752494957504952535764103109971051084699111109/Programming-in-C00-using-boostasio.php
struct image_metadata_t {
int width;
int height;
size_t image_size_bytes;
};
image_metadata_t parse_header(boost::asio::streambuf &buffer){
std::string data_buff_str = std::string(boost::asio::buffer_cast<const char*>(buffer.data()));
cout << data_buff_str << endl;
int width_pos = data_buff_str.find("W");
int x_pos = data_buff_str.find("x");
int height_pos = data_buff_str.find("H");
int comma_pos = data_buff_str.find(",");
std::cout << "data_buff_str.substr(x_pos+1,height_pos) " << data_buff_str.substr(x_pos+1,height_pos) <<std::endl;
image_metadata_t meta_data;
meta_data.width = std::stoi(data_buff_str.substr(0,width_pos));
meta_data.height = std::stoi(data_buff_str.substr(x_pos+1,height_pos));
meta_data.image_size_bytes = std::stoi(data_buff_str.substr(data_buff_str.find(",") + 1));
return meta_data;
}
cv::Mat GetImageFromMemory(uchar* image, int length, int flag) {
std::vector<uchar> data = std::vector<uchar>(image, image + length);
cv::Mat ImMat = imdecode(data, flag);
return ImMat;
}
int main()
{
try{
boost::asio::io_service io_service;
tcp::endpoint end_point(boost::asio::ip::address::from_string("127.0.0.1"), 3200);
tcp::socket socket(io_service);
socket.connect(end_point);
boost::system::error_code ignored_error;
boost::asio::streambuf receive_buffer;
// Now we retrieve the message header of 64 bytes
size_t header_size = 64;
boost::asio::read(socket, receive_buffer, boost::asio::transfer_exactly(header_size), ignored_error);
if ( ignored_error && ignored_error != boost::asio::error::eof ) {
cout << "first receive failed: " << ignored_error.message() << endl;
} else {
image_metadata_t header_data = parse_header(receive_buffer);
const int IMAGE_WIDTH = header_data.width;
const int IMAGE_HEIGHT = header_data.height;
// Now we retrieve the frame itself
std::cout << "Now asing for image bytes of size " << std::to_string(header_data.image_size_bytes) << std::endl;
//boost::asio::streambuf second_receive_buffer;
std::vector<std::uint8_t> buff(header_data.image_size_bytes);
boost::asio::read(socket, boost::asio::buffer(buff), boost::asio::transfer_exactly(header_data.image_size_bytes), ignored_error);
if( ignored_error && ignored_error != boost::asio::error::eof ) {
cout << "SECOND receive failed: " << ignored_error.message() << endl;
}
else {
std::cout << "[IMAGE_WIDTH] = " << std::to_string(IMAGE_WIDTH) << std::endl;
std::cout << "[IMAGE_HEIGHT] = " << std::to_string(IMAGE_HEIGHT) << std::endl;
std::cout << "image bytes = " << std::to_string(header_data.image_size_bytes) << std::endl;
cv::Mat img(cv::imdecode(buff, cv::IMREAD_ANYCOLOR));
imshow("client", img);
waitKey(5000);
}
}
} catch (std::exception& e) {
std::cerr << e.what() << std::endl;
}
return 0;
}

PyTorch and TorchVision FasterRCNN interpreting the output in C++ GenericDict

I'm trying to interpret the output of FasterRCNN in C++ and I'm fighting with the GenericDict type.
My code is as follows:
#include <opencv4/opencv2/opencv.hpp>
#include <opencv4/opencv2/shape.hpp>
#include <opencv4/opencv2/imgcodecs.hpp>
#include <opencv4/opencv2/highgui.hpp>
#include <opencv4/opencv2/imgproc.hpp>
#include <opencv4/opencv2/core/utility.hpp>
#include <opencv4/opencv2/core/mat.hpp>
#include <c10/cuda/CUDAStream.h>
#include <torch/csrc/autograd/grad_mode.h>
#include <torch/csrc/api/include/torch/torch.h>
#include <torch/script.h>
#include <torchvision/vision.h>
#include <torchvision/nms.h>
#include <iostream>
#include <memory>
#include <string>
int main(int argc, const char* argv[])
{
if (argc != 3)
{
printf("usage: %s <path-to-exported-script-module> <image_to_test>\n",argv[0]);
return -1;
}
std::string module_filename = argv[1];
std::string image_file = argv[2];
try
{
cv::Mat input_img = cv::imread(image_file, cv::IMREAD_GRAYSCALE);
torch::autograd::AutoGradMode guard(false);
// Deserialize the ScriptModule from a file using torch::jit::load().
torch::jit::script::Module module = torch::jit::load(module_filename);
assert(module.buffers().size() > 0);
module.eval();
// Assume that the entire model is on the same device.
// We just put input to this device.
auto device = (*std::begin(module.buffers())).device();
const int height = input_img.rows;
const int width = input_img.cols;
const int channels = 1;
auto input = torch::from_blob(input_img.data, {height, width, channels}, torch::kUInt8);
// HWC to CHW
// input = input.to(device, torch::kFloat).permute({2, 0, 1}).contiguous();
input = input.to(device, torch::kFloat).permute({2, 0, 1}).contiguous();
// run the network
std::vector<at::Tensor> inputs;
inputs.push_back(input);
auto output = module.forward({inputs});
if (device.is_cuda())
c10::cuda::getCurrentCUDAStream().synchronize();
std::cout << "output: " << output << std::endl;
auto outputs = output.toTuple()->elements();
std::cout << "outputs: " << outputs << std::endl;
for( auto& elem : outputs )
{
std::cout << "elem: " << elem << std::endl;
if( elem.isGenericDict() )
{
std::cout << "elem is generic dict: " << elem << std::endl;
c10::Dict<c10::IValue, c10::IValue> dict = elem.toGenericDict();
auto elem_vector_0 = dict.at(c10::IValue("scores")).toIntVector();
auto elem_vector_1 = dict.at(c10::IValue("boxes")).toIntVector();
auto elem_vector_2 = dict.at(c10::IValue("labels")).toIntVector();
for( auto& ee0 : elem_vector_0 )
{
std::cout << "elem_vector_0" << ee0 << std::endl;
}
for( auto& ee0 : elem_vector_1 )
{
std::cout << "elem_vector_1" << ee0 << std::endl;
}
for( auto& ee0 : elem_vector_2 )
{
std::cout << "elem_vector_2" << ee0 << std::endl;
}
}
}
cv::namedWindow("Display Image", cv::WINDOW_AUTOSIZE );
cv::imshow("Display Image", input_img);
cv::waitKey(0);
}
catch(const c10::Error& e)
{
std::cerr << e.what() << std::endl;
return -1;
}
catch(const cv::Exception& e)
{
std::cerr << e.what() << std::endl;
return -1;
}
catch(const std::exception& e)
{
std::cerr << e.what() << std::endl;
return -1;
}
catch(...)
{
std::cerr << "Unknown error" << std::endl;
return -1;
}
std::cout << "ok\n";
return 0;
}
and the output is:
(base) fstrati#fstrati-desktop:~/libtorch_shared_cuda_10.1/load_and_run_model/Release$ ./load_and_run_model ./torch_script_v0.2.pt test_img.png
[W faster_rcnn.py:95] Warning: RCNN always returns a (Losses, Detections) tuple in scripting (function )
output: ({}, [{boxes: [ CPUFloatType{0,4} ], labels: [ CPULongType{0} ], scores: [ CPUFloatType{0} ]}])
outputs: {} [{boxes: [ CPUFloatType{0,4} ], labels: [ CPULongType{0} ], scores: [ CPUFloatType{0} ]}]
elem: {}
elem is generic dict: {}
Argument passed to at() was not in the map.
I'm struggling to find a way to extract the boxes, labels and scores from the dictionary
GenericDict.
This map is strange, I cannot iterate on it and I cannot access first and second types...
with it->first it->second
Any ideas ?
Thanks in advance
I think the following method can resolve the main problem here,
output = module.forward(inputs);
auto detections = output.toTuple()->elements().at(1).toList().get(0).toGenericDict();
std::cout << ">>> detections labels: " << detections.at("labels") << std::endl;
std::cout << ">>> detections boxes: " << detections.at("boxes") << std::endl;
std::cout << ">>> detections scores: " << detections.at("scores") << std::endl;
Besides, I've added an executable file https://github.com/zhiqwang/yolov5-rt-stack/tree/master/deployment/libtorch to show how libtorch works.

How to use ONNX model in C++ code on Linux?

I train some Unet-based model in Pytorch. It take an image as an input, and return a mask.
After training i save it to ONNX format, run it with onnxruntime python module and it worked like a charm.
Now, i want to use this model in C++ code in Linux.
Is there simple tutorial (Hello world) when explained:
How to incorporate onnxruntime module to C++ program in Ubuntu
(install shared lib and so on)?
How to properly load an image and pass it to model?
P.S. I found only this: https://www.onnxruntime.ai/docs/tutorials/samples_catalog.html#cc
But there no info about loading image and converting it to ONNX - compatible format in C++ code.
For installation on the Linux, you should refer to https://www.onnxruntime.ai/.
You can refer to the following code to get help regarding how to load and run the ONNX model.
#include <algorithm> // std::generate
#include <assert.h>
#include <iostream>
#include <sstream>
#include <vector>
#include <experimental_onnxruntime_cxx_api.h>
// pretty prints a shape dimension vector
std::string print_shape(const std::vector<int64_t>& v) {
std::stringstream ss("");
for (size_t i = 0; i < v.size() - 1; i++)
ss << v[i] << "x";
ss << v[v.size() - 1];
return ss.str();
}
int calculate_product(const std::vector<int64_t>& v) {
int total = 1;
for (auto& i : v) total *= i;
return total;
}
using namespace std;
int main(int argc, char** argv) {
if (argc != 2) {
cout << "Usage: ./onnx-api-example <onnx_model.onnx>" << endl;
return -1;
}
#ifdef _WIN32
std::string str = argv[1];
std::wstring wide_string = std::wstring(str.begin(), str.end());
std::basic_string<ORTCHAR_T> model_file = std::basic_string<ORTCHAR_T>(wide_string);
#else
std::string model_file = argv[1];
#endif
// onnxruntime setup
Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "example-model-explorer");
Ort::SessionOptions session_options;
Ort::Experimental::Session session = Ort::Experimental::Session(env, model_file, session_options); // access experimental components via the Experimental namespace
// print name/shape of inputs
std::vector<std::string> input_names = session.GetInputNames();
std::vector<std::vector<int64_t> > input_shapes = session.GetInputShapes();
cout << "Input Node Name/Shape (" << input_names.size() << "):" << endl;
for (size_t i = 0; i < input_names.size(); i++) {
cout << "\t" << input_names[i] << " : " << print_shape(input_shapes[i]) << endl;
}
// print name/shape of outputs
std::vector<std::string> output_names = session.GetOutputNames();
std::vector<std::vector<int64_t> > output_shapes = session.GetOutputShapes();
cout << "Output Node Name/Shape (" << output_names.size() << "):" << endl;
for (size_t i = 0; i < output_names.size(); i++) {
cout << "\t" << output_names[i] << " : " << print_shape(output_shapes[i]) << endl;
}
// Assume model has 1 input node and 1 output node.
assert(input_names.size() == 1 && output_names.size() == 1);
// Create a single Ort tensor of random numbers
auto input_shape = input_shapes[0];
int total_number_elements = calculate_product(input_shape);
std::vector<float> input_tensor_values(total_number_elements);
std::generate(input_tensor_values.begin(), input_tensor_values.end(), [&] { return rand() % 255; }); // generate random numbers in the range [0, 255]
std::vector<Ort::Value> input_tensors;
input_tensors.push_back(Ort::Experimental::Value::CreateTensor<float>(input_tensor_values.data(), input_tensor_values.size(), input_shape));
// double-check the dimensions of the input tensor
assert(input_tensors[0].IsTensor() &&
input_tensors[0].GetTensorTypeAndShapeInfo().GetShape() == input_shape);
cout << "\ninput_tensor shape: " << print_shape(input_tensors[0].GetTensorTypeAndShapeInfo().GetShape()) << endl;
// pass data through model
cout << "Running model...";
try {
auto output_tensors = session.Run(session.GetInputNames(), input_tensors, session.GetOutputNames());
cout << "done" << endl;
// double-check the dimensions of the output tensors
// NOTE: the number of output tensors is equal to the number of output nodes specifed in the Run() call
assert(output_tensors.size() == session.GetOutputNames().size() &&
output_tensors[0].IsTensor());
cout << "output_tensor_shape: " << print_shape(output_tensors[0].GetTensorTypeAndShapeInfo().GetShape()) << endl;
} catch (const Ort::Exception& exception) {
cout << "ERROR running model inference: " << exception.what() << endl;
exit(-1);
}
}

How to convert tensor to image array?

I would like to convert a tensor to image array and use tensor.data() method.
But it doesn't work.
#include <torch/script.h> // One-stop header.
#include <iostream>
#include <memory>
#include <sstream>
#include <string>
#include <vector>
#include "itkImage.h"
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkImageRegionIterator.h"
//////////////////////////////////////////////////////
//Goal: load jit script model and segment myocardium
//Step: 1. load jit script model
// 2. load input image
// 3. predict by model
// 4. save the result to file
//////////////////////////////////////////////////////
typedef short PixelType;
const unsigned int Dimension = 3;
typedef itk::Image<PixelType, Dimension> ImageType;
typedef itk::ImageFileReader<ImageType> ReaderType;
typedef itk::ImageRegionIterator<ImageType> IteratorType;
bool itk2tensor(ImageType::Pointer itk_img, torch::Tensor &tensor_img) {
typename ImageType::RegionType region = itk_img->GetLargestPossibleRegion();
const typename ImageType::SizeType size = region.GetSize();
std::cout << "Input size: " << size[0] << ", " << size[1]<< ", " << size[2] << std::endl;
int len = size[0] * size[1] * size[2];
short rowdata[len];
int count = 0;
IteratorType iter(itk_img, itk_img->GetRequestedRegion());
// convert itk to array
for (iter.GoToBegin(); !iter.IsAtEnd(); ++iter) {
rowdata[count] = iter.Get();
count++;
}
std::cout << "Convert itk to array DONE!" << std::endl;
// convert array to tensor
tensor_img = torch::from_blob(rowdata, {1, 1, (int)size[0], (int)size[1], (int)size[2]}, torch::kShort).clone();
tensor_img = tensor_img.toType(torch::kFloat);
tensor_img = tensor_img.to(torch::kCUDA);
tensor_img.set_requires_grad(0);
return true;
}
bool tensor2itk(torch::Tensor &t, ImageType::Pointer itk_img) {
std::cout << "tensor dtype = " << t.dtype() << std::endl;
std::cout << "tensor size = " << t.sizes() << std::endl;
t = t.toType(torch::kShort);
short * array = t.data<short>();
ImageType::IndexType start;
start[0] = 0; // first index on X
start[1] = 0; // first index on Y
start[2] = 0; // first index on Z
ImageType::SizeType size;
size[0] = t.size(2);
size[1] = t.size(3);
size[2] = t.size(4);
ImageType::RegionType region;
region.SetSize( size );
region.SetIndex( start );
itk_img->SetRegions( region );
itk_img->Allocate();
int len = size[0] * size[1] * size[2];
IteratorType iter(itk_img, itk_img->GetRequestedRegion());
int count = 0;
// convert array to itk
std::cout << "start!" << std::endl;
for (iter.GoToBegin(); !iter.IsAtEnd(); ++iter) {
short temp = *array++; // ERROR!
std::cout << temp << " ";
iter.Set(temp);
count++;
}
std::cout << "end!" << std::endl;
return true;
}
int main(int argc, const char* argv[]) {
int a, b, c;
if (argc != 4) {
std::cerr << "usage: automyo input jitmodel output\n";
return -1;
}
std::cout << "========= jit start =========\n";
// 1. load jit script model
std::cout << "Load script module: " << argv[2] << std::endl;
std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(argv[2]);
module->to(at::kCUDA);
// assert(module != nullptr);
std::cout << "Load script module DONE" << std::endl;
// 2. load input image
const char* img_path = argv[1];
std::cout << "Load image: " << img_path << std::endl;
ReaderType::Pointer reader = ReaderType::New();
if (!img_path) {
std::cout << "Load input file error!" << std::endl;
return false;
}
reader->SetFileName(img_path);
reader->Update();
std::cout << "Load image DONE!" << std::endl;
ImageType::Pointer itk_img = reader->GetOutput();
torch::Tensor tensor_img;
if (!itk2tensor(itk_img, tensor_img)) {
std::cerr << "itk2tensor ERROR!" << std::endl;
}
else {
std::cout << "Convert array to tensor DONE!" << std::endl;
}
std::vector<torch::jit::IValue> inputs;
inputs.push_back(tensor_img);
// 3. predict by model
torch::Tensor y = module->forward(inputs).toTensor();
std::cout << "Inference DONE!" << std::endl;
// 4. save the result to file
torch::Tensor seg = y.gt(0.5);
// std::cout << seg << std::endl;
ImageType::Pointer out_itk_img = ImageType::New();
if (!tensor2itk(seg, out_itk_img)) {
std::cerr << "tensor2itk ERROR!" << std::endl;
}
else {
std::cout << "Convert tensor to itk DONE!" << std::endl;
}
std::cout << out_itk_img << std::endl;
return true;
}
The runtime log is showed below:
Load script module:model_myo_jit.pt
Load script module DONE Load
image: patch_6.nii.gz
Load image DONE!
Input size: 128, 128, 128
Convert itk to array DONE!
Convert array to tensor DONE!
Inference DONE!
tensor dtype = unsigned char
tensor size = [1, 1, 96, 96, 96]
start!
Segmentation fault (core dumped)
Why and how to convert?
I have found the solution. When I convert the y to kCPU, it works. Because it in CUDA before.