Skip to main content

viser_persegment/
lib.rs

1//! Segment-level CRF adaptation for the `viser` video-encoding-optimizer workspace.
2//!
3//! Analyzes spatial/temporal complexity across the source, then for each short (typically
4//! 1-second) segment runs a closed-loop binary search over CRF, encoding and measuring VMAF
5//! until the target quality is met within tolerance. The entry point `adapt` returns
6//! per-segment CRF choices along with overall bitrate/VMAF averages and the complexity profile.
7
8use serde::{Deserialize, Serialize};
9use std::sync::Arc;
10use std::time::{Duration, Instant};
11use viser_complexity::{self, AnalyzeOpts, Profile};
12use viser_ffmpeg::{self, Codec, EncodeJob, Resolution};
13use viser_quality::{self, MeasureOpts, Metric};
14
15/// Config for segment-level CRF adaptation.
16#[derive(Debug, Clone)]
17pub struct Config {
18    /// Target VMAF score each segment's CRF search aims to achieve.
19    pub target_vmaf: f64,
20    /// Acceptable absolute VMAF deviation from the target before stopping the search.
21    pub tolerance: f64,
22    /// Lowest (highest-quality) CRF the search will consider.
23    pub min_crf: i32,
24    /// Highest (lowest-quality) CRF the search will consider.
25    pub max_crf: i32,
26    /// Codec used for segment encodes.
27    pub codec: Codec,
28    /// Output resolution; `None` keeps the source resolution.
29    pub resolution: Option<Resolution>,
30    /// Encoder preset (e.g. "medium").
31    pub preset: String,
32    /// Length of each adapted segment.
33    pub segment_duration: Duration,
34    /// Maximum binary-search iterations per segment.
35    pub max_iterations: i32,
36}
37
38impl Default for Config {
39    fn default() -> Self {
40        Self {
41            target_vmaf: 93.0,
42            tolerance: 2.0,
43            min_crf: 15,
44            max_crf: 45,
45            codec: Codec::X264,
46            resolution: None,
47            preset: "medium".into(),
48            segment_duration: Duration::from_secs(1),
49            max_iterations: 5,
50        }
51    }
52}
53
54/// Adaptation result for a single segment.
55#[derive(Debug, Clone, Serialize, Deserialize)]
56pub struct SegmentResult {
57    /// Start timestamp of the segment within the source.
58    pub start: Duration,
59    /// End timestamp of the segment within the source.
60    pub end: Duration,
61    /// Final CRF selected for the segment.
62    pub crf: i32,
63    /// Measured bitrate at the selected CRF.
64    pub bitrate: f64,
65    /// Measured VMAF at the selected CRF.
66    pub vmaf: f64,
67    /// Complexity score of the segment from the complexity analysis.
68    pub complexity: f64,
69    /// Number of search iterations performed for the segment.
70    pub iterations: i32,
71}
72
73/// Complete output of a segment-level CRF adaptation.
74#[derive(Debug, Clone, Serialize, Deserialize)]
75pub struct Result {
76    /// Path to the analyzed source video.
77    pub source: String,
78    /// Per-segment adaptation results in time order.
79    pub segments: Vec<SegmentResult>,
80    /// Duration-weighted average bitrate across all segments.
81    pub avg_bitrate: f64,
82    /// Duration-weighted average VMAF across all segments.
83    pub avg_vmaf: f64,
84    /// The target VMAF the adaptation aimed for.
85    pub target_vmaf: f64,
86    /// Wall-clock duration of the adaptation.
87    pub duration: Duration,
88    /// Complexity profile computed for the source.
89    pub complexity_profile: Profile,
90}
91
92/// Runs segment-level adaptation.
93pub async fn adapt(source: &str, cfg: Config) -> anyhow::Result<Result> {
94    let start = Instant::now();
95
96    // Step 1: Analyze complexity
97    let profile = viser_complexity::analyze(
98        source,
99        AnalyzeOpts { segment_duration: cfg.segment_duration, subsample: 2 },
100    )
101    .await?;
102
103    // Step 2: Map complexity to initial CRF
104    let segments: Vec<SegmentResult> = profile
105        .segments
106        .iter()
107        .map(|seg| SegmentResult {
108            start: seg.start,
109            end: seg.end,
110            crf: complexity_to_crf(seg.score, cfg.min_crf, cfg.max_crf),
111            bitrate: 0.0,
112            vmaf: 0.0,
113            complexity: seg.score,
114            iterations: 0,
115        })
116        .collect();
117
118    // Step 3: Temp directory
119    let tmp_dir = tempfile::Builder::new().prefix("viser-persegment-").tempdir()?;
120
121    // Step 4: Encode and verify each segment in parallel
122    let parallel = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(4);
123    let semaphore = Arc::new(tokio::sync::Semaphore::new(parallel));
124    let source = Arc::new(source.to_string());
125
126    let mut handles = Vec::new();
127    for (i, seg) in segments.iter().enumerate() {
128        let sem = semaphore.clone();
129        let source = source.clone();
130        let seg = seg.clone();
131        let cfg_shared = Config {
132            min_crf: cfg.min_crf,
133            max_crf: cfg.max_crf,
134            codec: cfg.codec,
135            preset: cfg.preset.clone(),
136            resolution: cfg.resolution,
137            target_vmaf: cfg.target_vmaf,
138            tolerance: cfg.tolerance,
139            max_iterations: cfg.max_iterations,
140            segment_duration: Duration::from_secs(0),
141        };
142        let tmp_dir_path = tmp_dir.path().to_path_buf();
143
144        handles.push(tokio::spawn(async move {
145            let _permit = sem.acquire().await.unwrap();
146
147            let mut crf_low = cfg_shared.min_crf;
148            let mut crf_high = cfg_shared.max_crf;
149            let mut seg = seg;
150
151            for iter in 0..cfg_shared.max_iterations {
152                seg.iterations = iter + 1;
153
154                let seg_source = tmp_dir_path.join(format!("seg_{i:03}_src.mkv"));
155                let seg_encoded = tmp_dir_path.join(format!("seg_{i:03}_crf{}.mp4", seg.crf));
156
157                let dur_secs = (seg.end - seg.start).as_secs_f64();
158                viser_ffmpeg::extract(
159                    &source,
160                    &seg_source.to_string_lossy(),
161                    seg.start.as_secs_f64(),
162                    dur_secs,
163                )
164                .await?;
165
166                let job = EncodeJob {
167                    input: seg_source.to_string_lossy().to_string(),
168                    output: seg_encoded.to_string_lossy().to_string(),
169                    codec: cfg_shared.codec,
170                    crf: seg.crf,
171                    preset: cfg_shared.preset.clone(),
172                    resolution: cfg_shared.resolution,
173                    rate_control: viser_ffmpeg::RateControlMode::Crf,
174                    target_bitrate: 0.0,
175                    max_bitrate: 0.0,
176                    bufsize: 0.0,
177                    hwaccel: None,
178                    extra_args: vec![],
179                };
180
181                let enc_result = viser_ffmpeg::encode(job, None).await?;
182                seg.bitrate = enc_result.bitrate;
183
184                let q_result = viser_quality::measure(
185                    &seg_source.to_string_lossy(),
186                    &seg_encoded.to_string_lossy(),
187                    MeasureOpts { metrics: vec![Metric::Vmaf], subsample: 5, ..Default::default() },
188                )
189                .await?;
190                seg.vmaf = q_result.vmaf;
191
192                let _ = std::fs::remove_file(&seg_encoded);
193                let _ = std::fs::remove_file(&seg_source);
194
195                if (seg.vmaf - cfg_shared.target_vmaf).abs() <= cfg_shared.tolerance {
196                    break;
197                }
198
199                if seg.vmaf > cfg_shared.target_vmaf + cfg_shared.tolerance {
200                    crf_low = seg.crf;
201                } else {
202                    crf_high = seg.crf;
203                }
204                seg.crf = (crf_low + crf_high) / 2;
205
206                if crf_high - crf_low <= 1 {
207                    break;
208                }
209            }
210
211            Ok::<_, anyhow::Error>((i, seg))
212        }));
213    }
214
215    let mut seg_results: Vec<Option<SegmentResult>> = vec![None; segments.len()];
216    for h in handles {
217        let (i, seg) = h.await??;
218        seg_results[i] = Some(seg);
219    }
220    let segments: Vec<SegmentResult> = seg_results.into_iter().flatten().collect();
221
222    // Compute averages
223    let mut total_bitrate = 0.0;
224    let mut total_vmaf = 0.0;
225    let mut total_dur = 0.0;
226    for seg in &segments {
227        let dur = (seg.end - seg.start).as_secs_f64();
228        total_bitrate += seg.bitrate * dur;
229        total_vmaf += seg.vmaf * dur;
230        total_dur += dur;
231    }
232
233    // Guard against zero total duration (no segments, or all zero-length) so
234    // averages are 0.0 rather than NaN.
235    let (avg_bitrate, avg_vmaf) = if total_dur > 0.0 {
236        (total_bitrate / total_dur, total_vmaf / total_dur)
237    } else {
238        (0.0, 0.0)
239    };
240
241    Ok(Result {
242        source: source.to_string(),
243        segments,
244        avg_bitrate,
245        avg_vmaf,
246        target_vmaf: cfg.target_vmaf,
247        duration: start.elapsed(),
248        complexity_profile: profile,
249    })
250}
251
252fn complexity_to_crf(score: f64, min_crf: i32, max_crf: i32) -> i32 {
253    let crf = max_crf as f64 - (score / 100.0) * (max_crf - min_crf) as f64;
254    crf.round().clamp(min_crf as f64, max_crf as f64) as i32
255}