1use 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#[derive(Debug, Clone)]
17pub struct Config {
18 pub target_vmaf: f64,
20 pub tolerance: f64,
22 pub min_crf: i32,
24 pub max_crf: i32,
26 pub codec: Codec,
28 pub resolution: Option<Resolution>,
30 pub preset: String,
32 pub segment_duration: Duration,
34 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#[derive(Debug, Clone, Serialize, Deserialize)]
56pub struct SegmentResult {
57 pub start: Duration,
59 pub end: Duration,
61 pub crf: i32,
63 pub bitrate: f64,
65 pub vmaf: f64,
67 pub complexity: f64,
69 pub iterations: i32,
71}
72
73#[derive(Debug, Clone, Serialize, Deserialize)]
75pub struct Result {
76 pub source: String,
78 pub segments: Vec<SegmentResult>,
80 pub avg_bitrate: f64,
82 pub avg_vmaf: f64,
84 pub target_vmaf: f64,
86 pub duration: Duration,
88 pub complexity_profile: Profile,
90}
91
92pub async fn adapt(source: &str, cfg: Config) -> anyhow::Result<Result> {
94 let start = Instant::now();
95
96 let profile = viser_complexity::analyze(
98 source,
99 AnalyzeOpts { segment_duration: cfg.segment_duration, subsample: 2 },
100 )
101 .await?;
102
103 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 let tmp_dir = tempfile::Builder::new().prefix("viser-persegment-").tempdir()?;
120
121 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 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 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}