1use std::{
2 env,
3 path::PathBuf,
4 time::{Duration, Instant},
5};
6
7#[cfg(all(feature = "cuda", target_os = "linux"))]
8use lumen_ffmpeg::CudaDriver;
9use lumen_ffmpeg::{
10 DecodeMode, GpuBackend, GpuVideoFrame, InputContext, VideoDecoder, VideoDecoderConfig,
11};
12
13fn main() -> Result<(), Box<dyn std::error::Error>> {
14 let mut args = env::args().skip(1);
15 let path = args.next().map(PathBuf::from).ok_or(
16 "usage: decode_file <path> [max_frames] [cpu|cuda|metal|vulkan] [gpu|rgba|cuda-rgba]",
17 )?;
18 let next = args.next();
19 let (max_frames, mode_arg) = match next {
20 Some(value) => match value.parse::<usize>() {
21 Ok(max_frames) => (Some(max_frames), args.next()),
22 Err(_) => (None, Some(value)),
23 },
24 None => (None, None),
25 };
26 let mode = match mode_arg.as_deref() {
27 Some("cuda") => DecodeMode::Gpu(GpuBackend::Cuda),
28 Some("metal") => DecodeMode::Gpu(GpuBackend::Metal),
29 Some("vulkan") => DecodeMode::Gpu(GpuBackend::Vulkan),
30 Some("cpu") | None => DecodeMode::Cpu,
31 Some(other) => return Err(format!("unknown decode mode `{other}`").into()),
32 };
33 let receive = match args.next().as_deref() {
34 Some("gpu") => ReceiveMode::Gpu,
35 Some("cuda-rgba") => ReceiveMode::CudaRgba,
36 Some("rgba") | None => ReceiveMode::Rgba,
37 Some(other) => return Err(format!("unknown receive mode `{other}`").into()),
38 };
39
40 let started = Instant::now();
41 let mut input = InputContext::open(path.to_string_lossy().to_string())?;
42 let stream = input.best_video_stream()?;
43 let open_elapsed = started.elapsed();
44 let mut decoder = VideoDecoder::open(
45 &input,
46 VideoDecoderConfig {
47 stream_index: stream.stream_index,
48 mode,
49 },
50 )?;
51
52 let decode_started = Instant::now();
53 let mut cuda = CudaRgbaState::new(receive, stream.width, stream.height)?;
54 let mut frames = 0_usize;
55 let mut bytes = 0_usize;
56
57 'decode: while let Some(packet) = input.read_packet()? {
58 decoder.send_packet(&packet)?;
59 while let Some(frame_bytes) = receive_frame(&mut decoder, mode, receive, cuda.as_mut())? {
60 frames = frames.saturating_add(1);
61 bytes = bytes.saturating_add(frame_bytes);
62 if max_frames.is_some_and(|limit| frames >= limit) {
63 break 'decode;
64 }
65 }
66 }
67
68 if max_frames.is_none_or(|limit| frames < limit) {
69 decoder.send_eof()?;
70 while let Some(frame_bytes) = receive_frame(&mut decoder, mode, receive, cuda.as_mut())? {
71 frames = frames.saturating_add(1);
72 bytes = bytes.saturating_add(frame_bytes);
73 if max_frames.is_some_and(|limit| frames >= limit) {
74 break;
75 }
76 }
77 }
78
79 let decode_elapsed = decode_started.elapsed();
80 let total_elapsed = started.elapsed();
81 let usage = usage();
82 println!("path={}", path.display());
83 println!("codec={:?}", stream.codec);
84 println!("dimensions={}x{}", stream.width, stream.height);
85 println!("frames={frames}");
86 println!("mode={mode:?}");
87 println!("receive={receive:?}");
88 println!("decoded_frame_bytes={bytes}");
89 println!("open_ms={}", millis(open_elapsed));
90 println!("decode_ms={}", millis(decode_elapsed));
91 println!("total_ms={}", millis(total_elapsed));
92 println!(
93 "fps={:.2}",
94 frames as f64 / decode_elapsed.as_secs_f64().max(1e-9)
95 );
96 println!("user_cpu_ms={}", millis(usage.user));
97 println!("system_cpu_ms={}", millis(usage.system));
98 println!("max_rss_platform_units={}", usage.max_rss);
99 Ok(())
100}
101
102fn receive_frame(
103 decoder: &mut VideoDecoder,
104 mode: DecodeMode,
105 receive: ReceiveMode,
106 cuda: Option<&mut CudaRgbaState>,
107) -> lumen_ffmpeg::Result<Option<usize>> {
108 match receive {
109 ReceiveMode::Rgba => Ok(decoder.receive_rgba_frame()?.map(|frame| frame.data.len())),
110 ReceiveMode::CudaRgba => match decoder.receive_gpu_frame()? {
111 Some(frame) => {
112 let Some(cuda) = cuda else {
113 return Err(lumen_ffmpeg::FfmpegError::new(
114 "decode_file",
115 "cuda-rgba receive mode requires a CUDA conversion state",
116 ));
117 };
118 cuda.convert(&frame)?;
119 Ok(Some(frame.estimated_rgba_bytes() as usize))
120 }
121 None => Ok(None),
122 },
123 ReceiveMode::Gpu => match mode {
124 DecodeMode::Cpu => Ok(decoder.receive_cpu_frame()?.map(|frame| frame.data.len())),
125 DecodeMode::Gpu(_) => Ok(decoder.receive_gpu_frame()?.map(|_| 0)),
126 },
127 }
128}
129
130#[derive(Debug, Clone, Copy)]
131enum ReceiveMode {
132 Rgba,
133 Gpu,
134 CudaRgba,
135}
136
137#[cfg(all(feature = "cuda", target_os = "linux"))]
138struct CudaRgbaState {
139 converter: lumen_ffmpeg::CudaNv12ToRgbaConverter<'static>,
140 destination: lumen_ffmpeg::CudaDeviceAllocation<'static>,
141 _context: lumen_ffmpeg::CudaContext<'static>,
142}
143
144#[cfg(all(feature = "cuda", target_os = "linux"))]
145impl CudaRgbaState {
146 fn new(
147 receive: ReceiveMode,
148 width: u32,
149 height: u32,
150 ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
151 if !matches!(receive, ReceiveMode::CudaRgba) {
152 return Ok(None);
153 }
154 let driver = Box::leak(Box::new(CudaDriver::load()?));
155 let context = driver.create_primary_context()?;
156 let converter = driver.create_nv12_to_rgba_converter(&context)?;
157 let destination = driver.allocate_rgba_frame(width, height)?;
158 Ok(Some(Self {
159 converter,
160 destination,
161 _context: context,
162 }))
163 }
164
165 fn convert(&self, frame: &GpuVideoFrame) -> lumen_ffmpeg::Result<()> {
166 let GpuVideoFrame::Cuda(frame) = frame else {
167 return Err(lumen_ffmpeg::FfmpegError::new(
168 "decode_file",
169 "cuda-rgba receive mode requires CUDA decoded frames",
170 ));
171 };
172 self.converter
173 .convert(frame, &self.destination)
174 .map_err(|error| lumen_ffmpeg::FfmpegError::new("nv12_to_rgba8", error))
175 }
176}
177
178#[cfg(not(all(feature = "cuda", target_os = "linux")))]
179struct CudaRgbaState;
180
181#[cfg(not(all(feature = "cuda", target_os = "linux")))]
182impl CudaRgbaState {
183 fn new(
184 receive: ReceiveMode,
185 _width: u32,
186 _height: u32,
187 ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
188 if matches!(receive, ReceiveMode::CudaRgba) {
189 return Err(
190 "cuda-rgba receive mode requires a Linux build with the cuda feature".into(),
191 );
192 }
193 Ok(None)
194 }
195
196 fn convert(&self, _frame: &GpuVideoFrame) -> lumen_ffmpeg::Result<()> {
197 Err(lumen_ffmpeg::FfmpegError::new(
198 "decode_file",
199 "cuda-rgba receive mode requires a Linux build with the cuda feature",
200 ))
201 }
202}
203
204fn millis(duration: Duration) -> u128 {
205 duration.as_millis()
206}
207
208struct Usage {
209 user: Duration,
210 system: Duration,
211 max_rss: i64,
212}
213
214fn usage() -> Usage {
215 unsafe {
216 let mut value = std::mem::zeroed();
217 if libc::getrusage(libc::RUSAGE_SELF, &mut value) != 0 {
218 return Usage {
219 user: Duration::ZERO,
220 system: Duration::ZERO,
221 max_rss: 0,
222 };
223 }
224 Usage {
225 user: timeval_to_duration(value.ru_utime),
226 system: timeval_to_duration(value.ru_stime),
227 max_rss: value.ru_maxrss,
228 }
229 }
230}
231
232fn timeval_to_duration(value: libc::timeval) -> Duration {
233 Duration::new(
234 value.tv_sec.max(0) as u64,
235 (value.tv_usec.max(0) as u32) * 1_000,
236 )
237}