GreenScreen/videoSubscriber.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 "videoSubscriber.h"
// Use for display rotation matrix
extern "C" {
#include <libavutil/display.h>
}
#include <accel.h>
#include <frameUtils.h>

// Opencv processing
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>

// LOGGING
#include <pluglog.h>

const std::string TAG = "FORESEG";
const char sep = separator();

namespace jami {

VideoSubscriber::VideoSubscriber(const std::string& dataPath, const std::string& model, const std::string& backgroundImage, bool acc)
    : path_ {dataPath}
    , pluginProcessor {dataPath, model, backgroundImage, acc}
{
    /**
     * Waits for new frames and then process them
     * Writes the predictions in computedPredictions
     **/
    processFrameThread = std::thread([this] {
        while (running) {
            std::unique_lock<std::mutex> l(inputLock);
            inputCv.wait(l, [this] { return not running or newFrame; });
            if (not running) {
                break;
            }

            pluginProcessor.feedInput(fcopy.resizedFrameRGB);
            newFrame = false;
            /** Unclock the mutex, this way we let the other thread
             *  copy new data while we are processing the old one
             **/
            l.unlock();
            pluginProcessor.computePredictions();
        }
    });
}

VideoSubscriber::~VideoSubscriber()
{
    std::ostringstream oss;
    oss << "~MediaProcessor" << std::endl;
    stop();
    processFrameThread.join();
    Plog::log(Plog::LogPriority::INFO, TAG, oss.str());
}

void
VideoSubscriber::update(jami::Observable<AVFrame*>*, AVFrame* const& pluginFrame)
{
    if (!observable_ || !pluginProcessor.isAllocated() || !pluginProcessor.hasBackground() || !pluginFrame)
        return;

    //======================================================================================
    // GET FRAME ROTATION
    AVFrameSideData* side_data = av_frame_get_side_data(pluginFrame,
                                                        AV_FRAME_DATA_DISPLAYMATRIX);

    int angle {0};
    if (side_data) {
        auto matrix_rotation = reinterpret_cast<int32_t*>(side_data->data);
        angle = static_cast<int>(av_display_rotation_get(matrix_rotation));
    }

    //======================================================================================
    // GET RAW FRAME
    // Use a non-const Frame
    // Convert input frame to RGB
    int inputHeight = pluginFrame->height;
    int inputWidth = pluginFrame->width;

    fcopy.originalSize = cv::Size {inputWidth, inputHeight};

    AVFrame* temp = transferToMainMemory(pluginFrame, AV_PIX_FMT_NV12);
    AVFrame* bgrFrame = scaler.convertFormat(temp, AV_PIX_FMT_RGB24);
    av_frame_unref(temp);
    av_frame_free(&temp);
    if (!bgrFrame)
        return;
    cv::Mat frame = cv::Mat {bgrFrame->height,
                                bgrFrame->width,
                                CV_8UC3,
                                bgrFrame->data[0],
                                static_cast<size_t>(bgrFrame->linesize[0])};
    // First clone the frame as the original one is unusable because of
    // linespace

    cv::Mat clone = frame.clone();
    //======================================================================================

    pluginProcessor.setBackgroundRotation(angle);

    if (firstRun) {
        fcopy.resizedSize = cv::Size {257, 257};
        pluginProcessor.resetInitValues(fcopy.resizedSize);

        cv::resize(clone, fcopy.resizedFrameRGB, fcopy.resizedSize);
        pluginProcessor.rotateFrame(angle, fcopy.resizedFrameRGB);

        cv::resize(pluginProcessor.backgroundImage,
                    pluginProcessor.backgroundImage,
                    fcopy.resizedSize);

        firstRun = false;
    }

    if (!newFrame) {
        std::lock_guard<std::mutex> l(inputLock);
        cv::resize(clone, fcopy.resizedFrameRGB, fcopy.resizedSize);
        pluginProcessor.rotateFrame(angle, fcopy.resizedFrameRGB);
        newFrame = true;
        inputCv.notify_all();
    }

    fcopy.predictionsFrameRGB = frame;
    fcopy.predictionsResizedFrameRGB = fcopy.resizedFrameRGB.clone();
    pluginProcessor.rotateFrame(-angle, fcopy.predictionsResizedFrameRGB);
    pluginProcessor.drawMaskOnFrame(fcopy.predictionsFrameRGB,
                                    fcopy.predictionsResizedFrameRGB,
                                    pluginProcessor.computedMask,
                                    bgrFrame->linesize[0],
                                    angle);

    //======================================================================================
    // REPLACE AVFRAME DATA WITH FRAME DATA
    if (bgrFrame->data[0]) {
        uint8_t* frameData = bgrFrame->data[0];
        if (angle == 90 || angle == -90) {
            std::memmove(frameData,
                            fcopy.predictionsFrameRGB.data,
                            static_cast<size_t>(pluginFrame->width * pluginFrame->height * 3)
                                * sizeof(uint8_t));
        }

        av_frame_copy_props(bgrFrame, pluginFrame);
        moveFrom(pluginFrame, bgrFrame);
    }
    av_frame_unref(bgrFrame);
    av_frame_free(&bgrFrame);
}

void
VideoSubscriber::attached(jami::Observable<AVFrame*>* observable)
{
    std::ostringstream oss;
    oss << "::Attached ! " << std::endl;
    Plog::log(Plog::LogPriority::INFO, TAG, oss.str());
    observable_ = observable;
}

void
VideoSubscriber::detached(jami::Observable<AVFrame*>*)
{
    firstRun = true;
    observable_ = nullptr;
    std::ostringstream oss;
    oss << "::Detached()" << std::endl;
    Plog::log(Plog::LogPriority::INFO, TAG, oss.str());
}

void
VideoSubscriber::detach()
{
    if (observable_) {
        firstRun = true;
        std::ostringstream oss;
        oss << "::Calling detach()" << std::endl;
        Plog::log(Plog::LogPriority::INFO, TAG, oss.str());
        observable_->detach(this);
    }
}

void
VideoSubscriber::stop()
{
    running = false;
    inputCv.notify_all();
}

void
VideoSubscriber::setBackground(const std::string& backgroundPath)
{
    pluginProcessor.setBackgroundImage(backgroundPath);
}
} // namespace jami