sources/gst-plugins/gst-nvdsvideotemplate/ds-sparse4d/infer/engine_builder.h

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2025-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 ENGINE_BUILDER_H
#define ENGINE_BUILDER_H
 
#include <string>
#include <fstream>
#include <iostream>
#include <memory>
#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <cuda_runtime.h>
#include "../utils/log_utils.h"
 
 
// Define TensorRT 10.3 constants if needed
#define TRT_MAJOR_VERSION 10
#define TRT_MINOR_VERSION 3
#define TRT_PATCH_VERSION 0
 
// Smart pointer types for TensorRT objects to ensure proper cleanup
struct TRTDeleter {
    template <typename T>
    void operator()(T* obj) const {
        if (obj) {
            delete obj;
        }
    }
};
 
template <typename T>
using TRTUniquePtr = std::unique_ptr<T, TRTDeleter>;
 
// Logger for TensorRT with verbosity control
class Logger : public nvinfer1::ILogger {
public:
    // Default constructor sets verbosity to ERROR (only show errors and fatals)
    Logger(Severity verbosity = Severity::kERROR) : mVerbosity(verbosity) {}
 
    void log(Severity severity, const char* msg) noexcept override {
        // Only log messages at or above the configured verbosity level
        if (severity <= mVerbosity) {
            // Add a prefix to identify the source of the message
            const char* severityStr = "";
            switch (severity) {
                case Severity::kINTERNAL_ERROR: severityStr = "[F] "; break;
                case Severity::kERROR:          severityStr = "[E] "; break;
                case Severity::kWARNING:        severityStr = "[W] "; break;
                case Severity::kINFO:           severityStr = "[I] "; break;
                case Severity::kVERBOSE:        severityStr = "[V] "; break;
            }
            std::cerr << severityStr << msg << std::endl;
        }
    }
 
private:
    Severity mVerbosity;
} static logger;
 
bool buildEngineFromOnnx(
    const std::string& onnxFilePath,
    const std::string& engineFilePath,
    int batchSize,
    int numSensors,
    bool fp16Mode,
    PrettyLogger* logger = nullptr
);
 
#endif // ENGINE_BUILDER_H