feat: add template json
This commit is contained in:
imeepos
423
cargos/tvai-v2/src/ffmpeg.rs
Normal file
423
cargos/tvai-v2/src/ffmpeg.rs
Normal file
@@ -0,0 +1,423 @@
|
||||
use crate::{Template, TvaiError, TvaiResult};
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// FFmpeg command generator for Topaz Video AI templates
|
||||
#[derive(Debug, Default)]
|
||||
pub struct FfmpegCommandGenerator {
|
||||
/// Model mappings from template models to FFmpeg filter models
|
||||
model_mappings: HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl FfmpegCommandGenerator {
|
||||
/// Create a new FFmpeg command generator
|
||||
pub fn new() -> Self {
|
||||
let mut generator = Self {
|
||||
model_mappings: HashMap::new(),
|
||||
};
|
||||
|
||||
// Initialize default model mappings
|
||||
generator.init_model_mappings();
|
||||
generator
|
||||
}
|
||||
|
||||
/// Initialize model mappings from template models to FFmpeg filter models
|
||||
fn init_model_mappings(&mut self) {
|
||||
// Enhancement models
|
||||
self.model_mappings.insert("prob-3".to_string(), "ahq-12".to_string());
|
||||
self.model_mappings.insert("prob-4".to_string(), "ahq-12".to_string());
|
||||
self.model_mappings.insert("nyx-3".to_string(), "nyx-3".to_string());
|
||||
|
||||
// Frame interpolation models
|
||||
self.model_mappings.insert("apo-8".to_string(), "chr-2".to_string());
|
||||
self.model_mappings.insert("chf-3".to_string(), "chr-2".to_string());
|
||||
|
||||
// Motion blur models
|
||||
self.model_mappings.insert("thm-2".to_string(), "thm-2".to_string());
|
||||
|
||||
// Stabilization models (mapped to ref-2 for tvai_stb)
|
||||
// Note: Template stabilization uses different approach than tvai_stb filter
|
||||
}
|
||||
|
||||
/// Add custom model mapping
|
||||
pub fn add_model_mapping(&mut self, template_model: String, ffmpeg_model: String) {
|
||||
self.model_mappings.insert(template_model, ffmpeg_model);
|
||||
}
|
||||
|
||||
/// Generate FFmpeg command from template
|
||||
pub fn generate(&self, template: &Template, input_file: &str, output_file: &str) -> TvaiResult<String> {
|
||||
let mut filters = Vec::new();
|
||||
let settings = &template.settings;
|
||||
|
||||
// Generate stabilization filter
|
||||
if settings.stabilize.active {
|
||||
let stb_filter = self.generate_stabilization_filter(&settings.stabilize)?;
|
||||
filters.push(stb_filter);
|
||||
}
|
||||
|
||||
// Generate enhancement filter
|
||||
if settings.enhance.active {
|
||||
let up_filter = self.generate_enhancement_filter(&settings.enhance)?;
|
||||
filters.push(up_filter);
|
||||
}
|
||||
|
||||
// Generate frame interpolation filter
|
||||
if settings.slow_motion.active {
|
||||
let fi_filter = self.generate_frame_interpolation_filter(&settings.slow_motion, &settings.output)?;
|
||||
filters.push(fi_filter);
|
||||
}
|
||||
|
||||
// Generate grain filter (using FFmpeg's noise filter as approximation)
|
||||
if settings.grain.active {
|
||||
let grain_filter = self.generate_grain_filter(&settings.grain)?;
|
||||
filters.push(grain_filter);
|
||||
}
|
||||
|
||||
// Build the complete FFmpeg command
|
||||
let mut command = format!("ffmpeg -i \"{}\"", input_file);
|
||||
|
||||
if !filters.is_empty() {
|
||||
command.push_str(" -vf \"");
|
||||
command.push_str(&filters.join(","));
|
||||
command.push('"');
|
||||
}
|
||||
|
||||
// Add output settings
|
||||
command.push_str(&self.generate_output_settings(&settings.output)?);
|
||||
|
||||
command.push_str(&format!(" \"{}\"", output_file));
|
||||
|
||||
Ok(command)
|
||||
}
|
||||
|
||||
/// Generate stabilization filter (tvai_stb)
|
||||
fn generate_stabilization_filter(&self, settings: &crate::StabilizeSettings) -> TvaiResult<String> {
|
||||
let mut params = Vec::new();
|
||||
|
||||
// Use default model for stabilization
|
||||
params.push("model=ref-2".to_string());
|
||||
|
||||
// Map smoothness (0-100 in template to 0-16 in filter)
|
||||
let smoothness = (settings.smooth as f64 / 100.0 * 16.0).round() as i32;
|
||||
params.push(format!("smoothness={}", smoothness));
|
||||
|
||||
// Map method (0: auto crop, 1: full frame)
|
||||
let full = if settings.method == 1 { 1 } else { 0 };
|
||||
params.push(format!("full={}", full));
|
||||
|
||||
// Rolling shutter correction
|
||||
if settings.rsc {
|
||||
params.push("roll=1".to_string());
|
||||
}
|
||||
|
||||
// Reduce motion
|
||||
if settings.reduce_motion {
|
||||
params.push(format!("reduce={}", settings.reduce_motion_iteration));
|
||||
}
|
||||
|
||||
Ok(format!("tvai_stb={}", params.join(":")))
|
||||
}
|
||||
|
||||
/// Generate enhancement filter (tvai_up)
|
||||
fn generate_enhancement_filter(&self, settings: &crate::EnhanceSettings) -> TvaiResult<String> {
|
||||
let mut params = Vec::new();
|
||||
|
||||
// Map model
|
||||
let model = self.model_mappings.get(&settings.model)
|
||||
.ok_or_else(|| TvaiError::ModelMappingError(format!("Unknown enhancement model: {}", settings.model)))?;
|
||||
params.push(format!("model={}", model));
|
||||
|
||||
// Map enhancement parameters to tvai_up parameters
|
||||
if settings.compress != 0 {
|
||||
let compression = settings.compress as f64 / 100.0;
|
||||
params.push(format!("compression={:.2}", compression));
|
||||
}
|
||||
|
||||
if settings.detail != 0 {
|
||||
let details = settings.detail as f64 / 100.0;
|
||||
params.push(format!("details={:.2}", details));
|
||||
}
|
||||
|
||||
if settings.denoise != 0 {
|
||||
let noise = settings.denoise as f64 / 100.0;
|
||||
params.push(format!("noise={:.2}", noise));
|
||||
}
|
||||
|
||||
if settings.dehalo != 0 {
|
||||
let halo = settings.dehalo as f64 / 100.0;
|
||||
params.push(format!("halo={:.2}", halo));
|
||||
}
|
||||
|
||||
if settings.deblur != 0 {
|
||||
let blur = settings.deblur as f64 / 100.0;
|
||||
params.push(format!("blur={:.2}", blur));
|
||||
}
|
||||
|
||||
if settings.sharpen != 0 {
|
||||
let preblur = -(settings.sharpen as f64 / 100.0);
|
||||
params.push(format!("preblur={:.2}", preblur));
|
||||
}
|
||||
|
||||
Ok(format!("tvai_up={}", params.join(":")))
|
||||
}
|
||||
|
||||
/// Generate frame interpolation filter (tvai_fi)
|
||||
fn generate_frame_interpolation_filter(&self, slowmo_settings: &crate::SlowMotionSettings, output_settings: &crate::OutputSettings) -> TvaiResult<String> {
|
||||
let mut params = Vec::new();
|
||||
|
||||
// Map model
|
||||
let model = self.model_mappings.get(&slowmo_settings.model)
|
||||
.ok_or_else(|| TvaiError::ModelMappingError(format!("Unknown frame interpolation model: {}", slowmo_settings.model)))?;
|
||||
params.push(format!("model={}", model));
|
||||
|
||||
// Set slowmo factor
|
||||
if slowmo_settings.factor != 1.0 {
|
||||
params.push(format!("slowmo={}", slowmo_settings.factor));
|
||||
}
|
||||
|
||||
// Set output FPS if specified
|
||||
if output_settings.out_fps > 0.0 {
|
||||
params.push(format!("fps={}", output_settings.out_fps));
|
||||
}
|
||||
|
||||
// Duplicate frame threshold
|
||||
if !slowmo_settings.duplicate {
|
||||
params.push("rdt=-0.01".to_string()); // Disable duplicate frame removal
|
||||
} else {
|
||||
let threshold = slowmo_settings.duplicate_threshold as f64 / 1000.0;
|
||||
params.push(format!("rdt={:.3}", threshold));
|
||||
}
|
||||
|
||||
Ok(format!("tvai_fi={}", params.join(":")))
|
||||
}
|
||||
|
||||
/// Generate grain filter (approximation using FFmpeg noise filter)
|
||||
fn generate_grain_filter(&self, settings: &crate::GrainSettings) -> TvaiResult<String> {
|
||||
let strength = settings.grain as f64 / 100.0 * 20.0; // Scale to reasonable noise level
|
||||
Ok(format!("noise=alls={}:allf=t", strength))
|
||||
}
|
||||
|
||||
/// Generate output settings
|
||||
fn generate_output_settings(&self, settings: &crate::OutputSettings) -> TvaiResult<String> {
|
||||
let mut output_params = Vec::new();
|
||||
|
||||
// Handle output size method
|
||||
match settings.out_size_method {
|
||||
7 => {
|
||||
// 4K output
|
||||
output_params.push("-s 3840x2160".to_string());
|
||||
},
|
||||
6 => {
|
||||
// FHD output
|
||||
output_params.push("-s 1920x1080".to_string());
|
||||
},
|
||||
5 => {
|
||||
// HD output
|
||||
output_params.push("-s 1280x720".to_string());
|
||||
},
|
||||
_ => {
|
||||
// Keep original size (method 0) or other methods
|
||||
}
|
||||
}
|
||||
|
||||
// Handle frame rate
|
||||
if settings.out_fps > 0.0 {
|
||||
output_params.push(format!("-r {}", settings.out_fps));
|
||||
}
|
||||
|
||||
// Add default encoding settings for quality
|
||||
output_params.push("-c:v libx264".to_string());
|
||||
output_params.push("-crf 18".to_string());
|
||||
output_params.push("-preset slow".to_string());
|
||||
|
||||
if output_params.is_empty() {
|
||||
Ok(String::new())
|
||||
} else {
|
||||
Ok(format!(" {}", output_params.join(" ")))
|
||||
}
|
||||
}
|
||||
|
||||
/// Get available model mappings
|
||||
pub fn get_model_mappings(&self) -> &HashMap<String, String> {
|
||||
&self.model_mappings
|
||||
}
|
||||
|
||||
/// Validate template for FFmpeg generation
|
||||
pub fn validate_template(&self, template: &Template) -> TvaiResult<()> {
|
||||
let settings = &template.settings;
|
||||
|
||||
// Check if any processing is enabled
|
||||
if !settings.enhance.active && !settings.slow_motion.active && !settings.stabilize.active && !settings.grain.active {
|
||||
return Err(TvaiError::ValidationError("No processing modules are active".to_string()));
|
||||
}
|
||||
|
||||
// Validate enhancement model if active
|
||||
if settings.enhance.active && !self.model_mappings.contains_key(&settings.enhance.model) {
|
||||
return Err(TvaiError::ModelMappingError(format!("Unknown enhancement model: {}", settings.enhance.model)));
|
||||
}
|
||||
|
||||
// Validate frame interpolation model if active
|
||||
if settings.slow_motion.active && !self.model_mappings.contains_key(&settings.slow_motion.model) {
|
||||
return Err(TvaiError::ModelMappingError(format!("Unknown frame interpolation model: {}", settings.slow_motion.model)));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Generate advanced FFmpeg command with hardware acceleration
|
||||
pub fn generate_with_hardware_acceleration(&self, template: &Template, input_file: &str, output_file: &str, gpu_device: Option<&str>) -> TvaiResult<String> {
|
||||
let mut command = self.generate(template, input_file, output_file)?;
|
||||
|
||||
// Add GPU device specification for TVAI filters
|
||||
if let Some(device) = gpu_device {
|
||||
command = command.replace("tvai_up=", &format!("tvai_up=device={}:", device));
|
||||
command = command.replace("tvai_fi=", &format!("tvai_fi=device={}:", device));
|
||||
command = command.replace("tvai_stb=", &format!("tvai_stb=device={}:", device));
|
||||
}
|
||||
|
||||
Ok(command)
|
||||
}
|
||||
|
||||
/// Generate command with custom output codec
|
||||
pub fn generate_with_codec(&self, template: &Template, input_file: &str, output_file: &str, codec: &str, quality: Option<i32>) -> TvaiResult<String> {
|
||||
let mut command = self.generate(template, input_file, output_file)?;
|
||||
|
||||
// Replace default codec settings
|
||||
command = command.replace("-c:v libx264", &format!("-c:v {}", codec));
|
||||
|
||||
if let Some(q) = quality {
|
||||
match codec {
|
||||
"hevc_nvenc" | "h264_nvenc" => {
|
||||
command = command.replace("-crf 18", &format!("-cq {}", q));
|
||||
},
|
||||
"h264_amf" | "hevc_amf" => {
|
||||
command = command.replace("-crf 18", &format!("-qp {}", q));
|
||||
},
|
||||
"h264_qsv" | "hevc_qsv" => {
|
||||
command = command.replace("-crf 18", &format!("-global_quality {}", q));
|
||||
},
|
||||
_ => {
|
||||
command = command.replace("-crf 18", &format!("-crf {}", q));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(command)
|
||||
}
|
||||
|
||||
/// Generate batch processing commands
|
||||
pub fn generate_batch_commands(&self, template: &Template, input_files: &[String], output_dir: &str) -> TvaiResult<Vec<String>> {
|
||||
let mut commands = Vec::new();
|
||||
|
||||
for input_file in input_files {
|
||||
let input_path = std::path::Path::new(input_file);
|
||||
let filename = input_path.file_stem()
|
||||
.ok_or_else(|| TvaiError::FfmpegError("Invalid input filename".to_string()))?
|
||||
.to_string_lossy();
|
||||
|
||||
let output_file = format!("{}/{}_processed.mp4", output_dir, filename);
|
||||
let command = self.generate(template, input_file, &output_file)?;
|
||||
commands.push(command);
|
||||
}
|
||||
|
||||
Ok(commands)
|
||||
}
|
||||
}
|
||||
|
||||
/// FFmpeg command builder for more complex scenarios
|
||||
#[derive(Debug, Default)]
|
||||
pub struct FfmpegCommandBuilder {
|
||||
input_file: String,
|
||||
output_file: String,
|
||||
filters: Vec<String>,
|
||||
codec: String,
|
||||
quality: Option<i32>,
|
||||
hardware_accel: Option<String>,
|
||||
custom_params: Vec<String>,
|
||||
}
|
||||
|
||||
impl FfmpegCommandBuilder {
|
||||
/// Create new command builder
|
||||
pub fn new(input_file: &str, output_file: &str) -> Self {
|
||||
Self {
|
||||
input_file: input_file.to_string(),
|
||||
output_file: output_file.to_string(),
|
||||
filters: Vec::new(),
|
||||
codec: "libx264".to_string(),
|
||||
quality: Some(18),
|
||||
hardware_accel: None,
|
||||
custom_params: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Add filter
|
||||
pub fn add_filter(mut self, filter: &str) -> Self {
|
||||
self.filters.push(filter.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
/// Set codec
|
||||
pub fn codec(mut self, codec: &str) -> Self {
|
||||
self.codec = codec.to_string();
|
||||
self
|
||||
}
|
||||
|
||||
/// Set quality
|
||||
pub fn quality(mut self, quality: i32) -> Self {
|
||||
self.quality = Some(quality);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set hardware acceleration
|
||||
pub fn hardware_accel(mut self, accel: &str) -> Self {
|
||||
self.hardware_accel = Some(accel.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
/// Add custom parameter
|
||||
pub fn custom_param(mut self, param: &str) -> Self {
|
||||
self.custom_params.push(param.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
/// Build the command
|
||||
pub fn build(self) -> String {
|
||||
let mut command = format!("ffmpeg -i \"{}\"", self.input_file);
|
||||
|
||||
if let Some(accel) = &self.hardware_accel {
|
||||
command.push_str(&format!(" -hwaccel {}", accel));
|
||||
}
|
||||
|
||||
if !self.filters.is_empty() {
|
||||
command.push_str(" -vf \"");
|
||||
command.push_str(&self.filters.join(","));
|
||||
command.push('"');
|
||||
}
|
||||
|
||||
command.push_str(&format!(" -c:v {}", self.codec));
|
||||
|
||||
if let Some(quality) = self.quality {
|
||||
match self.codec.as_str() {
|
||||
"hevc_nvenc" | "h264_nvenc" => {
|
||||
command.push_str(&format!(" -cq {}", quality));
|
||||
},
|
||||
"h264_amf" | "hevc_amf" => {
|
||||
command.push_str(&format!(" -qp {}", quality));
|
||||
},
|
||||
"h264_qsv" | "hevc_qsv" => {
|
||||
command.push_str(&format!(" -global_quality {}", quality));
|
||||
},
|
||||
_ => {
|
||||
command.push_str(&format!(" -crf {}", quality));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for param in &self.custom_params {
|
||||
command.push_str(&format!(" {}", param));
|
||||
}
|
||||
|
||||
command.push_str(&format!(" \"{}\"", self.output_file));
|
||||
|
||||
command
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user