#include "ncnn/net.h"
#if defined(USE_NCNN_SIMPLEOCV)
#include "simpleocv.h"
#else
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#endif
#include <stdlib.h>
#include <float.h>
#include <stdio.h>
#include <vector>
#include <opencv2/imgproc/types_c.h>
#include<iostream>
using namespace std;
void toPng(cv::Mat& src, cv::Mat& dst, int mark)
{
cv::Mat cv_input = src.clone();
if (cv_input.channels() != 4)
{
cv::cvtColor(cv_input, dst, CV_BGR2BGRA);
}
else
{
return;
}
for (int y = 0; y < dst.rows; ++y)
{
for (int x = 0; x < dst.cols; ++x)
{
cv::Vec4b& pixel = dst.at<cv::Vec4b>(y, x);
// std::cout <<"pixel[0]= "<< pixel[0] << "pixel[1]= " << pixel[1] << "pixel[2]= " << pixel[2] << std::endl;
if (pixel[0] == mark && pixel[1] == mark && pixel[2] == mark)
{
pixel[3] = 0;
}
}
}
}
static void draw_objects(const cv::Mat& bgr, const cv::Mat& fgr, const cv::Mat& pha)
{
cv::Mat fgr8U;
fgr.convertTo(fgr8U, CV_8UC3, 255.0, 0);
cv::Mat pha8U;
pha.convertTo(pha8U, CV_8UC1, 255.0, 0);
cv::Mat comp;
cv::resize(bgr, comp, pha.size(), 0, 0, 1);
for (int i = 0; i < pha8U.rows; i++)
{
for (int j = 0; j < pha8U.cols; j++)
{
uchar data = pha8U.at<uchar>(i, j);
float alpha = (float)data / 255;
comp.at<cv::Vec3b>(i, j)[0] = fgr8U.at<cv::Vec3b>(i, j)[0] * alpha + (1 - alpha) * 255;//155
comp.at<cv::Vec3b>(i, j)[1] = fgr8U.at<cv::Vec3b>(i, j)[1] * alpha + (1 - alpha) * 255;//255
comp.at<cv::Vec3b>(i, j)[2] = fgr8U.at<cv::Vec3b>(i, j)[2] * alpha + (1 - alpha) * 255;//120
// std::cout <<(int)comp.at<cv::Vec3b>(i, j)[0] << " " << (int)comp.at<cv::Vec3b>(i, j)[1] <<" "<<(int)comp.at<cv::Vec3b>(i, j)[2] << std::endl;
}
}
// cv::imshow("pha", pha8U);
//cv::imshow("fgr", fgr8U);
//cv::Mat img1_t1(comp, cv::Rect(0, 0, comp.cols, comp.rows));
// toPng(comp, img1_t1, 0);
// cv::imshow("comp", comp);
cv::Mat img_alpha_0;
//
cv::resize(comp, comp,cv::Size(bgr.cols,bgr.rows), 0, 0, cv::INTER_LINEAR);
toPng(comp, img_alpha_0, 255);
imwrite("F:\\boost\\b.jpg", img_alpha_0);
cv::waitKey(0);
}
static int detect_rvm(const cv::Mat& bgr, cv::Mat& pha, cv::Mat& fgr)
{
const float downsample_ratio = 0.5f;
const int target_width = 512;
const int target_height = 512;
ncnn::Net net;
net.opt.use_vulkan_compute = false;
//original pretrained model from https://github.com/PeterL1n/RobustVideoMatting
//ncnn model https://pan.baidu.com/s/11iEY2RGfzWFtce8ue7T3JQ password: d9t6
net.load_param("G:\\rvm_ncnn_models\\rvm_512.param");
net.load_model("G:\\rvm_ncnn_models\\rvm_512.bin");
//if you use another input size,pleaze change input shape
ncnn::Mat r1i = ncnn::Mat(128, 128, 16);
ncnn::Mat r2i = ncnn::Mat(64, 64, 20);
ncnn::Mat r3i = ncnn::Mat(32, 32, 40);
ncnn::Mat r4i = ncnn::Mat(16, 16, 64);
r1i.fill(0.0f);
r2i.fill(0.0f);
r3i.fill(0.0f);
r4i.fill(0.0f);
ncnn::Extractor ex = net.create_extractor();
const float mean_vals1[3] = { 123.675f, 116.28f, 103.53f };
const float norm_vals1[3] = { 0.01712475f, 0.0175f, 0.01742919f };
const float mean_vals2[3] = { 0, 0, 0 };
const float norm_vals2[3] = { 1 / 255.0, 1 / 255.0, 1 / 255.0 };
ncnn::Mat ncnn_in2 = ncnn::Mat::from_pixels_resize(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, bgr.cols, bgr.rows, target_width, target_height);
ncnn::Mat ncnn_in1 = ncnn::Mat::from_pixels_resize(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, bgr.cols, bgr.rows, target_width * downsample_ratio, target_height * downsample_ratio);
ncnn_in1.substract_mean_normalize(mean_vals1, norm_vals1);
ncnn_in2.substract_mean_normalize(mean_vals2, norm_vals2);
ex.input("src1", ncnn_in1);
ex.input("src2", ncnn_in2);
ex.input("r1i", r1i);
ex.input("r2i", r2i);
ex.input("r3i", r3i);
ex.input("r4i", r4i);
//if use video matting,these output will be input of next infer
ex.extract("r4o", r4i);
ex.extract("r3o", r3i);
ex.extract("r2o", r2i);
ex.extract("r1o", r1i);
ncnn::Mat pha_;
ex.extract("pha", pha_);
ncnn::Mat fgr_;
ex.extract("fgr", fgr_);
cv::Mat cv_pha = cv::Mat(pha_.h, pha_.w, CV_32FC1, (float*)pha_.data);
cv::Mat cv_fgr = cv::Mat(fgr_.h, fgr_.w, CV_32FC3);
float* fgr_data = (float*)fgr_.data;
for (int i = 0; i < fgr_.h; i++)
{
for (int j = 0; j < fgr_.w; j++)
{
cv_fgr.at<cv::Vec3f>(i, j)[2] = fgr_data[0 * fgr_.h * fgr_.w + i * fgr_.w + j];
cv_fgr.at<cv::Vec3f>(i, j)[1] = fgr_data[1 * fgr_.h * fgr_.w + i * fgr_.w + j];
cv_fgr.at<cv::Vec3f>(i, j)[0] = fgr_data[2 * fgr_.h * fgr_.w + i * fgr_.w + j];
}
}
cv_pha.copyTo(pha);
cv_fgr.copyTo(fgr);
return 0;
}
int main(int argc, char** argv)
{
cv::Mat m = cv::imread("F:\\boost\\a.jpg");
if (m.empty())
{
return -1;
}
cv::Mat fgr, pha;
detect_rvm(m, pha, fgr);
draw_objects(m, fgr, pha);
return 0;
}
