GreenScreen/pluginProcessor.h

/**
 *  Copyright (C) 2020-2021 Savoir-faire Linux Inc.
 *
 *  Author: Aline Gondim Santos <aline.gondimsantos@savoirfairelinux.com>
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301
 * USA.
 */

#pragma once

#include <map>
#include <vector>
extern "C" {
#include <libavutil/frame.h>
}

#include <plugin/jamiplugin.h>
#include <plugin/mediahandler.h>

#include <opencv2/core.hpp>
#include <onnxruntime_cxx_api.h>
#ifdef NVIDIA
#include <cuda_provider_factory.h>
#endif
#ifdef __ANDROID__
#include <nnapi_provider_factory.h>
#endif

#include <frameFilter.h>
#include <frameUtils.h>
#include <mediaStream.h>
#include <functional>

namespace jami {

static const char* modelInputNames[8] = {"image:0"};
static const char* modelOutputNames[11] = {"Identity:0"};

class PluginProcessor
{
public:
    PluginProcessor(const std::string& model, bool acc);
    ~PluginProcessor();

    void initModel(const std::string& modelPath, bool activateAcc);
    /**
     * @brief feedInput
     * Takes a frame and feeds it to the model storage for predictions
     * @param frame
     */
    void feedInput(AVFrame* input);

    /**
     * @brief computePredictions
     * Uses the model to compute the predictions and store them in
     * computedPredictions
     */
    void computePredictions();

    void printMask();
    void drawMaskOnFrame(AVFrame* frame, AVFrame* frameReduced, int angle);
    bool isAllocated() { return isAllocated_; }
    void setBlur(bool isBlur) { isBlur_ = isBlur; }
    void setBlurLevel(const std::string& blurLevel) { blurLevel_ = blurLevel; }
    void setBackgroundImage(const std::string& backgroundPath) { backgroundPath_ = backgroundPath; }
    void initFilters(const std::pair<int, int>& inputSize, int format, int angle);

    std::pair<int, int> modelInputDimensions {257, 257};
    std::map<int, std::string> rotation = {{90, "-PI/2"},
                                           {-90, "PI/2"},
                                           {-180, "-PI"},
                                           {180, "PI"},
                                           {0, "0"}};

private:
    void resetInitValues();
    void loadBackground();
    MediaStream getbgAVFrameInfos();

    bool isBlur_ {false};
    std::string blurLevel_;
    std::string backgroundPath_;
    cv::Mat previousMasks_[2];
    std::vector<float> computedMask_;
    cv::Mat cvFrame_;

    // process variables
    cv::Size kSize_;
    int count_ {0};
    cv::Mat bgdModel_, fgdModel_;
    int grabCutMode_ {1}; // cv::GC_INIT_WITH_MASK = 1;
    int grabCutIterations_ {4};
    int grabcutClass_ {3};
    int frameCount_ {5};
    float smoothFactors_[3] = {0.6f, 0.3f, 0.1f};
    float kernelSize_ {0.05f};

    // filters
    std::string mainFilterDescription_;
    FrameFilter mainFilter_;
    std::unique_ptr<AVFormatContext, std::function<void(AVFormatContext*)>> pFormatCtx_
        = {avformat_alloc_context(), [](AVFormatContext* ptr) {
               avformat_close_input(&ptr);
               avformat_free_context(ptr);
           }};
    int videoStream_ {-1};
    MediaStream ims_, ims2_, maskms_;

    // onnx related
    bool isAllocated_ {false};
    Ort::Env env_ {ORT_LOGGING_LEVEL_WARNING, "test"};
    Ort::Session* session_ {};
    Ort::SessionOptions sessOpt_;

    Ort::Value input_tensor_ {nullptr};
    std::array<int64_t, 3> input_shape_ {257, 257, 3};
    std::array<float, 257 * 257 * 3> input_image_ {};

    Ort::Value output_tensor_ {nullptr};
    std::array<int64_t, 4> output_shape_ {1, 17, 17, 1};
    std::array<float, 17 * 17> results_ {};
};
} // namespace jami