agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 1 | /** |
| 2 | * Copyright (C) 2020 Savoir-faire Linux Inc. |
| 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" |
| 23 | // System includes |
| 24 | #include <algorithm> |
| 25 | #include <cstring> |
| 26 | // OpenCV headers |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 27 | #include <opencv2/core.hpp> |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 28 | #include <opencv2/imgcodecs.hpp> |
| 29 | #include <opencv2/imgproc.hpp> |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 30 | // Logger |
| 31 | #include <pluglog.h> |
| 32 | |
| 33 | extern "C" { |
| 34 | #include <libavutil/display.h> |
| 35 | } |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 36 | const char sep = separator(); |
| 37 | |
| 38 | const std::string TAG = "FORESEG"; |
| 39 | |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 40 | namespace jami { |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 41 | |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 42 | PluginProcessor::PluginProcessor(const std::string& dataPath, const std::string& model, const std::string& backgroundImage, bool acc) |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 43 | { |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 44 | activateAcc_ = acc; |
| 45 | initModel(dataPath+sep+"model/"+model); |
| 46 | setBackgroundImage(backgroundImage); |
| 47 | } |
| 48 | |
| 49 | PluginProcessor::~PluginProcessor() |
| 50 | { |
| 51 | Plog::log(Plog::LogPriority::INFO, TAG, "~pluginprocessor"); |
| 52 | if (session_) |
| 53 | delete session_; |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 54 | } |
| 55 | |
| 56 | void |
agsantos | 9dcf430 | 2020-09-01 18:21:48 -0400 | [diff] [blame] | 57 | PluginProcessor::setBackgroundImage(const std::string& backgroundPath) |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 58 | { |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 59 | cv::Size size = cv::Size {0, 0}; |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 60 | |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 61 | if (!backgroundImage.empty()) |
| 62 | size = backgroundImage.size(); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 63 | |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 64 | cv::Mat newBackgroundImage = cv::imread(backgroundPath); |
| 65 | if (newBackgroundImage.cols == 0) { |
| 66 | Plog::log(Plog::LogPriority::ERR, TAG, "Background image not Loaded"); |
| 67 | } else { |
| 68 | Plog::log(Plog::LogPriority::INFO, TAG, "Background image Loaded"); |
| 69 | cv::cvtColor(newBackgroundImage, newBackgroundImage, cv::COLOR_BGR2RGB); |
| 70 | newBackgroundImage.convertTo(newBackgroundImage, CV_32FC3); |
| 71 | if (size.height) { |
| 72 | cv::resize(newBackgroundImage, newBackgroundImage, size); |
| 73 | backgroundRotation = 0; |
| 74 | } |
| 75 | backgroundImage = newBackgroundImage.clone(); |
| 76 | newBackgroundImage.release(); |
| 77 | hasBackground_ = true; |
| 78 | } |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 79 | } |
| 80 | |
| 81 | void |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 82 | PluginProcessor::initModel(const std::string& modelPath) |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 83 | { |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 84 | try { |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 85 | auto allocator_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); |
| 86 | input_tensor_ = Ort::Value::CreateTensor<float>(allocator_info, input_image_.data(), input_image_.size(), input_shape_.data(), input_shape_.size()); |
| 87 | output_tensor_ = Ort::Value::CreateTensor<float>(allocator_info, results_.data(), results_.size(), output_shape_.data(), output_shape_.size()); |
| 88 | sessOpt_ = Ort::SessionOptions(); |
| 89 | |
| 90 | #ifdef NVIDIA |
| 91 | if (activateAcc_) |
| 92 | Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CUDA(sessOpt_, 0)); |
| 93 | #endif |
| 94 | #ifdef ANDROID |
| 95 | if (activateAcc_) |
| 96 | Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_Nnapi(sessOpt_, 0)); |
| 97 | #endif |
| 98 | |
| 99 | sessOpt_.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL); |
| 100 | #ifdef WIN32 |
| 101 | std::wstring wsTmp(modelPath.begin(), modelPath.end()); |
| 102 | session_ = new Ort::Session(env, wsTmp.c_str(), sessOpt_); |
| 103 | #else |
| 104 | session_ = new Ort::Session(env, modelPath.c_str(), sessOpt_); |
| 105 | #endif |
| 106 | isAllocated_ = true; |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 107 | } catch (std::exception& e) { |
| 108 | Plog::log(Plog::LogPriority::ERR, TAG, e.what()); |
| 109 | } |
| 110 | std::ostringstream oss; |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 111 | oss << "Model is allocated " << isAllocated_; |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 112 | Plog::log(Plog::LogPriority::INFO, TAG, oss.str()); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 113 | } |
| 114 | |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 115 | bool |
| 116 | PluginProcessor::isAllocated() |
| 117 | { |
| 118 | return isAllocated_; |
| 119 | } |
| 120 | |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 121 | void |
| 122 | PluginProcessor::feedInput(const cv::Mat& frame) |
| 123 | { |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 124 | cv::Mat temp(frame.rows, frame.cols, CV_32FC3, input_image_.data()); |
| 125 | frame.convertTo(temp, CV_32FC3); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 126 | } |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 127 | |
| 128 | int |
| 129 | PluginProcessor::getBackgroundRotation() |
| 130 | { |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 131 | return backgroundRotation; |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 132 | } |
| 133 | |
| 134 | void |
| 135 | PluginProcessor::setBackgroundRotation(int angle) |
| 136 | { |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 137 | if (backgroundRotation != angle && (backgroundRotation - angle) != 0) { |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 138 | rotateFrame(backgroundRotation - angle, backgroundImage); |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 139 | backgroundRotation = angle; |
| 140 | } |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 141 | } |
| 142 | |
| 143 | void |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 144 | PluginProcessor::computePredictions() |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 145 | { |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 146 | if (count == 0) { |
| 147 | // Run the graph |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 148 | session_->Run(Ort::RunOptions{nullptr}, input_names, &input_tensor_, 1, output_names, &output_tensor_, 1); |
| 149 | computedMask = std::vector(results_.begin(), results_.end()); |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 150 | } |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 151 | } |
| 152 | |
| 153 | void |
| 154 | PluginProcessor::printMask() |
| 155 | { |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 156 | for (size_t i = 0; i < computedMask.size(); i++) { |
| 157 | // Log the predictions |
| 158 | std::ostringstream oss; |
| 159 | oss << "\nclass: " << computedMask[i] << std::endl; |
| 160 | Plog::log(Plog::LogPriority::INFO, TAG, oss.str()); |
| 161 | } |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 162 | } |
| 163 | |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 164 | void |
agsantos | 31d1a0b | 2020-10-23 14:05:53 -0400 | [diff] [blame] | 165 | PluginProcessor::resetInitValues(const cv::Size& modelInputSize) |
| 166 | { |
| 167 | previousMasks[0] = cv::Mat(modelInputSize.height, modelInputSize.width, CV_32FC1, double(0.)); |
| 168 | previousMasks[1] = cv::Mat(modelInputSize.height, modelInputSize.width, CV_32FC1, double(0.)); |
| 169 | kSize = cv::Size(modelInputSize.width * kernelSize, modelInputSize.height * kernelSize); |
| 170 | if (kSize.height % 2 == 0) { |
| 171 | kSize.height -= 1; |
| 172 | } |
| 173 | if (kSize.width % 2 == 0) { |
| 174 | kSize.width -= 1; |
| 175 | } |
| 176 | count = 0; |
| 177 | grabCutMode = cv::GC_INIT_WITH_MASK; |
| 178 | grabCutIterations = 5; |
| 179 | } |
| 180 | |
| 181 | void |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 182 | copyByLine(uchar* frameData, uchar* applyMaskData, const int lineSize, cv::Size size) |
| 183 | { |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 184 | if (3 * size.width == lineSize) { |
| 185 | std::memcpy(frameData, applyMaskData, size.height * size.width * 3); |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 186 | } else { |
| 187 | int rows = size.height; |
| 188 | int offset = 0; |
| 189 | int maskoffset = 0; |
| 190 | for (int i = 0; i < rows; i++) { |
| 191 | std::memcpy(frameData + offset, applyMaskData + maskoffset, lineSize); |
| 192 | offset += lineSize; |
| 193 | maskoffset += 3 * size.width; |
| 194 | } |
| 195 | } |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 196 | } |
| 197 | |
| 198 | void |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 199 | PluginProcessor::drawMaskOnFrame( |
| 200 | cv::Mat& frame, cv::Mat& frameReduced, std::vector<float> computedMask, int lineSize, int angle) |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 201 | { |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 202 | if (computedMask.empty()) { |
| 203 | return; |
| 204 | } |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 205 | |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 206 | if (count == 0) { |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 207 | int maskSize = static_cast<int>(std::sqrt(computedMask.size())); |
| 208 | cv::Mat maskImg(maskSize, maskSize, CV_32FC1, computedMask.data()); |
| 209 | cv::Mat* applyMask = &frameReduced; |
| 210 | |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 211 | rotateFrame(-angle, maskImg); |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 212 | cv::resize(maskImg, maskImg, cv::Size(frameReduced.cols, frameReduced.rows)); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 213 | |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 214 | double m, M; |
| 215 | cv::minMaxLoc(maskImg, &m, &M); |
| 216 | |
| 217 | if (M < 2) { // avoid detection if there is any one in frame |
| 218 | maskImg = 0. * maskImg; |
| 219 | } else { |
| 220 | for (int i = 0; i < maskImg.cols; i++) { |
| 221 | for (int j = 0; j < maskImg.rows; j++) { |
| 222 | maskImg.at<float>(j, i) = (maskImg.at<float>(j, i) - m) / (M - m); |
| 223 | |
| 224 | if (maskImg.at<float>(j, i) < 0.4) |
| 225 | maskImg.at<float>(j, i) = 0.; |
| 226 | else if (maskImg.at<float>(j, i) < 0.7) { |
| 227 | float value = maskImg.at<float>(j, i) * smoothFactors[0] |
| 228 | + previousMasks[0].at<float>(j, i) * smoothFactors[1] |
| 229 | + previousMasks[1].at<float>(j, i) * smoothFactors[2]; |
| 230 | maskImg.at<float>(j, i) = 0.; |
| 231 | if (value > 0.7) |
| 232 | maskImg.at<float>(j, i) = 1.; |
| 233 | } else |
| 234 | maskImg.at<float>(j, i) = 1.; |
| 235 | } |
| 236 | } |
| 237 | } |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 238 | if (cv::countNonZero(maskImg) != 0) { |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 239 | cv::Mat dilate; |
| 240 | cv::dilate(maskImg, |
| 241 | dilate, |
| 242 | cv::getStructuringElement(cv::MORPH_ELLIPSE, kSize), |
| 243 | cv::Point(-1, -1), |
| 244 | 2); |
| 245 | cv::erode(maskImg, |
| 246 | maskImg, |
| 247 | cv::getStructuringElement(cv::MORPH_ELLIPSE, kSize), |
| 248 | cv::Point(-1, -1), |
| 249 | 2); |
| 250 | for (int i = 0; i < maskImg.cols; i++) { |
| 251 | for (int j = 0; j < maskImg.rows; j++) { |
| 252 | if (dilate.at<float>(j, i) != maskImg.at<float>(j, i)) |
| 253 | maskImg.at<float>(j, i) = grabcutClass; |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 254 | } |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 255 | } |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 256 | maskImg.convertTo(maskImg, CV_8UC1); |
| 257 | applyMask->convertTo(*applyMask, CV_8UC1); |
| 258 | cv::Rect rect(1, 1, maskImg.rows, maskImg.cols); |
| 259 | cv::grabCut(*applyMask, |
| 260 | maskImg, |
| 261 | rect, |
| 262 | bgdModel, |
| 263 | fgdModel, |
| 264 | grabCutIterations, |
| 265 | grabCutMode); |
| 266 | |
| 267 | grabCutMode = cv::GC_EVAL; |
| 268 | grabCutIterations = 1; |
| 269 | |
| 270 | maskImg = maskImg & 1; |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 271 | maskImg.convertTo(maskImg, CV_32FC1); |
| 272 | maskImg *= 255.; |
agsantos | 31d1a0b | 2020-10-23 14:05:53 -0400 | [diff] [blame] | 273 | GaussianBlur(maskImg, maskImg, cv::Size(7, 7), 0); // float mask from 0 to 255. |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 274 | maskImg = maskImg / 255.; |
| 275 | } |
| 276 | previousMasks[1] = previousMasks[0].clone(); |
| 277 | previousMasks[0] = maskImg.clone(); |
| 278 | } |
| 279 | |
| 280 | cv::Mat roiMaskImg = previousMasks[0].clone(); |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 281 | cv::Mat roiMaskImgComplementary = 1. - roiMaskImg; // mask from 1. to 0 |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 282 | |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 283 | std::vector<cv::Mat> channels; |
| 284 | std::vector<cv::Mat> channelsComplementary; |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 285 | |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 286 | channels.emplace_back(roiMaskImg); |
| 287 | channels.emplace_back(roiMaskImg); |
| 288 | channels.emplace_back(roiMaskImg); |
| 289 | channelsComplementary.emplace_back(roiMaskImgComplementary); |
| 290 | channelsComplementary.emplace_back(roiMaskImgComplementary); |
| 291 | channelsComplementary.emplace_back(roiMaskImgComplementary); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 292 | |
agsantos | ac1940d | 2020-09-17 10:18:40 -0400 | [diff] [blame] | 293 | cv::merge(channels, roiMaskImg); |
| 294 | cv::merge(channelsComplementary, roiMaskImgComplementary); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 295 | |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 296 | cv::Mat output; |
| 297 | frameReduced.convertTo(output, roiMaskImg.type()); |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 298 | output = output.mul(roiMaskImg); |
| 299 | output += backgroundImage.mul(roiMaskImgComplementary); |
agsantos | 796b5af | 2020-12-22 19:38:27 -0500 | [diff] [blame] | 300 | output.convertTo(output, frameReduced.type()); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 301 | |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 302 | cv::resize(output, output, cv::Size(frame.cols, frame.rows)); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 303 | |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 304 | copyByLine(frame.data, output.data, lineSize, cv::Size(frame.cols, frame.rows)); |
| 305 | count++; |
| 306 | count = count % frameCount; |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 307 | } |
| 308 | |
| 309 | void |
| 310 | PluginProcessor::rotateFrame(int angle, cv::Mat& mat) |
| 311 | { |
agsantos | b74f4cb | 2020-10-01 14:30:43 -0400 | [diff] [blame] | 312 | if (angle == -90) |
| 313 | cv::rotate(mat, mat, cv::ROTATE_90_COUNTERCLOCKWISE); |
| 314 | else if (std::abs(angle) == 180) |
| 315 | cv::rotate(mat, mat, cv::ROTATE_180); |
| 316 | else if (angle == 90) |
| 317 | cv::rotate(mat, mat, cv::ROTATE_90_CLOCKWISE); |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 318 | } |
agsantos | 9dcf430 | 2020-09-01 18:21:48 -0400 | [diff] [blame] | 319 | |
| 320 | bool |
| 321 | PluginProcessor::hasBackground() const |
| 322 | { |
| 323 | return hasBackground_; |
| 324 | } |
agsantos | 5aa3965 | 2020-08-11 18:18:04 -0400 | [diff] [blame] | 325 | } // namespace jami |