GreenScreen: blur functionality

GitLab: #15
Change-Id: Iafc8b410ebe0228eca9b780504382347b38fd8cc
diff --git a/GreenScreen/pluginProcessor.cpp b/GreenScreen/pluginProcessor.cpp
index 62de879..ff4955b 100644
--- a/GreenScreen/pluginProcessor.cpp
+++ b/GreenScreen/pluginProcessor.cpp
@@ -20,202 +20,161 @@
  */
 
 #include "pluginProcessor.h"
-// System includes
-#include <algorithm>
-#include <cstring>
-// OpenCV headers
-#include <opencv2/core.hpp>
-#include <opencv2/imgcodecs.hpp>
-#include <opencv2/imgproc.hpp>
-// Logger
-#include <pluglog.h>
 
+#include <opencv2/imgproc.hpp>
 extern "C" {
 #include <libavutil/display.h>
 }
+#include <frameUtils.h>
+#include <pluglog.h>
+#ifdef WIN32
+#include <string_utils.h>
+#endif
 const char sep = separator();
 
 const std::string TAG = "FORESEG";
 
 namespace jami {
 
-PluginProcessor::PluginProcessor(const std::string& dataPath, const std::string& model, const std::string& backgroundImage, bool acc)
+PluginProcessor::PluginProcessor(const std::string& model, bool acc)
 {
-    activateAcc_ = acc;
-    initModel(dataPath+sep+"model/"+model);
-    setBackgroundImage(backgroundImage);
+    initModel(model, acc);
 }
 
 PluginProcessor::~PluginProcessor()
 {
+    mainFilter_.clean();
     Plog::log(Plog::LogPriority::INFO, TAG, "~pluginprocessor");
     if (session_)
         delete session_;
 }
 
 void
-PluginProcessor::setBackgroundImage(const std::string& backgroundPath)
-{
-    cv::Size size = cv::Size {0, 0};
-
-    if (!backgroundImage.empty())
-        size = backgroundImage.size();
-
-    cv::Mat newBackgroundImage = cv::imread(backgroundPath);
-    if (newBackgroundImage.cols == 0) {
-        Plog::log(Plog::LogPriority::ERR, TAG, "Background image not Loaded");
-    } else {
-        Plog::log(Plog::LogPriority::INFO, TAG, "Background image Loaded");
-        cv::cvtColor(newBackgroundImage, newBackgroundImage, cv::COLOR_BGR2RGB);
-        newBackgroundImage.convertTo(newBackgroundImage, CV_32FC3);
-        if (size.height) {
-            cv::resize(newBackgroundImage, newBackgroundImage, size);
-            backgroundRotation = 0;
-        }
-        backgroundImage = newBackgroundImage.clone();
-        newBackgroundImage.release();
-        hasBackground_ = true;
-    }
-}
-
-void
-PluginProcessor::initModel(const std::string& modelPath)
+PluginProcessor::initModel(const std::string& modelPath, bool activateAcc)
 {
     try {
         auto allocator_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
-        input_tensor_ = Ort::Value::CreateTensor<float>(allocator_info, input_image_.data(), input_image_.size(), input_shape_.data(), input_shape_.size());
-        output_tensor_ = Ort::Value::CreateTensor<float>(allocator_info, results_.data(), results_.size(), output_shape_.data(), output_shape_.size());
-        sessOpt_ =  Ort::SessionOptions();
+        input_tensor_ = Ort::Value::CreateTensor<float>(allocator_info,
+                                                        input_image_.data(),
+                                                        input_image_.size(),
+                                                        input_shape_.data(),
+                                                        input_shape_.size());
+        output_tensor_ = Ort::Value::CreateTensor<float>(allocator_info,
+                                                         results_.data(),
+                                                         results_.size(),
+                                                         output_shape_.data(),
+                                                         output_shape_.size());
+        sessOpt_ = Ort::SessionOptions();
 
 #ifdef NVIDIA
-        if (activateAcc_)
+        if (activateAcc)
             Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CUDA(sessOpt_, 0));
 #endif
-#ifdef ANDROID
-        if (activateAcc_)
+#ifdef __ANDROID__
+        if (activateAcc)
             Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_Nnapi(sessOpt_, 0));
 #endif
 
         sessOpt_.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
 #ifdef WIN32
-        std::wstring wsTmp(modelPath.begin(), modelPath.end());
-        session_ = new Ort::Session(env, wsTmp.c_str(), sessOpt_);
+        session_ = new Ort::Session(env_, to_wstring(modelPath).c_str(), sessOpt_);
 #else
-        session_ = new Ort::Session(env, modelPath.c_str(), sessOpt_);
+        session_ = new Ort::Session(env_, modelPath.c_str(), sessOpt_);
 #endif
         isAllocated_ = true;
+        Plog::log(Plog::LogPriority::INFO, TAG, "Model is allocated");
     } catch (std::exception& e) {
         Plog::log(Plog::LogPriority::ERR, TAG, e.what());
     }
-    std::ostringstream oss;
-    oss << "Model is allocated " << isAllocated_;
-    Plog::log(Plog::LogPriority::INFO, TAG, oss.str());
-}
-
-bool
-PluginProcessor::isAllocated()
-{
-    return isAllocated_;
 }
 
 void
-PluginProcessor::feedInput(const cv::Mat& frame)
+PluginProcessor::feedInput(AVFrame* input)
 {
-    cv::Mat temp(frame.rows, frame.cols, CV_32FC3, input_image_.data());
-    frame.convertTo(temp, CV_32FC3);
-}
-
-int
-PluginProcessor::getBackgroundRotation()
-{
-    return backgroundRotation;
-}
-
-void
-PluginProcessor::setBackgroundRotation(int angle)
-{
-    if (backgroundRotation != angle && (backgroundRotation - angle) != 0) {
-        rotateFrame(backgroundRotation - angle, backgroundImage);
-        backgroundRotation = angle;
-    }
+    cv::Mat frame = cv::Mat {
+        input->height,
+        input->width,
+        CV_8UC3,
+        input->data[0],
+        static_cast<size_t>(
+            input->linesize[0])}; // not zero input->linesize[0] leads to non continuous data
+    cvFrame_ = frame.clone();     // this is done to have continuous data
+    cv::Mat temp(modelInputDimensions.first,
+                 modelInputDimensions.second,
+                 CV_32FC3,
+                 input_image_.data());
+    cvFrame_.convertTo(temp, CV_32FC3);
 }
 
 void
 PluginProcessor::computePredictions()
 {
-    if (count == 0) {
+    if (count_ == 0) {
         // Run the graph
-        session_->Run(Ort::RunOptions{nullptr}, input_names, &input_tensor_, 1, output_names, &output_tensor_, 1);
-        computedMask = std::vector(results_.begin(), results_.end());
+        session_->Run(Ort::RunOptions {nullptr},
+                      modelInputNames,
+                      &input_tensor_,
+                      1,
+                      modelOutputNames,
+                      &output_tensor_,
+                      1);
+        computedMask_ = std::vector(results_.begin(), results_.end());
     }
 }
 
 void
 PluginProcessor::printMask()
 {
-    for (size_t i = 0; i < computedMask.size(); i++) {
+    std::ostringstream oss;
+    for (size_t i = 0; i < computedMask_.size(); i++) {
         // Log the predictions
-        std::ostringstream oss;
-        oss << "\nclass: " << computedMask[i] << std::endl;
-        Plog::log(Plog::LogPriority::INFO, TAG, oss.str());
-    }
-}
-
-void
-PluginProcessor::resetInitValues(const cv::Size& modelInputSize)
-{
-    previousMasks[0] = cv::Mat(modelInputSize.height, modelInputSize.width, CV_32FC1, double(0.));
-    previousMasks[1] = cv::Mat(modelInputSize.height, modelInputSize.width, CV_32FC1, double(0.));
-    kSize = cv::Size(modelInputSize.width * kernelSize, modelInputSize.height * kernelSize);
-    if (kSize.height % 2 == 0) {
-        kSize.height -= 1;
-    }
-    if (kSize.width % 2 == 0) {
-        kSize.width -= 1;
-    }
-    count = 0;
-    grabCutMode = cv::GC_INIT_WITH_MASK;
-    grabCutIterations = 5;
-}
-
-void
-copyByLine(uchar* frameData, uchar* applyMaskData, const int lineSize, cv::Size size)
-{
-    if (3 * size.width == lineSize) {
-        std::memcpy(frameData, applyMaskData, size.height * size.width * 3);
-    } else {
-        int rows = size.height;
-        int offset = 0;
-        int maskoffset = 0;
-        for (int i = 0; i < rows; i++) {
-            std::memcpy(frameData + offset, applyMaskData + maskoffset, lineSize);
-            offset += lineSize;
-            maskoffset += 3 * size.width;
+        if (computedMask_[i] > 2) {
+            oss << computedMask_[i] << " " << std::endl;
         }
     }
+    Plog::log(Plog::LogPriority::INFO, TAG, oss.str());
 }
 
 void
-PluginProcessor::drawMaskOnFrame(
-    cv::Mat& frame, cv::Mat& frameReduced, std::vector<float> computedMask, int lineSize, int angle)
+PluginProcessor::resetInitValues()
 {
-    if (computedMask.empty()) {
-        return;
+    previousMasks_[0] = previousMasks_[1] = cv::Mat(modelInputDimensions.first,
+                                                    modelInputDimensions.second,
+                                                    CV_32FC1,
+                                                    double(0.));
+    kSize_ = cv::Size(modelInputDimensions.first * kernelSize_,
+                      modelInputDimensions.second * kernelSize_);
+    if (kSize_.height % 2 == 0) {
+        kSize_.height -= 1;
     }
+    if (kSize_.width % 2 == 0) {
+        kSize_.width -= 1;
+    }
+    count_ = 0;
+    grabCutMode_ = cv::GC_INIT_WITH_MASK;
+    grabCutIterations_ = 4;
+}
 
-    if (count == 0) {
-        int maskSize = static_cast<int>(std::sqrt(computedMask.size()));
-        cv::Mat maskImg(maskSize, maskSize, CV_32FC1, computedMask.data());
-        cv::Mat* applyMask = &frameReduced;
+void
+PluginProcessor::drawMaskOnFrame(AVFrame* frame, AVFrame* frameReduced, int angle)
+{
+    if (computedMask_.empty() || !mainFilter_.initialized_)
+        return;
 
-        rotateFrame(-angle, maskImg);
-        cv::resize(maskImg, maskImg, cv::Size(frameReduced.cols, frameReduced.rows));
+    if (count_ == 0) {
+        int maskSize = static_cast<int>(std::sqrt(computedMask_.size()));
+        cv::Mat maskImg(maskSize, maskSize, CV_32FC1, computedMask_.data());
+
+        cv::resize(maskImg,
+                   maskImg,
+                   cv::Size(modelInputDimensions.first, modelInputDimensions.second));
 
         double m, M;
         cv::minMaxLoc(maskImg, &m, &M);
-
-        if (M < 2) { // avoid detection if there is any one in frame
+        bool improveMask = !isBlur_;
+        if (M < 2) { // avoid detection if there isn't anyone in frame
             maskImg = 0. * maskImg;
+            improveMask = false;
         } else {
             for (int i = 0; i < maskImg.cols; i++) {
                 for (int j = 0; j < maskImg.rows; j++) {
@@ -224,9 +183,9 @@
                     if (maskImg.at<float>(j, i) < 0.4)
                         maskImg.at<float>(j, i) = 0.;
                     else if (maskImg.at<float>(j, i) < 0.7) {
-                        float value = maskImg.at<float>(j, i) * smoothFactors[0]
-                                      + previousMasks[0].at<float>(j, i) * smoothFactors[1]
-                                      + previousMasks[1].at<float>(j, i) * smoothFactors[2];
+                        float value = maskImg.at<float>(j, i) * smoothFactors_[0]
+                                      + previousMasks_[0].at<float>(j, i) * smoothFactors_[1]
+                                      + previousMasks_[1].at<float>(j, i) * smoothFactors_[2];
                         maskImg.at<float>(j, i) = 0.;
                         if (value > 0.7)
                             maskImg.at<float>(j, i) = 1.;
@@ -235,91 +194,230 @@
                 }
             }
         }
-        if (cv::countNonZero(maskImg) != 0) {
+
+        if (improveMask) {
             cv::Mat dilate;
             cv::dilate(maskImg,
-                        dilate,
-                        cv::getStructuringElement(cv::MORPH_ELLIPSE, kSize),
-                        cv::Point(-1, -1),
-                        2);
+                       dilate,
+                       cv::getStructuringElement(cv::MORPH_ELLIPSE, kSize_),
+                       cv::Point(-1, -1),
+                       2);
             cv::erode(maskImg,
-                        maskImg,
-                        cv::getStructuringElement(cv::MORPH_ELLIPSE, kSize),
-                        cv::Point(-1, -1),
-                        2);
+                      maskImg,
+                      cv::getStructuringElement(cv::MORPH_ELLIPSE, kSize_),
+                      cv::Point(-1, -1),
+                      2);
             for (int i = 0; i < maskImg.cols; i++) {
                 for (int j = 0; j < maskImg.rows; j++) {
                     if (dilate.at<float>(j, i) != maskImg.at<float>(j, i))
-                        maskImg.at<float>(j, i) = grabcutClass;
+                        maskImg.at<float>(j, i) = grabcutClass_;
                 }
             }
+            cv::Mat applyMask = cvFrame_.clone();
             maskImg.convertTo(maskImg, CV_8UC1);
-            applyMask->convertTo(*applyMask, CV_8UC1);
+            applyMask.convertTo(applyMask, CV_8UC1);
             cv::Rect rect(1, 1, maskImg.rows, maskImg.cols);
-            cv::grabCut(*applyMask,
+            cv::grabCut(applyMask,
                         maskImg,
                         rect,
-                        bgdModel,
-                        fgdModel,
-                        grabCutIterations,
-                        grabCutMode);
+                        bgdModel_,
+                        fgdModel_,
+                        grabCutIterations_,
+                        grabCutMode_);
 
-            grabCutMode = cv::GC_EVAL;
-            grabCutIterations = 1;
+            grabCutMode_ = cv::GC_EVAL;
+            grabCutIterations_ = 1;
 
             maskImg = maskImg & 1;
             maskImg.convertTo(maskImg, CV_32FC1);
             maskImg *= 255.;
-            GaussianBlur(maskImg, maskImg, cv::Size(7, 7), 0); // float mask from 0 to 255.
+            blur(maskImg, maskImg, cv::Size(7, 7)); // float mask from 0 to 255.
             maskImg = maskImg / 255.;
         }
-        previousMasks[1] = previousMasks[0].clone();
-        previousMasks[0] = maskImg.clone();
+
+        previousMasks_[1] = previousMasks_[0].clone();
+        previousMasks_[0] = maskImg.clone();
     }
 
-    cv::Mat roiMaskImg = previousMasks[0].clone();
-    cv::Mat roiMaskImgComplementary = 1. - roiMaskImg; // mask from 1. to 0
+    cv::Mat roiMaskImg = previousMasks_[0].clone() * 255.;
+    roiMaskImg.convertTo(roiMaskImg, CV_8UC1);
 
-    std::vector<cv::Mat> channels;
-    std::vector<cv::Mat> channelsComplementary;
+    gsFrame maskFrame = {av_frame_alloc(), frameFree};
+    maskFrame->format = AV_PIX_FMT_GRAY8;
+    maskFrame->width = roiMaskImg.cols;
+    maskFrame->height = roiMaskImg.rows;
+    maskFrame->linesize[0] = roiMaskImg.step;
+    if (av_frame_get_buffer(maskFrame.get(), 0) < 0)
+        return;
 
-    channels.emplace_back(roiMaskImg);
-    channels.emplace_back(roiMaskImg);
-    channels.emplace_back(roiMaskImg);
-    channelsComplementary.emplace_back(roiMaskImgComplementary);
-    channelsComplementary.emplace_back(roiMaskImgComplementary);
-    channelsComplementary.emplace_back(roiMaskImgComplementary);
+    maskFrame->data[0] = roiMaskImg.data;
+    maskFrame->pts = 1;
+    mainFilter_.feedInput(maskFrame.get(), "mask");
+    maskFrame.reset();
 
-    cv::merge(channels, roiMaskImg);
-    cv::merge(channelsComplementary, roiMaskImgComplementary);
+    if (isBlur_)
+        mainFilter_.feedInput(frameReduced, "input");
+    mainFilter_.feedInput(frame, "input2");
+    AVFrame* filteredFrame;
+    if ((filteredFrame = mainFilter_.readOutput())) {
+        moveFrom(frame, filteredFrame);
+        frameFree(filteredFrame);
+    }
+    count_++;
+    count_ = count_ % frameCount_;
+}
 
-    cv::Mat output;
-    frameReduced.convertTo(output, roiMaskImg.type());
-    output = output.mul(roiMaskImg);
-    output += backgroundImage.mul(roiMaskImgComplementary);
-    output.convertTo(output, frameReduced.type());
+MediaStream
+PluginProcessor::getbgAVFrameInfos()
+{
+    AVFormatContext* ctx = avformat_alloc_context();
+    // Open
+    if (avformat_open_input(&ctx, backgroundPath_.c_str(), NULL, NULL) != 0) {
+        avformat_free_context(ctx);
+        Plog::log(Plog::LogPriority::INFO, TAG, "Couldn't open input stream.");
+        return {};
+    }
+    pFormatCtx_.reset(ctx);
+    // Retrieve stream information
+    if (avformat_find_stream_info(pFormatCtx_.get(), NULL) < 0) {
+        Plog::log(Plog::LogPriority::INFO, TAG, "Couldn't find stream information.");
+        return {};
+    }
 
-    cv::resize(output, output, cv::Size(frame.cols, frame.rows));
+    // Dump valid information onto standard error
+    av_dump_format(pFormatCtx_.get(), 0, backgroundPath_.c_str(), false);
 
-    copyByLine(frame.data, output.data, lineSize, cv::Size(frame.cols, frame.rows));
-    count++;
-    count = count % frameCount;
+    // Find the video stream
+    for (int i = 0; i < static_cast<int>(pFormatCtx_->nb_streams); i++)
+        if (pFormatCtx_->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
+            videoStream_ = i;
+            break;
+        }
+
+    if (videoStream_ == -1) {
+        Plog::log(Plog::LogPriority::INFO, TAG, "Didn't find a video stream.");
+        return {};
+    }
+
+    rational<int> fr = pFormatCtx_->streams[videoStream_]->r_frame_rate;
+    return MediaStream("background",
+                       pFormatCtx_->streams[videoStream_]->codecpar->format,
+                       1 / fr,
+                       pFormatCtx_->streams[videoStream_]->codecpar->width,
+                       pFormatCtx_->streams[videoStream_]->codecpar->height,
+                       0,
+                       fr);
 }
 
 void
-PluginProcessor::rotateFrame(int angle, cv::Mat& mat)
+PluginProcessor::loadBackground()
 {
-    if (angle == -90)
-        cv::rotate(mat, mat, cv::ROTATE_90_COUNTERCLOCKWISE);
-    else if (std::abs(angle) == 180)
-        cv::rotate(mat, mat, cv::ROTATE_180);
-    else if (angle == 90)
-        cv::rotate(mat, mat, cv::ROTATE_90_CLOCKWISE);
+    if (backgroundPath_.empty())
+        return;
+
+    auto bgStream_ = getbgAVFrameInfos();
+    if (mainFilter_.initialize(mainFilterDescription_, {maskms_, bgStream_, ims2_}) < 0)
+        return;
+
+    int got_frame;
+    AVCodecContext* pCodecCtx;
+    AVCodec* pCodec;
+    AVPacket* packet;
+
+    pCodecCtx = pFormatCtx_->streams[videoStream_]->codec;
+    pCodec = avcodec_find_decoder(pCodecCtx->codec_id);
+    if (pCodec == nullptr) {
+        mainFilter_.clean();
+        pFormatCtx_.reset();
+        Plog::log(Plog::LogPriority::INFO, TAG, "Codec not found.");
+        return;
+    }
+
+    // Open codec
+    if (avcodec_open2(pCodecCtx, pCodec, NULL) < 0) {
+        mainFilter_.clean();
+        pFormatCtx_.reset();
+        Plog::log(Plog::LogPriority::INFO, TAG, "Could not open codec.");
+        return;
+    }
+
+    packet = av_packet_alloc();
+    if (av_read_frame(pFormatCtx_.get(), packet) < 0) {
+        mainFilter_.clean();
+        avcodec_close(pCodecCtx);
+        av_packet_free(&packet);
+        pFormatCtx_.reset();
+        Plog::log(Plog::LogPriority::INFO, TAG, "Could not read packet from context.");
+        return;
+    }
+
+    AVFrame* bgImage = av_frame_alloc();
+    avcodec_decode_video2(pCodecCtx, bgImage, &got_frame, packet);
+    if (got_frame) {
+        mainFilter_.feedInput(bgImage, "background");
+        mainFilter_.feedEOF("background");
+    } else
+        mainFilter_.clean();
+
+    frameFree(bgImage);
+    avcodec_close(pCodecCtx);
+    av_packet_free(&packet);
+    pFormatCtx_.reset();
 }
 
-bool
-PluginProcessor::hasBackground() const
+void
+PluginProcessor::initFilters(const std::pair<int, int>& inputSize, int format, int angle)
 {
-    return hasBackground_;
+    resetInitValues();
+    mainFilter_.clean();
+
+    std::string rotateSides = "";
+    std::string scaleSize = std::to_string(inputSize.first) + ":"
+                            + std::to_string(inputSize.second);
+    Plog::log(Plog::LogPriority::INFO, TAG, scaleSize);
+    if (std::abs(angle) == 90) {
+        rotateSides = ":out_w=ih:out_h=iw";
+        scaleSize = std::to_string(inputSize.second) + ":" + std::to_string(inputSize.first);
+    }
+
+    rational<int> fr(1, 1);
+    ims_ = MediaStream("input",
+                       AV_PIX_FMT_RGB24,
+                       1 / fr,
+                       modelInputDimensions.first,
+                       modelInputDimensions.second,
+                       0,
+                       fr);
+    ims2_ = MediaStream("input2", format, 1 / fr, inputSize.first, inputSize.second, 0, fr);
+
+    maskms_ = MediaStream("mask",
+                          AV_PIX_FMT_GRAY8,
+                          1 / fr,
+                          modelInputDimensions.first,
+                          modelInputDimensions.second,
+                          0,
+                          fr);
+
+    if (isBlur_) {
+        mainFilterDescription_ = "[mask]negate[negated],[input][negated]alphamerge,boxblur="
+                                 + blurLevel_ + ",scale=" + scaleSize
+                                 + "[blured],[input2]format=rgb24,rotate=" + rotation[-angle]
+                                 + rotateSides
+                                 + "[input2formated],[input2formated][blured]overlay,rotate="
+                                 + rotation[angle] + rotateSides;
+        Plog::log(Plog::LogPriority::INFO, TAG, mainFilterDescription_);
+        mainFilter_.initialize(mainFilterDescription_, {maskms_, ims_, ims2_});
+    } else {
+        mainFilterDescription_ = "[mask]scale=" + scaleSize
+                                 + "[fgmask],[fgmask]split=2[bg][fg],[bg]negate[bgn],"
+                                 + "[background]scale=" + scaleSize + "[backgroundformated],"
+                                 + "[backgroundformated][bgn]alphamerge[background2],"
+                                 + "[input2]rotate=" + rotation[-angle] + rotateSides
+                                 + "[input2formated],[input2formated][fg]alphamerge[foreground],"
+                                 + "[foreground][background2]overlay," + "rotate=" + rotation[angle]
+                                 + rotateSides;
+        Plog::log(Plog::LogPriority::INFO, TAG, mainFilterDescription_);
+        loadBackground();
+    }
 }
 } // namespace jami