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

use anyhow::bail;
use mp4::TrackType;
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).unwrap();

    let f = File::open("/home/andrey/demo/h264/terminator.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 = None;
    let mut sample_count = 0;

    for track in mp4.tracks().values() {
        if track.track_type().unwrap() == TrackType::Video {
            track_id = Some(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_id) = track_id {
        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();

                let frame = 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();

                if let Some(frame) = 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);
                        }
                    }
                }
            }
        }
    } else {
        println!("No video track found!");
    }

    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(())
}