NNData.hpp

Go to the documentation of this file.

#pragma once
 
#include <chrono>
#include <limits>
#include <unordered_map>
#include <vector>
 
#include "Buffer.hpp"
#include "depthai-shared/datatype/RawNNData.hpp"
 
namespace dai {
 
class NNData : public Buffer {
    static constexpr int DATA_ALIGNMENT = 64;
    std::shared_ptr<RawBuffer> serialize() const override;
    RawNNData& rawNn;
 
    // store the data
    // uint8_t
    std::unordered_map<std::string, std::vector<std::uint8_t>> u8Data;
    // FP16
    std::unordered_map<std::string, std::vector<std::uint16_t>> fp16Data;
 
   public:
    NNData();
    explicit NNData(std::shared_ptr<RawNNData> ptr);
    virtual ~NNData() = default;
 
    // Expose
    // uint8_t
    NNData& setLayer(const std::string& name, std::vector<std::uint8_t> data);
 
    NNData& setLayer(const std::string& name, const std::vector<int>& data);
 
    // fp16
    NNData& setLayer(const std::string& name, std::vector<float> data);
 
    NNData& setLayer(const std::string& name, std::vector<double> data);
 
    // getters
    std::vector<std::string> getAllLayerNames() const;
 
    std::vector<TensorInfo> getAllLayers() const;
 
    bool getLayer(const std::string& name, TensorInfo& tensor) const;
 
    bool hasLayer(const std::string& name) const;
 
    bool getLayerDatatype(const std::string& name, TensorInfo::DataType& datatype) const;
 
    // uint8
    std::vector<std::uint8_t> getLayerUInt8(const std::string& name) const;
 
    // fp16
    std::vector<float> getLayerFp16(const std::string& name) const;
 
    // int32
    std::vector<std::int32_t> getLayerInt32(const std::string& name) const;
 
    // first layer
    std::vector<std::uint8_t> getFirstLayerUInt8() const;
 
    std::vector<float> getFirstLayerFp16() const;
 
    std::vector<std::int32_t> getFirstLayerInt32() const;
 
    NNData& setTimestamp(std::chrono::time_point<std::chrono::steady_clock, std::chrono::steady_clock::duration> timestamp);
 
    NNData& setTimestampDevice(std::chrono::time_point<std::chrono::steady_clock, std::chrono::steady_clock::duration> timestamp);
 
    NNData& setSequenceNum(int64_t sequenceNum);
};
 
}  // namespace dai