sources/libs/nvmultiobjecttracker/include/modules/Segmenter/Segmenter.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 SEGMENTERLIB_SEGMENTER_HPP_
#define SEGMENTERLIB_SEGMENTER_HPP_
 
#include <deque>
#include <map>
#include <set>
 
#ifdef VIZ_DEBUG_SEG
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#endif
 
#include <cuda_runtime_api.h>
#include <cuda_runtime.h>
#include "cuda.h"
#include <NvInfer.h>
#include <NvOnnxParser.h>
#include "modules/Segmenter/ISegmenter.hpp"
#include "modules/FrameTransform/trtModule.hpp"
#include "modules/ReID/ReIDFrameTransformTexture.hpp"
#include "modules/ReID/ReIDCropScaler.hpp"
 
enum SAM_POINT_TYPE : int
{
    BACKGROUND = 0, 
    FOREGROUND = 1, 
    TOP_LEFT = 2, 
    BOTTOM_RIGHT = 3 
};
 
typedef std::shared_ptr<std::vector<float>> TensorPtr;
 
struct SegmentMemoryInfo
{
    TensorPtr pred_masks; 
    TensorPtr pred_masks_high_res; 
    TensorPtr obj_ptr; 
    TensorPtr maskmem_features; 
    TensorPtr maskmem_pos_enc; 
    TensorPtr temporal_code; 
};
 
struct MemoryEncoderStatus
{
    TensorPtr maskmem_features; 
    TensorPtr maskmem_pos_enc; 
    TensorPtr temporal_code; 
};
 
struct InferenceStatus
{
    NvMOTFrameNum currentFrameNum; 
    std::vector<TensorPtr> objPtrCond; 
    std::vector<MemoryEncoderStatus> statusCond; 
    std::vector<TensorPtr> objPtrUncond; 
    std::vector<MemoryEncoderStatus> statusUncond; 
};
 
struct CudaInferenceStatus
{
    NvMOTFrameNum currentFrameNum; 
    int numStatusUncond; 
    int numObjPtrUncond; 
};
 
struct SegmenterRunTimeParams
{
    int featuresEncoderSize; 
    int highResFeaturesEncoder0Size; 
    int highResFeaturesEncoder1Size; 
    int pixFeatWithMemSize; 
    int maxConcatMemFeatSize; 
    int maskMemFeatSize; 
    int temporalCodeSize; 
    int objPtrSize; 
    int basicMaskSize; 
    int highResMaskSize; 
    int maxPointNum; 
    int memChannel; 
    int batchSize; 
};
 
class SegmenterConfigParams : public NvTrackerParams
{
public:
    SegmenterConfigParams();
 
    SegmenterConfigParams(const std::string &configFilePath, const std::string moduleName);
 
    SegmenterConfigParams(const YAML::Node &fn);
 
    void read(const YAML::Node &fn) override;
 
    SEGMENTER_TYPES m_segmenterType; 
    int m_hiddenDim; 
    std::vector<int> m_inferDims; 
    int m_inputOrder; 
    int m_encoderFeatureDim; 
    int m_maxMemVecSize_legacy; 
    int m_width; 
    int m_height; 
    int m_channel; 
    int m_gpuId; 
    float m_netScaleFactor; 
    std::vector<float> m_offsets; 
    NvBufSurfaceColorFormat m_colorFormat; 
    int m_promptType; 
    int m_outputLabelNum; 
    bool m_useVPICropScaler; 
    std::string m_vizDir; 
    bool m_saveFrameViz; 
    bool m_saveIDMask; 
    bool m_outputMask; 
    bool m_useMemBank; 
    bool m_useCudaMemBank; 
    bool m_updateMemBankPrompt; 
    int m_numMaskMem; 
    int m_debugVerboseLevel; 
    float m_promptBboxPaddingScale; 
    int m_numObjPtr; 
    int m_numConditionalMem; 
    int m_maskMemoryUpdateInterval; 
    float m_minIouDiff4MemoryUpdate; 
    float m_minIouPrevFrame4ProgBbox; 
    float m_minMemoryUpdateConfidence; 
    int m_maskPromptUpdateInterval; 
    float m_minPromptUpdateConfidence; 
 
private:
    std::string m_trackerConfigPathForPathResolve; 
};
 
class SegmenterCudaMemBank
{
public:
    SegmenterCudaMemBank();
 
    ~SegmenterCudaMemBank();
 
    void initialize(
        const SourceFrameInfo &frameInfo, const uint maxTargetsPerStream,
        const SegmenterRunTimeParams &paramsRt, const SegmenterConfigParams &params, cudaStream_t& cuStream
    );
 
    void update(
        float* maskMemFeatures_d, float* maskMemPosEnc_d, float* temporalCode_d, float* objPtr_d,
        std::shared_ptr<NvTrackedObject> target, int indInBatch, CudaInferenceStatus& inferStatus
    );
 
    void reset(const TrackerIndex targetInd);
 
    std::vector<float*> getCudaMemBank(const TrackerIndex targetInd);
 
private:
    bool m_isInit; 
    cudaStream_t m_cuStream; 
    SegmenterRunTimeParams m_paramsRt; 
    SegmenterConfigParams m_params; 
    float* m_targetsMaskMemFeatures_d; 
    float* m_targetsMaskMemPosEnc_d; 
    float* m_targetsTemporalCode_d; 
    float* m_targetsObjPtr_d; 
    float* m_work_0_d; 
 
    void updateTensorBank(
        float* srcBatch, float* dst, int& numFilled, const int maxNumPerTarget, const int numElementsPerTensor,
        const TrackerIndex targetInd, const int indInBatch, const bool keepFirst=false
    );
};
 
class Segmenter: public ISegmenter
{
public:
    Segmenter(const std::string &configFilePath);
 
    Segmenter(const std::string &configFilePath, const std::string &modelName);
 
    virtual ~Segmenter();
 
    void initialize(
        const SourceFrameInfo &frameInfo, const uint maxTargetsPerStream,
        const std::shared_ptr<NvTrackedObjectManager> &targetManager
    ) override;
 
    void queryParams(NvMOTQuery &paramsQuery) override;
 
    void localize(const std::map<NvMOTStreamId, NvMOTFrame *>& inputMeta) override;
 
    void update(const std::map<NvMOTStreamId, NvMOTFrame *>& inputMeta) override;
 
    void outputMaskObj(const std::shared_ptr<NvTrackedObject> target, NvMOTMask& mask) override;
 
    void removeStream(const NvMOTStreamId videoStreamId) override;
 
    bool isInitialized() override { return m_isInit; }
 
    const SEGMENTER_TYPES getSegmenterType() override { return m_params.m_segmenterType; }
 
protected:
    bool m_isInit; 
    bool m_isForLocalization; 
    SegmenterConfigParams m_params; 
    SegmenterRunTimeParams m_paramsRt; 
    float *m_directly_add_no_mem_embed_d; 
    cudaStream_t m_cuStream; 
    cublasHandle_t m_cublasHandle; 
 
    // [TRT inference modules]
    TrtModule m_imageEncoder; 
    TrtModule m_maskDecoder; 
    TrtModule m_memoryAttention; 
    TrtModule m_memoryEncoder; 
 
    // [Image encoder]
    float *m_inputBuf_d, *m_inputBuf_h; 
    float *m_featuresEncoder_d, *m_highResFeaturesEncoder0_d, *m_highResFeaturesEncoder1_d; 
    float *m_visionPosEmbed_d, *m_visionFeats_d, *m_highResImgPosEmbed1_d; 
 
    // [Memory attention]
    float *m_work_0_d, *m_work_1_d, *m_work_2_d; 
    float *m_memory_obj_ptr_d, *m_memory_maskmem_features_d, *m_memory_d, *m_memory_pos_embed_d, *m_cond_frame_id_diff_d; 
    float *m_pix_feat_with_mem_d; 
    float *m_batchPixFeatWithMem_d; 
 
    // [Mask decoder]
    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_maskForMem_d; 
    float *m_obj_ptr_d; 
    float *m_outputIou_d; 
    float *m_outputOccLogit_d; 
    std::vector<float> m_outputIou_h; 
    std::vector<float> m_outputOccLogit_h; 
    int *m_outputSize_d; 
 
    // [Memory encoder]
    float *m_maskmem_features_d, *m_maskmem_pos_enc_d, *m_temporal_code_d; 
 
    // [Mask post-processing]
    float *m_batchMaskDecSize_d; 
    float *m_targetMaskOriginalSrcSize_d; 
 
    // [Image size]
    Eigen::Vector2f m_scaleMask2Src; 
    Eigen::Vector2f m_scaleInput2Original; 
    NvMotDataTypes::Size m_srcImageSize; 
    NvMotDataTypes::Size m_originalSrcImageSize; 
 
    // [Target related]
    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; 
    std::map<TrackerIndex, std::map<NvMOTFrameNum, SegmentMemoryInfo>> m_condFrameOutputs; 
    std::map<TrackerIndex, std::map<NvMOTFrameNum, SegmentMemoryInfo>> m_nonCondFrameOutputs; 
    std::map<TrackerIndex, InferenceStatus> m_inferStatus; 
    std::map<TrackerIndex, CudaInferenceStatus> m_cudaInferenceStatus; 
    std::map<TrackerIndex, std::vector<float>> m_targetScores; 
    SegmenterCudaMemBank m_cudaMemBank; 
    std::vector<void *> m_bindingsImgEncoder, m_bindingsMaskDecoder, m_bindingsMemAtten, m_bindingsMemEncoder; 
 
    void doInference(const std::map<NvMOTStreamId, NvMOTFrame *>& inputMeta);
 
    void extractImagePatchesBatch(const std::map<NvMOTStreamId, NvMOTFrame *> &inputMeta);
 
    void fillMemoryConditionedFeatures(TrackerIndex targetInd, int indInBatch);
 
    void fillCudaMemoryConditionedFeatures(TrackerIndex targetInd, int indInBatch);
 
    // [TensorRT inference functions]
    void initializeTrtBindings();
 
    void updateTrtBindingsPerStream(const uint inputRoiInd);
 
    void batchInfer(
        const std::map<uint32_t, std::shared_ptr<NvTrackedObject>> &targetList, bool usePrompt,
        int outputLabelIndMax, const NvMOTFrameNum frameNum, const uint inputRoiInd
    );
 
    void singleInfer(
        const std::map<uint32_t, std::shared_ptr<NvTrackedObject>> &targetList, bool usePrompt,
        int outputLabelIndMax, const NvMOTFrameNum frameNum, const uint inputRoiInd
    );
 
    void imgEncoderInfer();
 
    void maskDecoderInfer(
        bool usePrompt, std::vector<std::shared_ptr<NvTrackedObject>>& targetList,
        int& outputLabelIndMax, const uint inputStreamInd, const NvMOTFrameNum frameNum
    );
 
    void memAttenInfer(const std::shared_ptr<NvTrackedObject> target, const int indInBatch, const int inputStreamInd);
 
    void memAttenInferBatchParser(std::vector<std::shared_ptr<NvTrackedObject>>& targetList, const int inputStreamInd);
 
    void memEncoderInfer(const int batchSizeActual);
 
    void setMaskDecoderPointLabels(int promptType);
 
    void setMaskDecoderSinglePointLabelAsNegative(int pointNumPerTarget, int indInBatch);
 
    void setMaskDecoderPoints(const bool usePrompt, std::vector<std::shared_ptr<NvTrackedObject>> targetList);
 
    void postProcessing(
        const int indInBatch, const std::shared_ptr<NvTrackedObject> target, const NvMOTFrameNum frameNum
    );
 
    void updateInferStatus(
        const int indInBatch, const std::shared_ptr<NvTrackedObject> target, const NvMOTFrameNum frameNum
    );
 
    void removeInferStatus();
 
    void copyFromCuda(float* src, TensorPtr& dst, int numElements, cudaStream_t& cuStream);
 
    inline bool isInferStatusEmpty(TrackerIndex targetInd)
    {
        if (m_params.m_useCudaMemBank)
        {
            if (m_cudaInferenceStatus.find(targetInd) != m_cudaInferenceStatus.end())
            {
                CudaInferenceStatus& targetStatus = m_cudaInferenceStatus.at(targetInd);
                return targetStatus.numStatusUncond == 0;
            }
            else
            {
                return true;
            }
        }
        else
        {
            if (m_inferStatus.find(targetInd) != m_inferStatus.end())
            {
                InferenceStatus& targetStatus = m_inferStatus.at(targetInd);
                return targetStatus.statusCond.empty() && targetStatus.statusUncond.empty();
            }
            else
            {
                return true;
            }
        }
    }
 
    bool skipInfer(std::shared_ptr<NvTrackedObject> target, const NvMOTFrameNum frameNum);
 
    inline bool isGoodBboxProg(
        const std::shared_ptr<NvTrackedObject> target, const float confidence,
        const NvMotDataTypes::Rect2f bboxProg, const float iouThresh=0.5
    );
 
    #ifdef VIZ_DEBUG_SEG
        cv::Mat m_imageSumMask; 
        cv::Mat m_imageIDMask; 
        std::map<NvMOTStreamId, std::map<uint, cv::Mat>> m_vizSrcImages; 
 
        void updatePerTargetMaskInFrame(const std::shared_ptr<NvTrackedObject> target, bool initMask);
 
        void savePerFrameMask(
            NvMOTStreamId videoStreamId, NvMOTFrameNum frameNum,
            const std::map<uint32_t, std::shared_ptr<NvTrackedObject>> &targetList
        );
    #endif
};
 
#endif // SEGMENTERLIB_SEGMENTER_HPP_