infer_iprocess.h

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2018-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: LicenseRef-NvidiaProprietary
 *
 * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
 * property and proprietary rights in and to this material, related
 * documentation and any modifications thereto. Any use, reproduction,
 * disclosure or distribution of this material and related documentation
 * without an express license agreement from NVIDIA CORPORATION or
 * its affiliates is strictly prohibited.
 */
 
#ifndef __NVDSINFERSERVER_IPROCESS_H__
#define __NVDSINFERSERVER_IPROCESS_H__
 
#include <stdarg.h>
#include <condition_variable>
#include <functional>
#include <list>
#include <memory>
#include <mutex>
#include <queue>
#include <tuple>
 
#include "infer_common.h"
#include "infer_datatypes.h"
#include "infer_utils.h"
 
namespace nvdsinferserver {
 
class IPreprocessor {
public:
    using PreprocessDone =
        std::function<void(NvDsInferStatus, SharedBatchArray)>;
 
    IPreprocessor() = default;
    virtual ~IPreprocessor() = default;
 
    virtual NvDsInferStatus transform(SharedBatchArray src,
        SharedCuStream mainStream, PreprocessDone done) = 0;
 
private:
    DISABLE_CLASS_COPY(IPreprocessor);
};
 
class IPostprocessor {
public:
    using PostprocessDone =
        std::function<void(NvDsInferStatus, SharedBatchArray)>;
    IPostprocessor() = default;
    virtual ~IPostprocessor() = default;
 
    virtual NvDsInferStatus postCudaProcess(SharedBatchArray inBuf,
        SharedCuStream mainStream, PostprocessDone done) = 0;
 
    virtual NvDsInferStatus postHostProcess(SharedBatchArray inBuf,
        SharedCuStream mainStream, PostprocessDone done) = 0;
 
private:
    DISABLE_CLASS_COPY(IPostprocessor);
};
 
class BasePreprocessor : public IPreprocessor {
public:
    ~BasePreprocessor() override = default;
 
    /*
     * @brief Helper functions to access member variables.
     */
    void setUniqueId(int id) { m_UniqueId = id; }
    int uniqueId() const { return m_UniqueId; }
    void setTransformIdx(int idx) { m_TransformIdx = idx; }
    NvDsInferStatus transform(SharedBatchArray src, SharedCuStream mainStream,
        PreprocessDone done) override {
        assert(src && (int)src->getSize() > m_TransformIdx);
        SharedBatchBuf& from = src->buf(m_TransformIdx);
        SharedBatchBuf to = requestOutBuffer(from);
        if (!to) {
            return NVDSINFER_RESOURCE_ERROR;
        }
        to->setBufId(src->bufId());
        NvDsInferStatus status = transformImpl(from, to, mainStream);
        src->buf(m_TransformIdx) = to;
        mainStream.reset();
        if (status == NVDSINFER_SUCCESS) {
            done(status, std::move(src));
        }
        return status;
    }
 
    virtual NvDsInferStatus allocateResource(
        const std::vector<int>& devIds) = 0;
 
private:
    virtual SharedBatchBuf requestOutBuffer(SharedBatchBuf& inBuf) = 0;
    virtual NvDsInferStatus transformImpl(SharedBatchBuf& src,
        SharedBatchBuf& dst, SharedCuStream& mainStream) = 0;
 
private:
    int m_TransformIdx = 0;
    int m_UniqueId = 0;
};
 
class BasePostprocessor : public IPostprocessor {
public:
    BasePostprocessor(InferPostprocessType type, int uid)
        : m_UniqueId(uid), m_ProcessType(type) {}
    ~BasePostprocessor() override = default;
 
    /*
     * @brief Helper functions to access member variables.
     */
    void setUniqueId(int id) { m_UniqueId = id; }
    int uniqueId() const { return m_UniqueId; }
    InferPostprocessType networkType() const { return m_ProcessType; }
    NvDsInferStatus postCudaProcess(SharedBatchArray inBuf,
        SharedCuStream mainStream, PostprocessDone done) override {
        return genericPostProcess_(std::move(inBuf), std::move(mainStream),
            std::move(done), &BasePostprocessor::requestCudaOutBufs,
            &BasePostprocessor::postCudaImpl);
    }
 
    NvDsInferStatus postHostProcess(SharedBatchArray inBuf,
        SharedCuStream mainStream, PostprocessDone done) override {
        return genericPostProcess_(std::move(inBuf), std::move(mainStream),
            std::move(done), &BasePostprocessor::requestHostOutBufs,
            &BasePostprocessor::postHostImpl);
    }
 
    virtual NvDsInferStatus allocateResource(
        const std::vector<int>& devIds) = 0;
 
private:
    virtual NvDsInferStatus postCudaImpl(SharedBatchArray& inBuf,
        SharedBatchArray& outbuf, SharedCuStream& mainStream) = 0;
    virtual NvDsInferStatus postHostImpl(SharedBatchArray& inBuf,
        SharedBatchArray& outbuf, SharedCuStream& mainStream) = 0;
 
    virtual SharedBatchArray requestCudaOutBufs(const SharedBatchArray& inBuf) = 0;
    virtual SharedBatchArray requestHostOutBufs(const SharedBatchArray& inBuf) = 0;
 
    template <typename RequestBuf, typename Impl>
    NvDsInferStatus genericPostProcess_(SharedBatchArray inBuf,
        SharedCuStream mainStream, PostprocessDone done, RequestBuf reqBuf,
        Impl impl) {
        assert(inBuf);
        SharedBatchArray outBuf = (this->*reqBuf)(inBuf);
        if (!outBuf) {
            return NVDSINFER_UNKNOWN_ERROR;
        }
        outBuf->setBufId(inBuf->bufId());
        NvDsInferStatus status = (this->*impl)(inBuf, outBuf, mainStream);
        inBuf.reset();
        mainStream.reset();
        if (status == NVDSINFER_SUCCESS) {
            done(status, std::move(outBuf));
        }
        return status;
    }
 
private:
    int m_UniqueId = 0;
    InferPostprocessType m_ProcessType = InferPostprocessType::kOther;
};
 
template <class BasePreprocessorT>
class ThreadPreprocessor : public BasePreprocessorT {
private:
    using Item = std::tuple<SharedBatchArray, SharedCuStream,
        IPreprocessor::PreprocessDone>;
    enum { kBuf, kStream, kDone };
 
public:
    template <typename... Args>
    ThreadPreprocessor(Args&&... args)
        : BasePreprocessorT(std::forward<Args>(args)...),
          m_Worker(
              [this](Item i) -> bool {
                  IPreprocessor::PreprocessDone done =
                      std::move(std::get<kDone>(i));
                  NvDsInferStatus s = BasePreprocessorT::transform(
                      std::move(std::get<kBuf>(i)),
                      std::move(std::get<kStream>(i)), done);
                  if (s != NVDSINFER_SUCCESS) {
                      done(s, nullptr);
                  }
                  return true;
              },
              "Preproc")
    {
    }
    ~ThreadPreprocessor() override { m_Worker.join(); }
    void setThreadName(const std::string &name) {
        m_Worker.setThreadName(name);
    }
    NvDsInferStatus transform(SharedBatchArray src, SharedCuStream mainStream,
        IPreprocessor::PreprocessDone done) final {
        Item item = std::make_tuple(
            std::move(src), std::move(mainStream), std::move(done));
        if (m_Worker.queueItem(std::move(item))) {
            return NVDSINFER_SUCCESS;
        }
        return NVDSINFER_RESOURCE_ERROR;
    }
 
private:
    QueueThread<std::vector<Item>> m_Worker;
};
 
template <class BasePostprocessorT>
class ThreadCudaPostprocessor : public BasePostprocessorT {
    using ItemCuda = std::tuple<SharedBatchArray, SharedCuStream,
        IPostprocessor::PostprocessDone>;
    enum { kBuf, kStream, kDone };
 
public:
    template <typename... Args>
    ThreadCudaPostprocessor(Args&&... args)
        : BasePostprocessorT(std::forward<Args>(args)...),
          m_CudaWorker(
              [this](ItemCuda i) -> bool {
                  IPostprocessor::PostprocessDone done =
                      std::move(std::get<kDone>(i));
                  NvDsInferStatus s = BasePostprocessorT::postCudaProcess(
                      std::move(std::get<kBuf>(i)),
                      std::move(std::get<kStream>(i)), done);
                  if (s != NVDSINFER_SUCCESS) {
                      done(s, nullptr);
                  }
                  return true;
              },
              "PostCuda")
    {}
    ~ThreadCudaPostprocessor() { m_CudaWorker.join(); }
    NvDsInferStatus postCudaProcess(SharedBatchArray inBuf,
        SharedCuStream mainStream, IPostprocessor::PostprocessDone done) final {
        ItemCuda item = std::make_tuple(
            std::move(inBuf), std::move(mainStream), std::move(done));
        if (m_CudaWorker.queueItem(std::move(item))) {
            return NVDSINFER_SUCCESS;
        }
        return NVDSINFER_RESOURCE_ERROR;
    }
 
private:
    QueueThread<std::vector<ItemCuda>> m_CudaWorker;
};
 
template <typename BasePostprocessorT>
class ThreadHostPostprocessor : public BasePostprocessorT {
    using ItemHost = std::tuple<SharedBatchArray, SharedCuStream,
        IPostprocessor::PostprocessDone>;
    enum { kBuf, kStream, kDone };
 
public:
    template <typename... Args>
    ThreadHostPostprocessor(Args&&... args)
        : BasePostprocessorT(std::forward<Args>(args)...),
          m_HostWorker(
              [this](ItemHost i) -> bool {
                  IPostprocessor::PostprocessDone done =
                      std::move(std::get<kDone>(i));
                  NvDsInferStatus s = BasePostprocessorT::postHostProcess(
                      std::move(std::get<kBuf>(i)),
                      std::move(std::get<kStream>(i)), done);
                  if (s != NVDSINFER_SUCCESS) {
                      done(s, nullptr);
                  }
                  return true;
              },
              "PostHost")
    {
    }
 
    ~ThreadHostPostprocessor() { m_HostWorker.join(); }
 
    NvDsInferStatus postHostProcess(SharedBatchArray inBuf,
        SharedCuStream mainStream, IPostprocessor::PostprocessDone done) final {
        ItemHost item = std::make_tuple(
            std::move(inBuf), std::move(mainStream), std::move(done));
        if (m_HostWorker.queueItem(std::move(item))) {
            return NVDSINFER_SUCCESS;
        }
        return NVDSINFER_RESOURCE_ERROR;
    }
 
private:
    QueueThread<std::vector<ItemHost>> m_HostWorker;
};
 
} // namespace nvdsinferserver
 
#endif