Semantic Segmentation using ONNX Runtime

[abenchaaben]

Hi,

I'm a noob in using ONNX Runtime. I trained a model using PyTorch then I converted it to ONNX. I'm trying to use ONNX Runtime C++ api to perform a prediction using my model. The idea is to have an image mask as output. I tried to follow some examples of how we use the ONNX Runtime. But I did not succeed in having the right mask as output.

My input image is an RGB of 512x512

                   const int64_t batchSize = 1;
                    std::string instanceName{"image-classification-inference"};
                    Ort::Env env(OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING, instanceName.c_str());
                    Ort::SessionOptions sessionOptions;
                    sessionOptions.SetIntraOpNumThreads(1);
                    sessionOptions.SetGraphOptimizationLevel(
                    GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
                    Ort::Session session(env, model_path.c_str(), sessionOptions);
                    Ort::AllocatorWithDefaultOptions allocator;
                    size_t numInputNodes = session.GetInputCount();
                    size_t numOutputNodes = session.GetOutputCount();
                    const char* inputName = session.GetInputName(0, allocator);

                    Ort::TypeInfo inputTypeInfo = session.GetInputTypeInfo(0);
                    auto inputTensorInfo = inputTypeInfo.GetTensorTypeAndShapeInfo();
                    ONNXTensorElementDataType inputType = inputTensorInfo.GetElementType();
                    std::vector<int64_t> inputDims = inputTensorInfo.GetShape();
                    if (inputDims.at(0) == -1) {
                        std::cout << "Got dynamic batch size. Setting input batch size to "
                                << batchSize << "." << std::endl;
                        inputDims.at(0) = batchSize;
                    }
                    const char* outputName = session.GetOutputName(0, allocator);
                    Ort::TypeInfo outputTypeInfo = session.GetOutputTypeInfo(0);
                    auto outputTensorInfo = outputTypeInfo.GetTensorTypeAndShapeInfo();
                    ONNXTensorElementDataType outputType = outputTensorInfo.GetElementType();
                    std::vector<int64_t> outputDims = outputTensorInfo.GetShape();
                    if (outputDims.at(0) == -1)
                    {
                        std::cout << "Got dynamic batch size. Setting output batch size to "
                                << batchSize << "." << std::endl;
                                outputDims.at(0) = 256;
                                outputDims.at(1) = 256;
                                outputDims.at(2) = 2;
                    }

                    std::cout << "Number of Input Nodes: " << numInputNodes << std::endl;
                    std::cout << "Number of Output Nodes: " << numOutputNodes << std::endl;
                    std::cout << "Input Name: " << inputName << std::endl;
                    std::cout << "Input Type: " << inputType << std::endl;
                    std::cout << "Input Dimensions: " << inputDims << std::endl;
                    std::cout << "Output Name: " << outputName << std::endl;
                    std::cout << "Output Type: " << outputType << std::endl;
                    std::cout << "Output Dimensions: " << outputDims << std::endl;

                   cv::Mat resizedImageBGR, resizedImageRGB, resizedImage, preprocessedImage;
                    cv::resize(image, resizedImageBGR,
                            cv::Size(inputDims.at(3), inputDims.at(2)),
                            cv::InterpolationFlags::INTER_CUBIC);
                    cv::cvtColor(resizedImageBGR, resizedImageRGB,
                                cv::ColorConversionCodes::COLOR_BGR2RGB);
                    resizedImageRGB.convertTo(resizedImage, CV_32F, 1.0 / 255);

                    cv::imwrite("resizedImage.png", resizedImage);

                    // HWC to CHW
                    cv::dnn::blobFromImage(resizedImage, preprocessedImage);

                    size_t inputTensorSize = vectorProduct(inputDims);
                    std::vector<float> inputTensorValues(inputTensorSize);
                    // Make copies of the same image input.
                    for (int64_t i = 0; i < batchSize; ++i)
                    {
                        std::copy(preprocessedImage.begin<float>(),
                                preprocessedImage.end<float>(),
                                inputTensorValues.begin() + i * inputTensorSize / batchSize);
                    }

                    size_t outputTensorSize = vectorProduct(outputDims);
                    assert(("Output tensor size should equal to the label set size.",
                            labels.size() * batchSize == outputTensorSize));
                    std::vector<float> outputTensorValues(outputTensorSize);

                    std::vector<const char*> inputNames{inputName};
                    std::vector<const char*> outputNames{outputName};
                    std::vector<Ort::Value> inputTensors;
                    std::vector<Ort::Value> outputTensors;

                    Ort::MemoryInfo memoryInfo = Ort::MemoryInfo::CreateCpu(
                        OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);
                    inputTensors.push_back(Ort::Value::CreateTensor<float>(
                        memoryInfo, inputTensorValues.data(), inputTensorSize, inputDims.data(),
                        inputDims.size()));
                    outputTensors.push_back(Ort::Value::CreateTensor<float>(
                        memoryInfo, outputTensorValues.data(), outputTensorSize,
                        outputDims.data(), outputDims.size()));

                    session.Run(Ort::RunOptions{nullptr}, inputNames.data(),
                                inputTensors.data(), 1 , outputNames.data(),
                                outputTensors.data(), 1 );

                    cv::Mat out = cv::Mat::zeros(cv::Size(256,256),CV_32F);
                    memcpy(out.data, outputTensors[0].GetTensorMutableData<float>(), 256 * 256 * sizeof(float));

                    cv::imwrite("out.png", out);

[skottmckay]

Whatever pre-processing or post-processing you do needs to match exactly what the model requires and how it was trained.

e.g. is the ordering RGB or BGR? Is the input in NCHW layout or NHWC? When converting to float do the values need to be normalized and if so what values were used for that when training the model, and are there separate values for each channel or not?

I am not familiar with blobFromImage, however the documentation suggests it might swap the R and B channels, but the comment in your code implies it's changing the layout from channels last to channels first, which is a very different operation.

// HWC to CHW
cv::dnn::blobFromImage(resizedImage, preprocessedImage);