videoStreamInfer/src/infer/trt/depth/depth.hpp

#ifndef DEPTH_HPP__
#define DEPTH_HPP__

#include <vector>
#include <iomanip>
#include "common/memory.hpp"
#include "infer/infer.hpp"
#include "common/image.hpp"
#include "common/data.hpp"
#include "infer/trt/affine.hpp"

#ifdef TRT10
#include "common/tensorrt.hpp"
namespace TensorRT = TensorRT10;
#else
#include "common/tensorrt8.hpp"
namespace TensorRT = TensorRT8;
#endif

namespace depth
{

    class DepthModelImpl : public Infer 
    {
    public:
        ModelType model_type_;
    
        std::shared_ptr<TensorRT::Engine> trt_;
        std::string engine_file_;
        
        tensor::Memory<unsigned char> preprocess_buffer_;
    
        tensor::Memory<float> affine_matrix_;
        tensor::Memory<float> invert_affine_matrix_;
        tensor::Memory<float> input_buffer_, output_buffer_;
        tensor::Memory<float> depth_map_buffer_;
    
        int network_input_width_, network_input_height_;
        affine::Norm normalize_;

        bool isdynamic_model_ = false;
    
        DepthModelImpl() = default;
    
        virtual ~DepthModelImpl() = default;
    
        void adjust_memory(int width, int height)
        {
            depth_map_buffer_.gpu(width * height);
            depth_map_buffer_.cpu(width * height);
        }


        void adjust_memory() 
        {
            // the inference batch_size
            size_t input_numel = network_input_width_ * network_input_height_ * 3;
            input_buffer_.gpu(input_numel);

            output_buffer_.gpu(input_numel / 3);
            output_buffer_.cpu(input_numel / 3);

            affine_matrix_.gpu(6);
            affine_matrix_.cpu(6);

            invert_affine_matrix_.gpu(6);
            invert_affine_matrix_.cpu(6);
        }
    
        void preprocess(const tensor::Image &image, affine::LetterBoxMatrix &affine, void *stream = nullptr);
        void postprocess(int width, int height, void *stream = nullptr);
        
    
        bool load(const std::string &engine_file);
    
        virtual Result forward(const tensor::Image &image, void *stream = nullptr);
        virtual Result forward(const tensor::Image &image, int slice_width, int slice_height, float overlap_width_ratio, float overlap_height_ratio, void *stream = nullptr);
    
};

Infer *loadraw(const std::string &engine_file);

std::shared_ptr<Infer> load_depth(const std::string &engine_file, int gpu_id);

}   // namespace depth

#endif