mirror of
https://github.com/k4yt3x/video2x.git
synced 2026-02-04 03:22:07 +08:00
feat(realcugan): add support for Real-CUGAN ncnn Vulkan (#1268)
Signed-off-by: k4yt3x <i@k4yt3x.com>
This commit is contained in:
K4YT3X
207
src/filter_realcugan.cpp
Normal file
207
src/filter_realcugan.cpp
Normal file
@@ -0,0 +1,207 @@
|
||||
#include "filter_realcugan.h"
|
||||
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <filesystem>
|
||||
|
||||
#include <spdlog/spdlog.h>
|
||||
|
||||
#include "conversions.h"
|
||||
#include "fsutils.h"
|
||||
#include "logger_manager.h"
|
||||
|
||||
namespace video2x {
|
||||
namespace processors {
|
||||
|
||||
FilterRealcugan::FilterRealcugan(
|
||||
int gpuid,
|
||||
bool tta_mode,
|
||||
int scaling_factor,
|
||||
int noise_level,
|
||||
int num_threads,
|
||||
int syncgap,
|
||||
const fsutils::StringType model_name
|
||||
)
|
||||
: realcugan_(nullptr),
|
||||
gpuid_(gpuid),
|
||||
tta_mode_(tta_mode),
|
||||
scaling_factor_(scaling_factor),
|
||||
noise_level_(noise_level),
|
||||
num_threads_(num_threads),
|
||||
syncgap_(syncgap),
|
||||
model_name_(std::move(model_name)) {}
|
||||
|
||||
FilterRealcugan::~FilterRealcugan() {
|
||||
if (realcugan_) {
|
||||
delete realcugan_;
|
||||
realcugan_ = nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
int FilterRealcugan::init(AVCodecContext *dec_ctx, AVCodecContext *enc_ctx, AVBufferRef *) {
|
||||
// Construct the model paths using std::filesystem
|
||||
std::filesystem::path model_param_path;
|
||||
std::filesystem::path model_bin_path;
|
||||
|
||||
fsutils::StringType model_base_name =
|
||||
STR("up") + fsutils::to_string_type(scaling_factor_) + STR("x-");
|
||||
|
||||
switch (noise_level_) {
|
||||
case -1:
|
||||
model_base_name += STR("conservative");
|
||||
break;
|
||||
case 0:
|
||||
model_base_name += STR("no-denoise");
|
||||
break;
|
||||
default:
|
||||
model_base_name += STR("denoise") + fsutils::to_string_type(noise_level_) + STR("x");
|
||||
break;
|
||||
}
|
||||
|
||||
fsutils::StringType param_file_name = model_base_name + STR(".param");
|
||||
fsutils::StringType bin_file_name = model_base_name + STR(".bin");
|
||||
|
||||
// Find the model paths by model name if provided
|
||||
model_param_path =
|
||||
std::filesystem::path(STR("models")) / STR("realcugan") / model_name_ / param_file_name;
|
||||
model_bin_path =
|
||||
std::filesystem::path(STR("models")) / STR("realcugan") / model_name_ / bin_file_name;
|
||||
|
||||
// Get the full paths using a function that possibly modifies or validates the path
|
||||
std::filesystem::path model_param_full_path = fsutils::find_resource_file(model_param_path);
|
||||
std::filesystem::path model_bin_full_path = fsutils::find_resource_file(model_bin_path);
|
||||
|
||||
// Check if the model files exist
|
||||
if (!std::filesystem::exists(model_param_full_path)) {
|
||||
logger()->error("RealCUGAN model param file not found: {}", model_param_path.u8string());
|
||||
return -1;
|
||||
}
|
||||
if (!std::filesystem::exists(model_bin_full_path)) {
|
||||
logger()->error("RealCUGAN model bin file not found: {}", model_bin_path.u8string());
|
||||
return -1;
|
||||
}
|
||||
|
||||
// Create a new RealCUGAN instance
|
||||
realcugan_ = new RealCUGAN(gpuid_, tta_mode_, num_threads_);
|
||||
|
||||
// Store the time bases
|
||||
in_time_base_ = dec_ctx->time_base;
|
||||
out_time_base_ = enc_ctx->time_base;
|
||||
out_pix_fmt_ = enc_ctx->pix_fmt;
|
||||
|
||||
// Load the model
|
||||
if (realcugan_->load(model_param_full_path, model_bin_full_path) != 0) {
|
||||
logger()->error("Failed to load RealCUGAN model");
|
||||
return -1;
|
||||
}
|
||||
|
||||
// Set syncgap to 0 for models-nose
|
||||
if (model_name_.find(STR("models-nose")) != fsutils::StringType::npos) {
|
||||
syncgap_ = 0;
|
||||
}
|
||||
|
||||
// Set realcugan parameters
|
||||
realcugan_->scale = scaling_factor_;
|
||||
realcugan_->noise = noise_level_;
|
||||
realcugan_->prepadding = 10;
|
||||
|
||||
// Set prepadding based on scaling factor
|
||||
if (scaling_factor_ == 2) {
|
||||
realcugan_->prepadding = 18;
|
||||
}
|
||||
if (scaling_factor_ == 3) {
|
||||
realcugan_->prepadding = 14;
|
||||
}
|
||||
if (scaling_factor_ == 4) {
|
||||
realcugan_->prepadding = 19;
|
||||
}
|
||||
|
||||
// Calculate tilesize based on GPU heap budget
|
||||
uint32_t heap_budget = ncnn::get_gpu_device(gpuid_)->get_heap_budget();
|
||||
if (scaling_factor_ == 2) {
|
||||
if (heap_budget > 1300) {
|
||||
realcugan_->tilesize = 400;
|
||||
} else if (heap_budget > 800) {
|
||||
realcugan_->tilesize = 300;
|
||||
} else if (heap_budget > 400) {
|
||||
realcugan_->tilesize = 200;
|
||||
} else if (heap_budget > 200) {
|
||||
realcugan_->tilesize = 100;
|
||||
} else {
|
||||
realcugan_->tilesize = 32;
|
||||
}
|
||||
}
|
||||
if (scaling_factor_ == 3) {
|
||||
if (heap_budget > 3300) {
|
||||
realcugan_->tilesize = 400;
|
||||
} else if (heap_budget > 1900) {
|
||||
realcugan_->tilesize = 300;
|
||||
} else if (heap_budget > 950) {
|
||||
realcugan_->tilesize = 200;
|
||||
} else if (heap_budget > 320) {
|
||||
realcugan_->tilesize = 100;
|
||||
} else {
|
||||
realcugan_->tilesize = 32;
|
||||
}
|
||||
}
|
||||
if (scaling_factor_ == 4) {
|
||||
if (heap_budget > 1690) {
|
||||
realcugan_->tilesize = 400;
|
||||
} else if (heap_budget > 980) {
|
||||
realcugan_->tilesize = 300;
|
||||
} else if (heap_budget > 530) {
|
||||
realcugan_->tilesize = 200;
|
||||
} else if (heap_budget > 240) {
|
||||
realcugan_->tilesize = 100;
|
||||
} else {
|
||||
realcugan_->tilesize = 32;
|
||||
}
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
int FilterRealcugan::filter(AVFrame *in_frame, AVFrame **out_frame) {
|
||||
int ret;
|
||||
|
||||
// Convert the input frame to RGB24
|
||||
ncnn::Mat in_mat = conversions::avframe_to_ncnn_mat(in_frame);
|
||||
if (in_mat.empty()) {
|
||||
logger()->error("Failed to convert AVFrame to ncnn::Mat");
|
||||
return -1;
|
||||
}
|
||||
|
||||
// Allocate space for output ncnn::Mat
|
||||
int output_width = in_mat.w * realcugan_->scale;
|
||||
int output_height = in_mat.h * realcugan_->scale;
|
||||
ncnn::Mat out_mat = ncnn::Mat(output_width, output_height, static_cast<size_t>(3), 3);
|
||||
|
||||
ret = realcugan_->process(in_mat, out_mat);
|
||||
if (ret != 0) {
|
||||
logger()->error("RealCUGAN processing failed");
|
||||
return ret;
|
||||
}
|
||||
|
||||
// Convert ncnn::Mat to AVFrame
|
||||
*out_frame = conversions::ncnn_mat_to_avframe(out_mat, out_pix_fmt_);
|
||||
|
||||
// Rescale PTS to encoder's time base
|
||||
(*out_frame)->pts = av_rescale_q(in_frame->pts, in_time_base_, out_time_base_);
|
||||
|
||||
// Return the processed frame to the caller
|
||||
return ret;
|
||||
}
|
||||
|
||||
void FilterRealcugan::get_output_dimensions(
|
||||
const ProcessorConfig &,
|
||||
int in_width,
|
||||
int in_height,
|
||||
int &out_width,
|
||||
int &out_height
|
||||
) const {
|
||||
out_width = in_width * scaling_factor_;
|
||||
out_height = in_height * scaling_factor_;
|
||||
}
|
||||
|
||||
} // namespace processors
|
||||
} // namespace video2x
|
||||
Reference in New Issue
Block a user