mirror of
https://github.com/Tencent/ncnn.git
synced 2026-10-18 13:42:18 +08:00
nihui
337 lines
9.7 KiB
C++
337 lines
9.7 KiB
C++
// Copyright 2025 Tencent
|
|
// SPDX-License-Identifier: BSD-3-Clause
|
|
|
|
// ncnn model exported from https://github.com/PeterL1n/RobustVideoMatting
|
|
//
|
|
// import torch
|
|
// from torch import nn
|
|
// from model import MattingNetwork
|
|
// from model.fast_guided_filter import FastGuidedFilterRefiner
|
|
// from model.deep_guided_filter import DeepGuidedFilterRefiner
|
|
//
|
|
// class Model(nn.Module):
|
|
// def __init__(self):
|
|
// super().__init__()
|
|
//
|
|
// self.rvm = MattingNetwork('mobilenetv3').eval()
|
|
// self.rvm.load_state_dict(torch.load('rvm_mobilenetv3.pth'))
|
|
//
|
|
// self.refiner_deep = DeepGuidedFilterRefiner()
|
|
// self.refiner_fast = FastGuidedFilterRefiner()
|
|
//
|
|
// def forward_first_frame(self, src):
|
|
// return self.rvm(src)
|
|
//
|
|
// def forward(self, src, src_sm, r1, r2, r3, r4):
|
|
//
|
|
// f1, f2, f3, f4 = self.rvm.backbone(src_sm)
|
|
// f4 = self.rvm.aspp(f4)
|
|
// hid, *rec = self.rvm.decoder(src_sm, f1, f2, f3, f4, r1, r2, r3, r4)
|
|
//
|
|
// # downsample
|
|
// fgr_residual, pha = self.rvm.project_mat(hid).split([3, 1], dim=-3)
|
|
// fgr = fgr_residual + src_sm
|
|
//
|
|
// # downsample + refiner_deep
|
|
// fgr_residual_deep, pha_deep = self.refiner_deep(src, src_sm, fgr_residual, pha, hid)
|
|
// fgr_deep = fgr_residual_deep + src
|
|
//
|
|
// # downsample + refiner_fast
|
|
// fgr_residual_fast, pha_fast = self.refiner_fast(src, src_sm, fgr_residual, pha, hid)
|
|
// fgr_fast = fgr_residual_fast + src
|
|
//
|
|
// # downsample + segmentation
|
|
// seg = self.rvm.project_seg(hid)
|
|
//
|
|
// return fgr, pha, fgr_deep, pha_deep, fgr_fast, pha_fast, seg, *rec
|
|
//
|
|
// import pnnx
|
|
//
|
|
// model = Model().eval()
|
|
//
|
|
// x = torch.rand(1, 3, 512, 512)
|
|
// x2 = torch.rand(1, 3, 256, 256)
|
|
// x2_hr = torch.rand(1, 3, 1024, 1024)
|
|
//
|
|
// # generate feats via forward_first_frame, with different shapes
|
|
// fgr, pha, r1, r2, r3, r4 = model.forward_first_frame(x)
|
|
// fgr2, pha2, r12, r22, r32, r42 = model.forward_first_frame(x2)
|
|
//
|
|
// # export with dynamic shape
|
|
// pnnx.export(model, "rvm_mobilenetv3.pt", (x, x, r1, r2, r3, r4), (x2_hr, x2, r12, r22, r32, r42))
|
|
//
|
|
// and then fix refiner_fast fp16 overflow issue in ncnn.param via appending 31=1 layer feat mask
|
|
//
|
|
// BinaryOp div_58 2 1 401 399 402 0=3 31=1
|
|
//
|
|
|
|
#include "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
|
|
|
|
static int detect_rvm(const cv::Mat& bgr, cv::Mat& fgr, cv::Mat& pha, cv::Mat& seg)
|
|
{
|
|
ncnn::Net rvm;
|
|
|
|
rvm.opt.use_vulkan_compute = true;
|
|
|
|
// https://github.com/nihui/ncnn-android-rvm/tree/master/app/src/main/assets
|
|
// you shall also change r1,r2,r3,r4 shape below when model changed
|
|
if (rvm.load_param("rvm_mobilenetv3.ncnn.param"))
|
|
exit(-1);
|
|
if (rvm.load_model("rvm_mobilenetv3.ncnn.bin"))
|
|
exit(-1);
|
|
// if (rvm.load_param("rvm_resnet50.ncnn.param"))
|
|
// exit(-1);
|
|
// if (rvm.load_model("rvm_resnet50.ncnn.bin"))
|
|
// exit(-1);
|
|
|
|
const int w = bgr.cols;
|
|
const int h = bgr.rows;
|
|
|
|
const int target_size = 512;
|
|
const int max_stride = 16;
|
|
|
|
bool refine_deep = true;
|
|
// bool refine_fast = true;
|
|
|
|
const float norm_vals[3] = {1 / 255.f, 1 / 255.f, 1 / 255.f};
|
|
|
|
ncnn::Mat in_pad;
|
|
ncnn::Mat in_small_pad;
|
|
|
|
int wpad = 0;
|
|
int hpad = 0;
|
|
|
|
bool downsample = std::max(w, h) > target_size;
|
|
if (downsample)
|
|
{
|
|
// letterbox pad to multiple of max_stride
|
|
int w2 = w;
|
|
int h2 = h;
|
|
float scale = 1.f;
|
|
if (w > h)
|
|
{
|
|
scale = (float)target_size / w;
|
|
w2 = target_size;
|
|
h2 = h2 * scale;
|
|
}
|
|
else
|
|
{
|
|
scale = (float)target_size / h;
|
|
h2 = target_size;
|
|
w2 = w2 * scale;
|
|
}
|
|
|
|
ncnn::Mat in_small = ncnn::Mat::from_pixels_resize(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, w, h, w2, h2);
|
|
|
|
// letterbox pad to target_size rectangle
|
|
int w2pad = (w2 + max_stride - 1) / max_stride * max_stride - w2;
|
|
int h2pad = (h2 + max_stride - 1) / max_stride * max_stride - h2;
|
|
ncnn::copy_make_border(in_small, in_small_pad, h2pad / 2, h2pad - h2pad / 2, w2pad / 2, w2pad - w2pad / 2, ncnn::BORDER_CONSTANT, 114.f);
|
|
|
|
in_small_pad.substract_mean_normalize(0, norm_vals);
|
|
|
|
int w3 = w;
|
|
int h3 = h;
|
|
if (w > h)
|
|
{
|
|
w3 = w;
|
|
h3 = in_small_pad.h / scale;
|
|
wpad = 0;
|
|
hpad = h3 - h;
|
|
}
|
|
else
|
|
{
|
|
h3 = h;
|
|
w3 = in_small_pad.w / scale;
|
|
wpad = w3 - w;
|
|
hpad = 0;
|
|
}
|
|
|
|
ncnn::Mat in = ncnn::Mat::from_pixels(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, w, h);
|
|
|
|
ncnn::copy_make_border(in, in_pad, hpad / 2, hpad - hpad / 2, wpad / 2, wpad - wpad / 2, ncnn::BORDER_CONSTANT, 114.f);
|
|
|
|
in_pad.substract_mean_normalize(0, norm_vals);
|
|
}
|
|
else
|
|
{
|
|
ncnn::Mat in = ncnn::Mat::from_pixels(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, w, h);
|
|
|
|
// letterbox pad to target_size rectangle
|
|
wpad = (w + max_stride - 1) / max_stride * max_stride - w;
|
|
hpad = (h + max_stride - 1) / max_stride * max_stride - h;
|
|
ncnn::copy_make_border(in, in_pad, hpad / 2, hpad - hpad / 2, wpad / 2, wpad - wpad / 2, ncnn::BORDER_CONSTANT, 114.f);
|
|
|
|
in_pad.substract_mean_normalize(0, norm_vals);
|
|
|
|
in_small_pad = in_pad;
|
|
}
|
|
|
|
// rvm_mobilenetv3
|
|
ncnn::Mat r1(in_small_pad.w / 2, in_small_pad.h / 2, 16);
|
|
ncnn::Mat r2(in_small_pad.w / 4, in_small_pad.h / 4, 20);
|
|
ncnn::Mat r3(in_small_pad.w / 8, in_small_pad.h / 8, 40);
|
|
ncnn::Mat r4(in_small_pad.w / 16, in_small_pad.h / 16, 64);
|
|
|
|
// rvm_resnet50
|
|
// ncnn::Mat r1(in_small_pad.w / 2, in_small_pad.h / 2, 16);
|
|
// ncnn::Mat r2(in_small_pad.w / 4, in_small_pad.h / 4, 32);
|
|
// ncnn::Mat r3(in_small_pad.w / 8, in_small_pad.h / 8, 64);
|
|
// ncnn::Mat r4(in_small_pad.w / 16, in_small_pad.h / 16, 128);
|
|
|
|
r1.fill(0.f);
|
|
r2.fill(0.f);
|
|
r3.fill(0.f);
|
|
r4.fill(0.f);
|
|
|
|
ncnn::Extractor ex = rvm.create_extractor();
|
|
|
|
ex.input("in0", in_pad);
|
|
ex.input("in1", in_small_pad);
|
|
|
|
ex.input("in2", r1);
|
|
ex.input("in3", r2);
|
|
ex.input("in4", r3);
|
|
ex.input("in5", r4);
|
|
|
|
ncnn::Mat out_fgr;
|
|
ncnn::Mat out_pha;
|
|
|
|
if (downsample)
|
|
{
|
|
if (refine_deep)
|
|
{
|
|
// downsample + refine deep
|
|
ex.extract("out2", out_fgr);
|
|
ex.extract("out3", out_pha);
|
|
}
|
|
else // if (refine_fast)
|
|
{
|
|
// downsample + refine fast
|
|
ex.extract("out4", out_fgr);
|
|
ex.extract("out5", out_pha);
|
|
}
|
|
}
|
|
else
|
|
{
|
|
// no downsample
|
|
ex.extract("out0", out_fgr);
|
|
ex.extract("out1", out_pha);
|
|
}
|
|
|
|
ncnn::Mat out_seg;
|
|
|
|
// segmentation
|
|
ex.extract("out6", out_seg);
|
|
|
|
// feats
|
|
ex.extract("out7", r1);
|
|
ex.extract("out8", r2);
|
|
ex.extract("out9", r3);
|
|
ex.extract("out10", r4);
|
|
|
|
const float denorm_vals[3] = {255.f, 255.f, 255.f};
|
|
|
|
out_fgr.substract_mean_normalize(0, denorm_vals);
|
|
fgr.create(out_fgr.h, out_fgr.w, CV_8UC3);
|
|
out_fgr.to_pixels(fgr.data, ncnn::Mat::PIXEL_RGB2BGR);
|
|
|
|
out_pha.substract_mean_normalize(0, denorm_vals);
|
|
pha.create(out_pha.h, out_pha.w, CV_8UC1);
|
|
out_pha.to_pixels(pha.data, ncnn::Mat::PIXEL_GRAY);
|
|
|
|
out_seg.substract_mean_normalize(0, denorm_vals);
|
|
seg.create(in_pad.h, in_pad.w, CV_8UC1);
|
|
out_seg.to_pixels_resize(seg.data, ncnn::Mat::PIXEL_GRAY, in_pad.w, in_pad.h);
|
|
|
|
// cut letterbox pad
|
|
fgr = fgr(cv::Rect(wpad / 2, hpad / 2, w, h));
|
|
pha = pha(cv::Rect(wpad / 2, hpad / 2, w, h));
|
|
seg = seg(cv::Rect(wpad / 2, hpad / 2, w, h));
|
|
|
|
return 0;
|
|
}
|
|
|
|
static void draw_objects(const cv::Mat& bgr, const cv::Mat& fgr, const cv::Mat& pha, const cv::Mat& seg)
|
|
{
|
|
const int w = bgr.cols;
|
|
const int h = bgr.rows;
|
|
|
|
// composite
|
|
cv::Mat comp(h, w, CV_8UC3);
|
|
for (int y = 0; y < h; y++)
|
|
{
|
|
const uchar* pf = fgr.ptr<const uchar>(y);
|
|
const uchar* pa = pha.ptr<const uchar>(y);
|
|
uchar* p = comp.ptr<uchar>(y);
|
|
for (int x = 0; x < w; x++)
|
|
{
|
|
const float alpha = pa[0] / 255.f;
|
|
p[0] = cv::saturate_cast<uchar>(pf[0] * alpha + (1 - alpha) * 155);
|
|
p[1] = cv::saturate_cast<uchar>(pf[1] * alpha + (1 - alpha) * 255);
|
|
p[2] = cv::saturate_cast<uchar>(pf[2] * alpha + (1 - alpha) * 120);
|
|
pf += 3;
|
|
pa += 1;
|
|
p += 3;
|
|
}
|
|
}
|
|
|
|
// composite seg
|
|
cv::Mat comp_seg(h, w, CV_8UC3);
|
|
for (int y = 0; y < h; y++)
|
|
{
|
|
const uchar* pb = bgr.ptr<const uchar>(y);
|
|
const uchar* ps = seg.ptr<const uchar>(y);
|
|
uchar* p = comp_seg.ptr<uchar>(y);
|
|
for (int x = 0; x < w; x++)
|
|
{
|
|
const float alpha = ps[0] / 255.f;
|
|
p[0] = cv::saturate_cast<uchar>(pb[0] * alpha + (1 - alpha) * 155);
|
|
p[1] = cv::saturate_cast<uchar>(pb[1] * alpha + (1 - alpha) * 255);
|
|
p[2] = cv::saturate_cast<uchar>(pb[2] * alpha + (1 - alpha) * 120);
|
|
pb += 3;
|
|
ps += 1;
|
|
p += 3;
|
|
}
|
|
}
|
|
|
|
cv::imshow("comp", comp);
|
|
cv::imshow("comp_seg", comp_seg);
|
|
cv::waitKey(0);
|
|
}
|
|
|
|
int main(int argc, char** argv)
|
|
{
|
|
if (argc != 2)
|
|
{
|
|
fprintf(stderr, "Usage: %s [imagepath]\n", argv[0]);
|
|
return -1;
|
|
}
|
|
|
|
const char* imagepath = argv[1];
|
|
|
|
cv::Mat m = cv::imread(imagepath, 1);
|
|
if (m.empty())
|
|
{
|
|
fprintf(stderr, "cv::imread %s failed\n", imagepath);
|
|
return -1;
|
|
}
|
|
|
|
cv::Mat fgr;
|
|
cv::Mat pha;
|
|
cv::Mat seg;
|
|
detect_rvm(m, fgr, pha, seg);
|
|
|
|
draw_objects(m, fgr, pha, seg);
|
|
|
|
return 0;
|
|
}
|