Skip to main content

decode_file/
decode_file.rs

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}