sources/gst-plugins/gst-nvdsvideotemplate/ds-sparse4d/postprocess/instance_bank.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 INSTANCE_BANK_H
#define INSTANCE_BANK_H
 
#include <torch/torch.h>
 
#include "../utils/log_utils.h"
#include "../utils/tensor.h"
 
class AnchorHandler {
public:
    virtual ~AnchorHandler() = default;
 
    virtual std::vector<torch::Tensor> anchor_projection(
        const torch::Tensor& anchors,
        const std::vector<torch::Tensor>& transforms,
        const std::vector<torch::Tensor>& time_intervals) = 0;
};
 
// Type definitions for metadata structures
// Use std::map for now for simplicity since it works
typedef std::map<std::string, torch::jit::IValue> MetaData;
typedef std::map<std::string, torch::Tensor> DNMetaData;
 
class InstanceBank : public torch::nn::Module {
public:
    InstanceBank(
        int num_anchor,
        int embed_dims,
        const torch::Tensor& anchor,
        AnchorHandler* anchor_handler = nullptr,
        int num_temp_instances = 0,
        float default_time_interval = 0.5,
        float confidence_decay = 0.6,
        bool anchor_grad = true,
        bool feat_grad = true,
        float max_time_interval = 2.0,
        int reid_dims = -1,
        int gpu_id = -1);
 
    void init_weight();
 
    void reset();
 
    std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
    get(int batch_size, const MetaData& metas, const DNMetaData* dn_metas = nullptr);
 
    std::tuple<torch::Tensor, torch::Tensor>
    update(torch::Tensor instance_feature, torch::Tensor anchor, torch::Tensor confidence);
 
    void cache(
        torch::Tensor instance_feature,
        torch::Tensor anchor,
        torch::Tensor confidence,
        const MetaData* metas = nullptr,
        const std::vector<torch::Tensor>* feature_maps = nullptr);
 
    torch::Tensor get_instance_id(
        torch::Tensor confidence,
        torch::Tensor anchor = torch::Tensor(),
        float threshold = -1.0);
 
    void update_instance_id(
        torch::Tensor instance_id_tensor,
        torch::Tensor confidence = torch::Tensor());
 
    torch::Tensor getMask() const {
        return m_mask;
    }
 
private:
    // Parameters
    int m_embed_dims;
    int m_reid_dims;
    int m_num_anchor;
    int m_num_temp_instances;
    float m_default_time_interval;
    float m_confidence_decay;
    float m_max_time_interval;
    AnchorHandler* m_anchor_handler;
 
    // Registered parameters
    torch::Tensor m_anchor;
    torch::Tensor m_anchor_init;
    torch::Tensor m_instance_feature;
    torch::Tensor m_reid_feature;
 
    // Internal state
    torch::Tensor m_cached_feature;
    torch::Tensor m_cached_anchor;
    MetaData m_metas;
    torch::Tensor m_mask;
    torch::Tensor m_confidence;
    torch::Tensor m_temp_confidence;
    torch::Tensor m_instance_id;
    int64_t m_prev_id;
 
    // Device to use for computations
    torch::Device m_device;
};
 
torch::Tensor load_anchor(const std::string& filename);
 
#endif // INSTANCE_BANK_H