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

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2023-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 ReIDCropScaler_HPP_
#define ReIDCropScaler_HPP_
 
#include <map>
#include <vector>
#include <unordered_map>
#include "nvbufsurface.h"
#include "vpi/Stream.h"
#include "vpi/Image.h"
#include "vpi/Array.h"
#include "vpi/algo/CropScaler.h"
#include "modules/utility/dataTypes.hpp"
#include "modules/NvMultiObjectTracker/nvdstracker.h"
 
#if defined(__aarch64__)
#include <EGL/egl.h>
#include <EGL/eglext.h>
#include "cudaEGL.h"
#endif
 
class ReIDCropScaler
{
public:
    ReIDCropScaler();
 
    ~ReIDCropScaler();
 
    void initialize(
        const int debugVerboseLevel, const int maxStreamNum, const int maxBufSurfAddrSize, const int maxTargetNum,
        const NvMotDataTypes::Size srcImgSize, const uint srcImgPitch, const NvBufSurfaceColorFormat colorFormatType,
        const NvMotDataTypes::Size dstImgSize, const NvBufSurfaceColorFormat dstColorFormatType, const int dimOrder
    );
 
    bool frameTransformInBatch(
        const std::map<NvMOTStreamId, std::vector<std::pair<uint, NvMotDataTypes::Rect2f> > > &roisInBatch,
        const int lastBboxInd, const float outputScaleFactor, const std::vector<float> outputOffsets,
        const std::map<NvMOTStreamId, NvMOTFrame*> &inputMeta
    );
 
    float* getOutputImagePatchPtrs() { return m_odata_d; }
 
    #if defined VIZ_DEBUG_REID || defined VIZ_DEBUG_SEG || defined VIZ_DEBUG_POSE
 
        std::map<NvMOTStreamId, std::map<uint, cv::Mat > > getImagePatches2Host_OCV_batch(
            const float outputScaleFactor, const std::vector<float> outputOffsets
        );
    #endif
 
private:
    void initialize4Output(
        const NvMotDataTypes::Size dstSize, const NvBufSurfaceColorFormat dstColorFormat, const int dimOrder
    );
 
    void allocateOutputMemories();
 
    void deAllocateOutputMemories();
 
    void allocateInputMemoriesBatch(const uint numFramesInBatch, const int nx, const int ny, const uint pitch_in);
 
    void deAllocateInputMemoriesBatch();
 
    void initialize4Input(
        uint numFramesInBatch, const NvMotDataTypes::Size frameSize, const uint srcImgPitch, const uint frameChannels
    );
 
    void checkInput(const std::map<NvMOTStreamId, NvMOTFrame*> &inputMeta);
 
    inline void getCudaMemPointers(NvBufSurfaceParams *baseSurface);
 
    bool m_isInit; 
 
    // [Output configuration specific to ReID]
    int m_width_out, m_height_out, m_channels_out; 
    int m_dimOrder; 
    NvBufSurfaceColorFormat m_dstColorFormatType; 
    VPIImageFormat m_vpiImageFormat; 
    cudaStream_t m_cuStream; 
    const int m_cudaBlockSize_2D = 32; 
    int m_debugVerboseLevel; 
    int m_maxTargetsToTrack; 
    uint m_lastInd; 
    uint m_maxBufSurfAddrSize; 
    uint m_numFramesInBatchMax; 
 
    // [Input image format]
    int m_width_in, m_height_in, m_pitch_in, m_channels_in; 
    NvBufSurfaceColorFormat m_colorFormatType; 
 
    // [Output memory specific to ReID]
    float *m_odata_d; 
    float *m_odata_host; 
    #if defined VIZ_DEBUG_REID || defined VIZ_DEBUG_SEG || defined VIZ_DEBUG_POSE
        float *m_odata_host_nchw; 
    #endif
    uint8_t *m_odata_unit8_d; 
    VPIImageData m_vpiObjPatchesData; 
    VPIImage m_vpiObjPatches; 
 
    // [Getting CUDA pointers from NvBufSurfaces]
    int m_integrated_gpu; 
    int m_deviceId; 
    uint8_t *m_surfCudaMemPtr[NVBUF_MAX_PLANES]; 
    #if defined(__aarch64__)
        std::unordered_map<const NvBufSurfaceParams *, CUgraphicsResource> m_graphicsResourceAddrMap; 
        std::unordered_map<const NvBufSurfaceParams *, CUeglFrame> m_eglFrameAddrMap; 
    #endif
 
    // [VPI]
    VPIStream m_vpiStream; 
    VPIBackend m_vpiBackend4CropnScale = VPI_BACKEND_CUDA; 
    VPIPayload m_vpiCropPayload; 
    int m_vpiMemFlags = VPI_BACKEND_CUDA | VPI_RESTRICT_MEM_USAGE; 
 
    // [Sequence of input frames]
    VPIImage *m_vpiInputFrameSequence; 
    VPIImageData *m_vpiInputFrameDataList; 
    std::vector<bool> m_isFrameCreated; 
 
    // [List of tracked objects]
    std::vector<VPIDCFTrackedBoundingBox> m_vpiDCFBboxes; 
    VPIArrayData m_vpiDCFBboxesData; 
    VPIArray m_vpiDCFBboxesArray; 
    int32_t m_vpiBboxesArraySize; 
 
    #if defined VIZ_DEBUG_REID || defined VIZ_DEBUG_SEG || defined VIZ_DEBUG_POSE
        std::map<uint, std::pair<NvMOTStreamId, uint>> m_targetNumInd2RoisInBatchMapper; 
    #endif
};
 
#endif // ReIDCropScaler_HPP_