agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 1 | /** |
Sébastien Blin | cb783e3 | 2021-02-12 11:34:10 -0500 | [diff] [blame] | 2 | * Copyright (C) 2020-2021 Savoir-faire Linux Inc. |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 3 | * |
| 4 | * Author: Aline Gondim Santos <aline.gondimsantos@savoirfairelinux.com> |
| 5 | * |
| 6 | * This program is free software; you can redistribute it and/or modify |
| 7 | * it under the terms of the GNU General Public License as published by |
| 8 | * the Free Software Foundation; either version 3 of the License, or |
| 9 | * (at your option) any later version. |
| 10 | * |
| 11 | * This program is distributed in the hope that it will be useful, |
| 12 | * but WITHOUT ANY WARRANTY; without even the implied warranty of |
| 13 | * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
| 14 | * GNU General Public License for more details. |
| 15 | * |
| 16 | * You should have received a copy of the GNU General Public License |
| 17 | * along with this program; if not, write to the Free Software |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 18 | * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 |
| 19 | * USA. |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 20 | */ |
| 21 | |
| 22 | #include "pluginProcessor.h" |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 23 | |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 24 | #include <opencv2/imgproc.hpp> |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 25 | extern "C" { |
| 26 | #include <libavutil/display.h> |
| 27 | } |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 28 | #include <frameUtils.h> |
| 29 | #include <pluglog.h> |
| 30 | #ifdef WIN32 |
| 31 | #include <string_utils.h> |
| 32 | #endif |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 33 | const char sep = separator(); |
| 34 | |
| 35 | const std::string TAG = "FORESEG"; |
| 36 | |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 37 | namespace jami { |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 38 | |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 39 | PluginProcessor::PluginProcessor(const std::string& model, bool acc) |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 40 | { |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 41 | initModel(model, acc); |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 42 | } |
| 43 | |
| 44 | PluginProcessor::~PluginProcessor() |
| 45 | { |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 46 | mainFilter_.clean(); |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 47 | Plog::log(Plog::LogPriority::INFO, TAG, "~pluginprocessor"); |
| 48 | if (session_) |
| 49 | delete session_; |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 50 | } |
| 51 | |
| 52 | void |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 53 | PluginProcessor::initModel(const std::string& modelPath, bool activateAcc) |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 54 | { |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 55 | try { |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 56 | auto allocator_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 57 | input_tensor_ = Ort::Value::CreateTensor<float>(allocator_info, |
| 58 | input_image_.data(), |
| 59 | input_image_.size(), |
| 60 | input_shape_.data(), |
| 61 | input_shape_.size()); |
| 62 | output_tensor_ = Ort::Value::CreateTensor<float>(allocator_info, |
| 63 | results_.data(), |
| 64 | results_.size(), |
| 65 | output_shape_.data(), |
| 66 | output_shape_.size()); |
| 67 | sessOpt_ = Ort::SessionOptions(); |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 68 | |
| 69 | #ifdef NVIDIA |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 70 | if (activateAcc) |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 71 | Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CUDA(sessOpt_, 0)); |
| 72 | #endif |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 73 | #ifdef __ANDROID__ |
| 74 | if (activateAcc) |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 75 | Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_Nnapi(sessOpt_, 0)); |
| 76 | #endif |
| 77 | |
| 78 | sessOpt_.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL); |
| 79 | #ifdef WIN32 |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 80 | session_ = new Ort::Session(env_, to_wstring(modelPath).c_str(), sessOpt_); |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 81 | #else |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 82 | session_ = new Ort::Session(env_, modelPath.c_str(), sessOpt_); |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 83 | #endif |
| 84 | isAllocated_ = true; |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 85 | Plog::log(Plog::LogPriority::INFO, TAG, "Model is allocated"); |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 86 | } catch (std::exception& e) { |
| 87 | Plog::log(Plog::LogPriority::ERR, TAG, e.what()); |
| 88 | } |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 89 | } |
| 90 | |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 91 | void |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 92 | PluginProcessor::feedInput(AVFrame* input) |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 93 | { |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 94 | cv::Mat frame = cv::Mat { |
| 95 | input->height, |
| 96 | input->width, |
| 97 | CV_8UC3, |
| 98 | input->data[0], |
| 99 | static_cast<size_t>( |
| 100 | input->linesize[0])}; // not zero input->linesize[0] leads to non continuous data |
| 101 | cvFrame_ = frame.clone(); // this is done to have continuous data |
| 102 | cv::Mat temp(modelInputDimensions.first, |
| 103 | modelInputDimensions.second, |
| 104 | CV_32FC3, |
| 105 | input_image_.data()); |
| 106 | cvFrame_.convertTo(temp, CV_32FC3); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 107 | } |
| 108 | |
| 109 | void |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 110 | PluginProcessor::computePredictions() |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 111 | { |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 112 | if (count_ == 0) { |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 113 | // Run the graph |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 114 | session_->Run(Ort::RunOptions {nullptr}, |
| 115 | modelInputNames, |
| 116 | &input_tensor_, |
| 117 | 1, |
| 118 | modelOutputNames, |
| 119 | &output_tensor_, |
| 120 | 1); |
| 121 | computedMask_ = std::vector(results_.begin(), results_.end()); |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 122 | } |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 123 | } |
| 124 | |
| 125 | void |
| 126 | PluginProcessor::printMask() |
| 127 | { |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 128 | std::ostringstream oss; |
| 129 | for (size_t i = 0; i < computedMask_.size(); i++) { |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 130 | // Log the predictions |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 131 | if (computedMask_[i] > 2) { |
| 132 | oss << computedMask_[i] << " " << std::endl; |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 133 | } |
| 134 | } |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 135 | Plog::log(Plog::LogPriority::INFO, TAG, oss.str()); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 136 | } |
| 137 | |
| 138 | void |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 139 | PluginProcessor::resetInitValues() |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 140 | { |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 141 | previousMasks_[0] = previousMasks_[1] = cv::Mat(modelInputDimensions.first, |
| 142 | modelInputDimensions.second, |
| 143 | CV_32FC1, |
| 144 | double(0.)); |
| 145 | kSize_ = cv::Size(modelInputDimensions.first * kernelSize_, |
| 146 | modelInputDimensions.second * kernelSize_); |
| 147 | if (kSize_.height % 2 == 0) { |
| 148 | kSize_.height -= 1; |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 149 | } |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 150 | if (kSize_.width % 2 == 0) { |
| 151 | kSize_.width -= 1; |
| 152 | } |
| 153 | count_ = 0; |
| 154 | grabCutMode_ = cv::GC_INIT_WITH_MASK; |
| 155 | grabCutIterations_ = 4; |
| 156 | } |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 157 | |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 158 | void |
| 159 | PluginProcessor::drawMaskOnFrame(AVFrame* frame, AVFrame* frameReduced, int angle) |
| 160 | { |
| 161 | if (computedMask_.empty() || !mainFilter_.initialized_) |
| 162 | return; |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 163 | |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 164 | if (count_ == 0) { |
| 165 | int maskSize = static_cast<int>(std::sqrt(computedMask_.size())); |
| 166 | cv::Mat maskImg(maskSize, maskSize, CV_32FC1, computedMask_.data()); |
| 167 | |
| 168 | cv::resize(maskImg, |
| 169 | maskImg, |
| 170 | cv::Size(modelInputDimensions.first, modelInputDimensions.second)); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 171 | |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 172 | double m, M; |
| 173 | cv::minMaxLoc(maskImg, &m, &M); |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 174 | bool improveMask = !isBlur_; |
| 175 | if (M < 2) { // avoid detection if there isn't anyone in frame |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 176 | maskImg = 0. * maskImg; |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 177 | improveMask = false; |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 178 | } else { |
| 179 | for (int i = 0; i < maskImg.cols; i++) { |
| 180 | for (int j = 0; j < maskImg.rows; j++) { |
| 181 | maskImg.at<float>(j, i) = (maskImg.at<float>(j, i) - m) / (M - m); |
| 182 | |
| 183 | if (maskImg.at<float>(j, i) < 0.4) |
| 184 | maskImg.at<float>(j, i) = 0.; |
| 185 | else if (maskImg.at<float>(j, i) < 0.7) { |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 186 | float value = maskImg.at<float>(j, i) * smoothFactors_[0] |
| 187 | + previousMasks_[0].at<float>(j, i) * smoothFactors_[1] |
| 188 | + previousMasks_[1].at<float>(j, i) * smoothFactors_[2]; |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 189 | maskImg.at<float>(j, i) = 0.; |
| 190 | if (value > 0.7) |
| 191 | maskImg.at<float>(j, i) = 1.; |
| 192 | } else |
| 193 | maskImg.at<float>(j, i) = 1.; |
| 194 | } |
| 195 | } |
| 196 | } |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 197 | |
| 198 | if (improveMask) { |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 199 | cv::Mat dilate; |
| 200 | cv::dilate(maskImg, |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 201 | dilate, |
| 202 | cv::getStructuringElement(cv::MORPH_ELLIPSE, kSize_), |
| 203 | cv::Point(-1, -1), |
| 204 | 2); |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 205 | cv::erode(maskImg, |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 206 | maskImg, |
| 207 | cv::getStructuringElement(cv::MORPH_ELLIPSE, kSize_), |
| 208 | cv::Point(-1, -1), |
| 209 | 2); |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 210 | for (int i = 0; i < maskImg.cols; i++) { |
| 211 | for (int j = 0; j < maskImg.rows; j++) { |
| 212 | if (dilate.at<float>(j, i) != maskImg.at<float>(j, i)) |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 213 | maskImg.at<float>(j, i) = grabcutClass_; |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 214 | } |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 215 | } |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 216 | cv::Mat applyMask = cvFrame_.clone(); |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 217 | maskImg.convertTo(maskImg, CV_8UC1); |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 218 | applyMask.convertTo(applyMask, CV_8UC1); |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 219 | cv::Rect rect(1, 1, maskImg.rows, maskImg.cols); |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 220 | cv::grabCut(applyMask, |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 221 | maskImg, |
| 222 | rect, |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 223 | bgdModel_, |
| 224 | fgdModel_, |
| 225 | grabCutIterations_, |
| 226 | grabCutMode_); |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 227 | |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 228 | grabCutMode_ = cv::GC_EVAL; |
| 229 | grabCutIterations_ = 1; |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 230 | |
| 231 | maskImg = maskImg & 1; |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 232 | maskImg.convertTo(maskImg, CV_32FC1); |
| 233 | maskImg *= 255.; |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 234 | blur(maskImg, maskImg, cv::Size(7, 7)); // float mask from 0 to 255. |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 235 | maskImg = maskImg / 255.; |
| 236 | } |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 237 | |
| 238 | previousMasks_[1] = previousMasks_[0].clone(); |
| 239 | previousMasks_[0] = maskImg.clone(); |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 240 | } |
| 241 | |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 242 | cv::Mat roiMaskImg = previousMasks_[0].clone() * 255.; |
| 243 | roiMaskImg.convertTo(roiMaskImg, CV_8UC1); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 244 | |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 245 | gsFrame maskFrame = {av_frame_alloc(), frameFree}; |
| 246 | maskFrame->format = AV_PIX_FMT_GRAY8; |
| 247 | maskFrame->width = roiMaskImg.cols; |
| 248 | maskFrame->height = roiMaskImg.rows; |
| 249 | maskFrame->linesize[0] = roiMaskImg.step; |
| 250 | if (av_frame_get_buffer(maskFrame.get(), 0) < 0) |
| 251 | return; |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 252 | |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 253 | maskFrame->data[0] = roiMaskImg.data; |
| 254 | maskFrame->pts = 1; |
| 255 | mainFilter_.feedInput(maskFrame.get(), "mask"); |
| 256 | maskFrame.reset(); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 257 | |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 258 | if (isBlur_) |
| 259 | mainFilter_.feedInput(frameReduced, "input"); |
| 260 | mainFilter_.feedInput(frame, "input2"); |
| 261 | AVFrame* filteredFrame; |
| 262 | if ((filteredFrame = mainFilter_.readOutput())) { |
| 263 | moveFrom(frame, filteredFrame); |
| 264 | frameFree(filteredFrame); |
| 265 | } |
| 266 | count_++; |
| 267 | count_ = count_ % frameCount_; |
| 268 | } |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 269 | |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 270 | MediaStream |
| 271 | PluginProcessor::getbgAVFrameInfos() |
| 272 | { |
| 273 | AVFormatContext* ctx = avformat_alloc_context(); |
| 274 | // Open |
| 275 | if (avformat_open_input(&ctx, backgroundPath_.c_str(), NULL, NULL) != 0) { |
| 276 | avformat_free_context(ctx); |
| 277 | Plog::log(Plog::LogPriority::INFO, TAG, "Couldn't open input stream."); |
| 278 | return {}; |
| 279 | } |
| 280 | pFormatCtx_.reset(ctx); |
| 281 | // Retrieve stream information |
| 282 | if (avformat_find_stream_info(pFormatCtx_.get(), NULL) < 0) { |
| 283 | Plog::log(Plog::LogPriority::INFO, TAG, "Couldn't find stream information."); |
| 284 | return {}; |
| 285 | } |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 286 | |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 287 | // Dump valid information onto standard error |
| 288 | av_dump_format(pFormatCtx_.get(), 0, backgroundPath_.c_str(), false); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 289 | |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 290 | // Find the video stream |
| 291 | for (int i = 0; i < static_cast<int>(pFormatCtx_->nb_streams); i++) |
| 292 | if (pFormatCtx_->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) { |
| 293 | videoStream_ = i; |
| 294 | break; |
| 295 | } |
| 296 | |
| 297 | if (videoStream_ == -1) { |
| 298 | Plog::log(Plog::LogPriority::INFO, TAG, "Didn't find a video stream."); |
| 299 | return {}; |
| 300 | } |
| 301 | |
| 302 | rational<int> fr = pFormatCtx_->streams[videoStream_]->r_frame_rate; |
| 303 | return MediaStream("background", |
| 304 | pFormatCtx_->streams[videoStream_]->codecpar->format, |
| 305 | 1 / fr, |
| 306 | pFormatCtx_->streams[videoStream_]->codecpar->width, |
| 307 | pFormatCtx_->streams[videoStream_]->codecpar->height, |
| 308 | 0, |
| 309 | fr); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 310 | } |
| 311 | |
| 312 | void |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 313 | PluginProcessor::loadBackground() |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 314 | { |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 315 | if (backgroundPath_.empty()) |
| 316 | return; |
| 317 | |
| 318 | auto bgStream_ = getbgAVFrameInfos(); |
| 319 | if (mainFilter_.initialize(mainFilterDescription_, {maskms_, bgStream_, ims2_}) < 0) |
| 320 | return; |
| 321 | |
| 322 | int got_frame; |
| 323 | AVCodecContext* pCodecCtx; |
| 324 | AVCodec* pCodec; |
| 325 | AVPacket* packet; |
| 326 | |
| 327 | pCodecCtx = pFormatCtx_->streams[videoStream_]->codec; |
| 328 | pCodec = avcodec_find_decoder(pCodecCtx->codec_id); |
| 329 | if (pCodec == nullptr) { |
| 330 | mainFilter_.clean(); |
| 331 | pFormatCtx_.reset(); |
| 332 | Plog::log(Plog::LogPriority::INFO, TAG, "Codec not found."); |
| 333 | return; |
| 334 | } |
| 335 | |
| 336 | // Open codec |
| 337 | if (avcodec_open2(pCodecCtx, pCodec, NULL) < 0) { |
| 338 | mainFilter_.clean(); |
| 339 | pFormatCtx_.reset(); |
| 340 | Plog::log(Plog::LogPriority::INFO, TAG, "Could not open codec."); |
| 341 | return; |
| 342 | } |
| 343 | |
| 344 | packet = av_packet_alloc(); |
| 345 | if (av_read_frame(pFormatCtx_.get(), packet) < 0) { |
| 346 | mainFilter_.clean(); |
| 347 | avcodec_close(pCodecCtx); |
| 348 | av_packet_free(&packet); |
| 349 | pFormatCtx_.reset(); |
| 350 | Plog::log(Plog::LogPriority::INFO, TAG, "Could not read packet from context."); |
| 351 | return; |
| 352 | } |
| 353 | |
| 354 | AVFrame* bgImage = av_frame_alloc(); |
| 355 | avcodec_decode_video2(pCodecCtx, bgImage, &got_frame, packet); |
| 356 | if (got_frame) { |
| 357 | mainFilter_.feedInput(bgImage, "background"); |
| 358 | mainFilter_.feedEOF("background"); |
| 359 | } else |
| 360 | mainFilter_.clean(); |
| 361 | |
| 362 | frameFree(bgImage); |
| 363 | avcodec_close(pCodecCtx); |
| 364 | av_packet_free(&packet); |
| 365 | pFormatCtx_.reset(); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 366 | } |
agsantos | 9dcf430 | 2020-09-01 18:21:48 -0400 | [diff] [blame] | 367 | |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 368 | void |
| 369 | PluginProcessor::initFilters(const std::pair<int, int>& inputSize, int format, int angle) |
agsantos | 9dcf430 | 2020-09-01 18:21:48 -0400 | [diff] [blame] | 370 | { |
agsantos | dd6a62a | 2021-03-29 17:13:27 -0400 | [diff] [blame] | 371 | resetInitValues(); |
| 372 | mainFilter_.clean(); |
| 373 | |
| 374 | std::string rotateSides = ""; |
| 375 | std::string scaleSize = std::to_string(inputSize.first) + ":" |
| 376 | + std::to_string(inputSize.second); |
| 377 | Plog::log(Plog::LogPriority::INFO, TAG, scaleSize); |
| 378 | if (std::abs(angle) == 90) { |
| 379 | rotateSides = ":out_w=ih:out_h=iw"; |
| 380 | scaleSize = std::to_string(inputSize.second) + ":" + std::to_string(inputSize.first); |
| 381 | } |
| 382 | |
| 383 | rational<int> fr(1, 1); |
| 384 | ims_ = MediaStream("input", |
| 385 | AV_PIX_FMT_RGB24, |
| 386 | 1 / fr, |
| 387 | modelInputDimensions.first, |
| 388 | modelInputDimensions.second, |
| 389 | 0, |
| 390 | fr); |
| 391 | ims2_ = MediaStream("input2", format, 1 / fr, inputSize.first, inputSize.second, 0, fr); |
| 392 | |
| 393 | maskms_ = MediaStream("mask", |
| 394 | AV_PIX_FMT_GRAY8, |
| 395 | 1 / fr, |
| 396 | modelInputDimensions.first, |
| 397 | modelInputDimensions.second, |
| 398 | 0, |
| 399 | fr); |
| 400 | |
| 401 | if (isBlur_) { |
| 402 | mainFilterDescription_ = "[mask]negate[negated],[input][negated]alphamerge,boxblur=" |
| 403 | + blurLevel_ + ",scale=" + scaleSize |
| 404 | + "[blured],[input2]format=rgb24,rotate=" + rotation[-angle] |
| 405 | + rotateSides |
| 406 | + "[input2formated],[input2formated][blured]overlay,rotate=" |
| 407 | + rotation[angle] + rotateSides; |
| 408 | Plog::log(Plog::LogPriority::INFO, TAG, mainFilterDescription_); |
| 409 | mainFilter_.initialize(mainFilterDescription_, {maskms_, ims_, ims2_}); |
| 410 | } else { |
| 411 | mainFilterDescription_ = "[mask]scale=" + scaleSize |
| 412 | + "[fgmask],[fgmask]split=2[bg][fg],[bg]negate[bgn]," |
| 413 | + "[background]scale=" + scaleSize + "[backgroundformated]," |
| 414 | + "[backgroundformated][bgn]alphamerge[background2]," |
| 415 | + "[input2]rotate=" + rotation[-angle] + rotateSides |
| 416 | + "[input2formated],[input2formated][fg]alphamerge[foreground]," |
| 417 | + "[foreground][background2]overlay," + "rotate=" + rotation[angle] |
| 418 | + rotateSides; |
| 419 | Plog::log(Plog::LogPriority::INFO, TAG, mainFilterDescription_); |
| 420 | loadBackground(); |
| 421 | } |
agsantos | 9dcf430 | 2020-09-01 18:21:48 -0400 | [diff] [blame] | 422 | } |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 423 | } // namespace jami |