GreenScreen/pluginProcessor.cpp

/**
 *  Copyright (C) 2020 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.
 */

#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>

extern "C" {
#include <libavutil/display.h>
}

const char sep = separator();

const std::string TAG = "FORESEG";

PluginParameters* mPluginParameters = getGlobalPluginParameters();

namespace jami {

PluginProcessor::PluginProcessor(const std::string& dataPath)
    : pluginInference {TFModel {dataPath + sep + "models" + sep + mPluginParameters->model}}
{
    initModel();
    setBackgroundImage(mPluginParameters->image);
}

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()
{
    try {
        pluginInference.init();
    } catch (std::exception& e) {
        Plog::log(Plog::LogPriority::ERR, TAG, e.what());
    }
    std::ostringstream oss;
    oss << "Model is allocated " << pluginInference.isAllocated();
    Plog::log(Plog::LogPriority::INFO, TAG, oss.str());
}

#ifdef TFLITE
void
PluginProcessor::feedInput(const cv::Mat& frame)
{
    auto pair = pluginInference.getInput();
    uint8_t* inputPointer = pair.first;

    cv::Mat temp(frame.rows, frame.cols, CV_8UC3, inputPointer);
    frame.convertTo(temp, CV_8UC3);

    inputPointer = nullptr;
}
#else
void
PluginProcessor::feedInput(const cv::Mat& frame)
{
    pluginInference.ReadTensorFromMat(frame);
}
#endif // TFLITE

int
PluginProcessor::getBackgroundRotation()
{
    return backgroundRotation;
}

void
PluginProcessor::setBackgroundRotation(int angle)
{
    if (backgroundRotation != angle && (backgroundRotation - angle) != 0) {
        rotateFrame(backgroundRotation - angle, backgroundImage);
        backgroundRotation = angle;
    }
}

void
PluginProcessor::computePredictions()
{
    if (count == 0) {
        // Run the graph
        pluginInference.runGraph();
        auto predictions = pluginInference.masksPredictions();

        // Save the predictions
        computedMask = predictions;
    }
}

void
PluginProcessor::printMask()
{
    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;
        }
    }
}

void
PluginProcessor::drawMaskOnFrame(
    cv::Mat& frame, cv::Mat& frameReduced, std::vector<float> computedMask, int lineSize, int angle)
{
    if (computedMask.empty()) {
        return;
    }

    int maskSize = static_cast<int>(std::sqrt(computedMask.size()));
    cv::Mat maskImg(maskSize, maskSize, CV_32FC1, computedMask.data());
    cv::Mat* applyMask = &frameReduced;
    cv::Mat output;

    if (count == 0) {
        rotateFrame(-angle, maskImg);
#ifdef TFLITE
        for (int i = 0; i < maskImg.cols; i++) {
            for (int j = 0; j < maskImg.rows; j++) {
                if (maskImg.at<float>(j, i) == 15)
                    maskImg.at<float>(j, i) = 1.;
                else
                    maskImg.at<float>(j, i) = smoothFactors[0] * previousMasks[0].at<float>(j, i)
                                              + smoothFactors[1] * previousMasks[1].at<float>(j, i);
            }
        }
        cv::morphologyEx(maskImg,
                         maskImg,
                         cv::MORPH_CLOSE,
                         cv::getStructuringElement(cv::MORPH_ELLIPSE, kSize),
                         cv::Point(-1, -1),
                         4);
#else
        cv::resize(maskImg, maskImg, cv::Size(frameReduced.cols, frameReduced.rows));

        double m, M;
        cv::minMaxLoc(maskImg, &m, &M);

        if (M < 2) { // avoid detection if there is any one in frame
            maskImg = 0. * maskImg;
        } else {
            for (int i = 0; i < maskImg.cols; i++) {
                for (int j = 0; j < maskImg.rows; j++) {
                    maskImg.at<float>(j, i) = (maskImg.at<float>(j, i) - m) / (M - m);

                    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];
                        maskImg.at<float>(j, i) = 0.;
                        if (value > 0.7)
                            maskImg.at<float>(j, i) = 1.;
                    } else
                        maskImg.at<float>(j, i) = 1.;
                }
            }
        }
#endif
        if (cv::countNonZero(maskImg) != 0) {
#ifdef TFLITE
            cv::Mat tfMask;
            tfMask = maskImg.clone();
            tfMask *= 255.;
            tfMask.convertTo(tfMask, CV_8UC1);
            cv::threshold(tfMask, tfMask, 127, 255, cv::THRESH_BINARY);
            if (cv::countNonZero(tfMask) != 0) {
#endif
                cv::Mat dilate;
                cv::dilate(maskImg,
                           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);
                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.convertTo(maskImg, CV_8UC1);
                applyMask->convertTo(*applyMask, CV_8UC1);
                cv::Rect rect(1, 1, maskImg.rows, maskImg.cols);
                cv::grabCut(*applyMask,
                            maskImg,
                            rect,
                            bgdModel,
                            fgdModel,
                            grabCutIterations,
                            grabCutMode);

                grabCutMode = cv::GC_EVAL;
                grabCutIterations = 1;

                maskImg = maskImg & 1;
#ifdef TFLITE
                cv::bitwise_and(maskImg, tfMask, maskImg);
            }
#endif
            maskImg.convertTo(maskImg, CV_32FC1);
            maskImg *= 255.;
            GaussianBlur(maskImg, maskImg, cv::Size(7, 7), 0); // float mask from 0 to 255.
            maskImg = maskImg / 255.;
        }
        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

    std::vector<cv::Mat> channels;
    std::vector<cv::Mat> channelsComplementary;

    channels.emplace_back(roiMaskImg);
    channels.emplace_back(roiMaskImg);
    channels.emplace_back(roiMaskImg);
    channelsComplementary.emplace_back(roiMaskImgComplementary);
    channelsComplementary.emplace_back(roiMaskImgComplementary);
    channelsComplementary.emplace_back(roiMaskImgComplementary);

    cv::merge(channels, roiMaskImg);
    cv::merge(channelsComplementary, roiMaskImgComplementary);

    int origType = frameReduced.type();
    int roiMaskType = roiMaskImg.type();

    frameReduced.convertTo(output, roiMaskType);
    output = output.mul(roiMaskImg);
    output += backgroundImage.mul(roiMaskImgComplementary);
    output.convertTo(output, origType);

    cv::resize(output, output, cv::Size(frame.cols, frame.rows));

    copyByLine(frame.data, output.data, lineSize, cv::Size(frame.cols, frame.rows));
    count++;
    count = count % frameCount;
}

void
PluginProcessor::rotateFrame(int angle, cv::Mat& mat)
{
    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);
}

bool
PluginProcessor::hasBackground() const
{
    return hasBackground_;
}
} // namespace jami