sources/libs/nvmultiobjecttracker/include/modules/FrameTransform/trtModule.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#ifndef TRTMODULE_HPP_
#define TRTMODULE_HPP_
 
#include <fstream>
#include <functional>
#include <memory>
#include <cstring>
#include <NvInfer.h>
#include <NvOnnxParser.h>
#include "nvdsinfer_context.h"
#include "modules/utility/cuda_common.hpp"
#include "modules/NvMultiObjectTracker/NvTrackerParams.hpp"
 
class TRTConfigParams : public NvTrackerParams
{
public:
    TRTConfigParams();
 
    TRTConfigParams(const std::string &configFilePath, const std::string moduleName, const std::string submoduleName);
 
    TRTConfigParams(const YAML::Node &fn);
 
    void read(const YAML::Node &fn) override;
 
    int m_maxBatchSize; 
    int m_inputOrder; 
    int m_width; 
    int m_height; 
    int m_channel; 
    int m_gpuId; 
    std::vector<int> m_inferDims; 
    size_t m_workspaceSize; 
    std::string m_tltEncodedModelFilePath; 
    std::string m_tltModelKey; 
    std::string m_onnxFilePath; 
    std::string m_engineFilePath; 
    std::string m_calibrationTableFilePath; 
    nvinfer1::DataType m_networkMode; 
    int m_builderOptimizationLevel; 
    int m_debugVerboseLevel; 
};
 
class TrackerInt8EntropyCalibrator : public nvinfer1::IInt8EntropyCalibrator
{
public:
    TrackerInt8EntropyCalibrator(std::string calibrationTableFile)
        : m_CalibrationTableFile(calibrationTableFile)
    {}
 
    ~TrackerInt8EntropyCalibrator() {}
 
    int getBatchSize() const noexcept override { return 0; }
 
    bool getBatch(void* bindings[], const char* names[], int nbBindings) noexcept override { return false; }
 
    const void*
    readCalibrationCache(size_t& length) noexcept override
    {
        m_CalibrationCache.clear();
        std::ifstream input(m_CalibrationTableFile, std::ios::binary);
        input >> std::noskipws;
        if (input.good())
            copy(
                std::istream_iterator<char>(input), std::istream_iterator<char>(),
                std::back_inserter(m_CalibrationCache)
            );
 
        length = m_CalibrationCache.size();
        return length ? m_CalibrationCache.data() : nullptr;
    }
 
    void writeCalibrationCache(const void* cache, size_t length) noexcept override {}
 
private:
    std::string m_CalibrationTableFile;
    std::vector<char> m_CalibrationCache;
};
 
class TrtModule
{
public:
    TrtModule();
 
    void initializeModel(
        std::function<void(
            nvinfer1::INetworkDefinition *network, nvinfer1::IBuilder *builder,
            nvinfer1::IBuilderConfig *config, int batchSize
        )> onnxProfileFunc,
        std::function<void(TrtModule*, const std::vector<int>& sizes)> setRuntimeDimFunc
    );
 
    void createInferEngine();
 
    void parseOnnxModel(
        nvinfer1::INetworkDefinition &network, nvinfer1::ILogger &logger, std::unique_ptr<nvonnxparser::IParser>& parser
    );
 
    nvinfer1::ICudaEngine* decodeTltModel(nvinfer1::IBuilder* const builder, nvinfer1::IBuilderConfig* const config);
 
    bool loadTRTEngine();
 
    void submitInput(std::vector<void *>& bindings, const std::vector<int>& batchSize, cudaStream_t& cuStream);
 
    void printEngineInfo();
 
    void setCudaStream(cudaStream_t cuStream) { m_cuStream = cuStream; };
 
    TRTConfigParams m_params; 
    std::shared_ptr<nvinfer1::ICudaEngine> m_engine; 
    std::shared_ptr<nvinfer1::IExecutionContext> m_context; 
    cudaStream_t m_cuStream; 
protected:
    std::shared_ptr<nvinfer1::IRuntime> m_runtime; 
    std::shared_ptr<TrackerInt8EntropyCalibrator> m_calibrator; 
    std::string m_absModelFilePath; 
    std::function<void(
        nvinfer1::INetworkDefinition *network, nvinfer1::IBuilder *builder, nvinfer1::IBuilderConfig *config,
        int batchSize)> m_onnxProfileFunc; 
    std::function<void(TrtModule *, const std::vector<int>&)> m_setRuntimeDimFunc; 
};
 
#endif // TRTMODULE_HPP_