sources/libs/nvmultiobjecttracker/include/modules/PoseEstimator/PoseInference.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 POSEESTIMATORLIB_POSEINFERENCE_HPP_
#define POSEESTIMATORLIB_POSEINFERENCE_HPP_
 
#include <deque>
#include <map>
 
#ifdef VIZ_DEBUG_POSE
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#endif
 
#include <NvInfer.h>
#include <NvOnnxParser.h>
#include "modules/PoseEstimator/IPoseEstimator.hpp"
#include "modules/FrameTransform/trtModule.hpp"
#include "modules/ReID/ReIDFrameTransformTexture.hpp"
#include "modules/ReID/ReIDCropScaler.hpp"
 
class PoseInferenceConfigParams : public NvTrackerParams
{
public:
    PoseInferenceConfigParams();
 
    PoseInferenceConfigParams(const std::string &configFilePath, const std::string moduleName);
 
    PoseInferenceConfigParams(const YAML::Node &fn);
 
    void read(const YAML::Node &fn) override;
 
    POSEESTIMATOR_TYPES m_poseEstimatorType; 
    int m_hiddenDim; 
    std::vector<int> m_inferDims; 
    int m_inputOrder; 
    int m_encoderFeatureDim; 
    int m_width; 
    int m_height; 
    int m_channel; 
    int m_maxBatchSize; 
    int m_gpuId; 
    int m_poseInferenceInterval; 
    float m_netScaleFactor; 
    std::vector<float> m_offsets; 
    NvBufSurfaceColorFormat m_colorFormat; 
    bool m_useVPICropScaler; 
    std::string m_vizDir; 
    int m_debugVerboseLevel; 
    bool m_saveFrameViz; 
    std::vector<int> m_operateOnClassIds; 
};
 
class PoseInference: public IPoseEstimator
{
public:
    PoseInference(const std::string &configFilePath);
 
    PoseInference(const std::string &configFilePath, const std::string &modelName);
 
    virtual ~PoseInference();
 
    void initializeRuntime(
        const SourceFrameInfo &frameInfo, const uint maxTargetsPerStream,
        const std::shared_ptr<NvTrackedObjectManager> &targetManager
    ) override;
 
    void initializeModel() override;
 
    void queryParams(NvMOTQuery &paramsQuery);
 
    void doInference(
        const std::map<NvMOTStreamId, NvMOTFrame *>& inputMeta,
        std::map<NvMOTStreamId, std::deque<DetectedObjInfo>> &detectedObjects, const bool isForReAssoc
    ) override;
 
    void update() override;
 
    void removeStream(const NvMOTStreamId videoStreamId) override;
 
    const POSEESTIMATOR_TYPES getPoseEstimatorType() override { return m_params.m_poseEstimatorType; }
 
protected:
    bool m_isRuntimeInit; 
    int m_maxPointNum; 
    int m_numBboxes; 
    float *m_inputBuf_d, *m_inputBuf_h; 
    float *m_points_d; 
    float *m_pointLabels_d; 
    float *m_inputMask_d; 
    float *m_inputMaskLabel_d; 
    float *m_targetMaskDecSize_d, *m_targetMaskDecSize_h; 
    float *m_batchMaskDecSize_d; 
    float *m_targetMaskOriginalSrcSize_d; 
    float *m_kInv_d, *m_tFormInv_d, *m_scaleNormalizedMeanLimbLengths_d, *m_meanLimbLengths_d; 
    float *m_pose2d_d, *m_pose2d_org_img_d, *m_pose25d_d, *m_pose3d_d; 
    float *m_pose25d_h, *m_pose3d_h; 
    float *m_outputIou_d; 
    PoseInferenceConfigParams m_params; 
    NvMotDataTypes::Size m_srcImageSize; 
    NvMotDataTypes::Size m_originalSrcImageSize; 
    std::vector<std::tuple<uint, NvMOTStreamId, NvMotDataTypes::Rect2f>> m_targetStreamIdList; 
 
    cudaStream_t m_cuStream; 
    std::vector<void*> m_bindings; 
    float *m_featuresEncoder_d; 
    std::shared_ptr<NvTrackedObjectManager> m_targetManager; 
    std::shared_ptr<ReIDFrameTransformTexture> m_subImageCropper2; 
    std::shared_ptr<ReIDCropScaler> m_subImageCropScaler; 
    std::map<NvMOTStreamId, std::vector<std::pair<uint, NvMotDataTypes::Rect2f>>> m_inputRoisInBatch; 
    TrtModule m_inference; 
    int m_maxTargetsToTrack; 
 
    void extractImagePatchesBatch(const std::map<NvMOTStreamId, NvMOTFrame *> &inputMeta);
 
    void setValues4PoseGeometryInput();
};
 
#endif // POSEESTIMATORLIB_POSEINFERENCE_HPP_