sources/libs/nvmultiobjecttracker/include/modules/ReID/ReID.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2021-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 REIDLIB_REID_HPP_
#define REIDLIB_REID_HPP_
 
#include <map>
#include <iterator>
#include <memory>
 
#ifdef VIZ_DEBUG_REID
#include <opencv2/core.hpp>
#endif
 
#include <NvInfer.h>
#include <NvOnnxParser.h>
#include "modules/ReID/IReID.hpp"
#include "modules/ReID/ReIDFeatureGallery.hpp"
#include "modules/ReID/ReIDFrameTransformTexture.hpp"
#include "modules/ReID/ReIDCropScaler.hpp"
#include "modules/NvMultiObjectTracker/NvMultiObjectTracker_common.hpp"
#include "modules/NvMultiObjectTracker/NvTrackerParams.hpp"
#include "modules/FrameTransform/trtModule.hpp"
 
class ReIDConfigParams : public NvTrackerParams
{
public:
    ReIDConfigParams();
 
    ReIDConfigParams(const std::string &configFilePath, const std::string moduleName, const std::string submoduleName);
 
    ReIDConfigParams(const YAML::Node &fn);
 
    void read(const YAML::Node &fn) override;
 
    bool m_keepAspc; 
    bool m_addFeatureNormalization; 
    bool m_outputReidTensor; 
    int m_maxBatchSize; 
    int m_reidFeatureSize; 
    int m_inputOrder; 
    int m_width; 
    int m_height; 
    int m_channel; 
    int m_reidHistorySize; 
    int m_gpuId; 
    float m_netScaleFactor; 
    float m_minVisibility4GalleryUpdate; 
    float m_minVisibility4ReidOutput; 
    size_t m_workspaceSize; 
    std::vector<int> m_inferDims; 
    std::vector<float> m_offsets; 
    NvBufSurfaceColorFormat m_colorFormat; 
    REID_TYPES m_reidType; 
    bool m_useVPICropScaler; 
    bool m_useSegmentationMask4Input; 
    int m_debugVerboseLevel; 
    #ifdef VIZ_DEBUG_REID
        bool m_useExternalImages; 
        std::string m_extImagePath; 
    #endif
};
 
class ReID : public IReID
{
public:
 
    ReID(const std::string &configFilePath);
 
    ReID(const std::string &configFilePath, const std::string &modelName, const bool calDisMaha=true);
 
    virtual ~ReID();
 
    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;
 
    void outputReIDTensor(NvDsReidTensorBatch& reidTensorBatch) override;
 
    int outputTargetReIDInd(NvMOTStreamId streamId, uint targetInd) override;
 
    void outputTargetReIDObj(
        NvMOTStreamId streamId, uint targetInd, NvDsReidTensorBatch& reidTensorBatch, NvDsObjReid& reidObj
    ) override;
 
    #ifdef VIZ_DEBUG_REID
 
        void vizImagePatches(
            const std::map<NvMOTStreamId, std::map<uint, cv::Mat > >& subImgs,
            const std::map<NvMOTStreamId, NvMOTFrame *> &inputMeta
        );
 
        void loadExtImagePatchesBatch();
    #endif
 
    void getFeatures2Host();
 
    bool isModelInitialized() override { return m_isModelInit; }
 
    bool isRuntimeInitialized() override { return m_isRuntimeInit; }
 
    const REID_TYPES getReIDType() override { return m_params.m_reidType; }
 
protected:
    bool m_isModelInit; 
    bool m_isRuntimeInit; 
    ReIDConfigParams m_params; 
    TrtModule m_inference; 
    std::shared_ptr<ReIDFeatureGallery> m_reidFeatureGallery; 
    std::shared_ptr<NvTrackedObjectManager> m_targetManager; 
    // Debug timers.
    #ifdef VIZ_DEBUG_REID
        int m_extFrameNum; 
        double m_timeTotal; 
        int m_timeCnt; 
    #endif
    bool m_calDisMaha; 
    int m_maxTargetsToTrack; 
    int m_maxTargetsPerStream; 
    int m_maxStreams; 
    int m_numBboxes; 
    NvMotDataTypes::Size m_srcImageSize; 
    std::map<NvMOTStreamId, std::vector<std::pair<uint, NvMotDataTypes::Rect2f>>> m_inputRoisInBatch; 
    std::shared_ptr<ReIDFrameTransformTexture> m_subImageCropper2; 
    std::shared_ptr<ReIDCropScaler> m_subImageCropScaler; 
    std::map<NvMOTStreamId, std::pair<float*, std::map<int, int>>> m_detObjReidFeatures; 
    std::map<NvMOTStreamId, std::map<int, int>> m_featureIndFrameNumMap; 
 
    // [ReID network parameters]
    int m_inputInd; 
    int m_outputInd; 
    float *m_reidFeaturesDet_d, *m_reidFeaturesDet_h; 
    float *m_inputBuf_d, *m_inputBuf_h; 
    std::vector<void*> m_bindings; 
    cudaStream_t m_cuStream; 
 
    // [ReID feature similarity computation]
    int m_maxReidProductSizePerStream; 
    float *m_reidProduct_d, *m_reidProduct_h; 
    int *m_targetReidFeatureInds_d, *m_targetReidFeatureInds_h; 
    int2 *m_reidProductSize_d, *m_reidProductSize_h; 
    int *m_targetWork_1_d, *m_targetWork_1_h; 
    int *m_targetWork_2_d, *m_targetWork_2_h; 
    int *m_targetWork_3_d, *m_targetWork_3_h; 
    float *m_confidenceReidDetTargets_d, *m_confidenceReidDetTargets_h; 
    float4 *m_inputRoisInBatch_d, *m_inputRoisInBatch_h; 
    int *m_targetInds4Rois_d, *m_targetInds4Rois_h; 
    int *m_nearestIndReidDetTargets_d, *m_nearestIndReidDetTargets_h; 
    int m_normFeatureCudaBlockSize; 
    cublasHandle_t m_cublasHandle; 
 
    // [Mahalanobis distance thread]
    std::shared_ptr<NvTrackerMultiThreadWorker<DisMahalanobisIn>> m_disMahaThreadWorker; 
 
    void extractImagePatchesBatch(
        const std::map<NvMOTStreamId, NvMOTFrame *>& inputMeta,
        std::map<NvMOTStreamId, std::deque<DetectedObjInfo>> &detectedObjects
    );
 
    void generateImagePatchesWithSegmentationMask(
        std::map<NvMOTStreamId, std::deque<DetectedObjInfo>> &detectedObjects
    );
 
    void removeTrackerReIDFeature(const NvMOTStreamId videoStreamId, const std::shared_ptr<NvTrackedObject> target);
 
    void computeReIDFeatureDist(const std::map<NvMOTStreamId, NvMOTFrame *> &inputMeta);
 
    void multiplyReIDFeatures(const int streamInd, const NvMOTStreamId videoStreamId);
 
    void normalizaReIDFeatures();
 
    void computeMahalanobisDist(DisMahalanobisIn in);
};
 
#endif // REIDLIB_REID_HPP_