nvidia-video-codec-rs/examples/demo.rs

use anyhow::bail;
use nvidia_video_codec::cuda::context::CuContext;
use nvidia_video_codec::cuda::device::CuDevice;
use nvidia_video_codec::decoder::CuvidRect;
use nvidia_video_codec::parser::CuvidVideoParser;
use nvidia_video_codec_sys::cuvid::{
    cuDriverGetVersion, cuInit, cudaVideoCodec_enum_cudaVideoCodec_H264,
};
use std::fs::File;
use std::io::BufReader;

fn init_and_version() -> CuContext {
    let ret = unsafe { cuInit(0) };
    println!("{:?}", ret);

    let ver = unsafe {
        let mut ver = 0;
        cuDriverGetVersion(&mut ver as *mut i32);
        ver
    };

    println!("Cuda Version: {}", ver);

    let dev = CuDevice::new(0).unwrap();
    println!(
        "Using: {} {:.2}Gb",
        dev.get_name().unwrap(),
        dev.get_total_mem().unwrap() as f64 / (1024 * 1024 * 1024) as f64
    );

    CuContext::new(dev, 0).unwrap()
}

fn main() -> mp4::Result<()> {
    let ctx = init_and_version();
    let conf = viuer::Config {
        // set offset
        x: 0,
        y: 0,
        // set dimensions
        width: Some(200),
        height: Some(40),
        ..Default::default()
    };

    let codec = nvidia_video_codec::decoder::CuvidCodec::H264;
    let config = nvidia_video_codec::parser::Config {
        crop: Some(CuvidRect {
            left: 280,
            top: 0,
            right: 1000,
            bottom: 720,
        }),
        // resize: Some((224, 224)),
        ..Default::default()
    };
    let mut parser = CuvidVideoParser::new(codec, config, move |res| match res {
        Ok(frame) => {
            println!("decoded {}x{}", frame.width, frame.height);
            if let Some(buf) = image::RgbImage::from_vec(frame.width, frame.height, frame.data) {
                let img = image::DynamicImage::ImageRgb8(buf);
                if let Err(err) = viuer::print(&img, &conf) {
                    println!("print image error: {}", err);
                }
            }

            true
        }
        Err(err) => {
            println!("err {}", err);
            true
        }
    })
    .unwrap();

    let f = File::open("/data/home/andrey/video_test/test_h264.mp4").unwrap();
    let size = f.metadata()?.len();
    let reader = BufReader::new(f);

    let mut mp4 = mp4::Mp4Reader::read_header(reader, size)?;

    // Print boxes.
    println!("major brand: {}", mp4.ftyp.major_brand);
    println!("timescale: {}", mp4.moov.mvhd.timescale);

    // Use available methods.
    println!("size: {}", mp4.size());

    let mut compatible_brands = String::new();
    for brand in mp4.compatible_brands().iter() {
        compatible_brands.push_str(&brand.to_string());
        compatible_brands.push_str(",");
    }
    println!("compatible brands: {}", compatible_brands);
    println!("duration: {:?}", mp4.duration());

    // Track info.
    let mut track_id = 0;
    let mut sample_count = 0;

    for track in mp4.tracks().values() {
        track_id = track.track_id();
        sample_count = track.sample_count();
        println!(
            "track: #{}({}) {} : {}; {}x{}; {:0.2}MBps; {}fps; {} samples; sps: {}, pps: {}; audio: {:?}",
            track.track_id(),
            track.language(),
            track.track_type()?,
            track.box_type()?,
            track.width(),
            track.height(),
            track.bitrate() as f64 / (1024 * 1024) as f64,
            track.frame_rate(),
            track.sample_count(),
            track.sequence_parameter_set().is_ok(),
            track.picture_parameter_set().is_ok(),
            track.audio_profile(),
        );
    }

    if let Some(track) = mp4.tracks().get(&track_id) {
        if let Ok(sps) = track.sequence_parameter_set() {
            let mut x: Vec<u8> = Vec::with_capacity(sps.len() + 4);
            x.extend_from_slice(&[0, 0, 0, 1]);
            x.extend_from_slice(sps);
            println!("sps: {:?}", x);

            parser.feed_data(&x, 0, false).unwrap();
        }

        if let Ok(pps) = track.picture_parameter_set() {
            let mut x: Vec<u8> = Vec::with_capacity(pps.len() + 4);
            x.extend_from_slice(&[0, 0, 0, 1]);
            x.extend_from_slice(pps);
            println!("pps: {:?}", x);

            parser.feed_data(&x, 0, false).unwrap();
        }
    }

    for sample_idx in 0..sample_count {
        let sample_id = sample_idx + 1;
        let sample = mp4.read_sample(track_id, sample_id).unwrap();

        if let Some(samp) = sample {
            let mut data: Vec<u8> = samp.bytes.into();
            convert_h264(&mut data).unwrap();

            parser
                .feed_data(
                    &data,
                    (((samp.start_time as f64 + samp.rendering_offset as f64) / 24000.0) * 1000.0)
                        as i64,
                    sample_idx == sample_count - 1,
                )
                .unwrap();
        }
    }

    Ok(())
}

fn convert_h264(data: &mut [u8]) -> anyhow::Result<()> {
    // TODO:
    // * For each IDR frame, copy the SPS and PPS from the stream's
    //   parameters, rather than depend on it being present in the frame
    //   already. In-band parameters aren't guaranteed. This is awkward
    //   with h264_reader v0.5's h264_reader::avcc::AvcDecoderRecord because it
    //   strips off the NAL header byte from each parameter. The next major
    //   version shouldn't do this.
    // * Copy only the slice data. In particular, don't copy SEI, which confuses
    //   Safari: <https://github.com/scottlamb/retina/issues/60#issuecomment-1178369955>

    let mut i = 0;
    while i < data.len() - 3 {
        // Replace each NAL's length with the Annex B start code b"\x00\x00\x00\x01".
        let bytes = &mut data[i..i + 4];
        let nalu_length = u32::from_be_bytes(bytes.try_into().unwrap()) as usize;
        bytes.copy_from_slice(&[0, 0, 0, 1]);

        i += 4 + nalu_length;

        if i > data.len() {
            bail!("partial nal body");
        }
    }

    if i < data.len() {
        bail!("partial nal body");
    }

    Ok(())
}