Skip to main content

vk_video/vulkan_video/
wgpu_api.rs

1use crate::{
2    DecoderError, DecoderEvent, EncodedInputChunk, EncodedOutputChunk, InputFrame, OutputFrame,
3    VulkanEncoderError,
4    parser::{
5        decoder_instructions::compile_to_decoder_instructions,
6        h264::{AccessUnit, H264Parser},
7        reference_manager::ReferenceContext,
8    },
9    vulkan_decoder::{FrameSorter, VulkanDecoder},
10    vulkan_encoder::VulkanEncoder,
11};
12
13/// A decoder that outputs frames stored as [`wgpu::Texture`]s
14pub struct WgpuTexturesDecoder {
15    pub(crate) vulkan_decoder: VulkanDecoder<'static>,
16    pub(crate) parser: H264Parser,
17    pub(crate) reference_ctx: ReferenceContext,
18    pub(crate) frame_sorter: FrameSorter<wgpu::Texture>,
19}
20
21impl WgpuTexturesDecoder {
22    /// The produced textures have the [`wgpu::TextureFormat::NV12`] format and can be used as a texture binding.
23    pub fn decode(
24        &mut self,
25        frame: EncodedInputChunk<'_>,
26    ) -> Result<Vec<OutputFrame<wgpu::Texture>>, DecoderError> {
27        self.process_event(DecoderEvent::DecodeChunk(frame))
28    }
29
30    /// Flush all frames from the decoder.
31    ///
32    /// Make sure that this is done when you have the knowledge that no more frames will be coming
33    /// that need to be presented before the already decoded frames.
34    pub fn flush(&mut self) -> Result<Vec<OutputFrame<wgpu::Texture>>, DecoderError> {
35        self.process_event(DecoderEvent::Flush)
36    }
37
38    /// Process a [`DecoderEvent`]. For most use cases, using [`Self::decode`] and [`Self::flush`] is enough.
39    /// Use this only when you need more fine-grained control.
40    /// May return a sequence of decoded frames in the [NV12 format](https://en.wikipedia.org/wiki/YCbCr#4:2:0).
41    pub fn process_event(
42        &mut self,
43        event: DecoderEvent<'_>,
44    ) -> Result<Vec<OutputFrame<wgpu::Texture>>, DecoderError> {
45        match event {
46            DecoderEvent::DecodeChunk(chunk) => {
47                let nalus = self.parser.parse(chunk.data, chunk.pts)?;
48                self.decode_access_units(nalus)
49            }
50            DecoderEvent::SignalFrameEnd => {
51                let access_units = self.parser.flush()?;
52                self.decode_access_units(access_units)
53            }
54            DecoderEvent::SignalDataLoss => {
55                self.reference_ctx.mark_missed_frames();
56                Ok(Vec::new())
57            }
58            DecoderEvent::Flush => {
59                let access_units = self.parser.flush()?;
60                let mut frames = self.decode_access_units(access_units)?;
61                frames.append(&mut self.frame_sorter.flush());
62                Ok(frames)
63            }
64        }
65    }
66
67    fn decode_access_units(
68        &mut self,
69        access_units: Vec<AccessUnit>,
70    ) -> Result<Vec<OutputFrame<wgpu::Texture>>, DecoderError> {
71        let instructions = compile_to_decoder_instructions(&mut self.reference_ctx, access_units)?;
72        let unsorted_frames = self.vulkan_decoder.decode_to_wgpu_textures(&instructions)?;
73        let sorted_frames = self.frame_sorter.put_frames(unsorted_frames);
74        Ok(sorted_frames)
75    }
76}
77
78/// An encoder that takes input frames as [`wgpu::Texture`]s (in [`wgpu::TextureFormat::NV12`])
79pub struct WgpuTexturesEncoder {
80    pub(crate) vulkan_encoder: VulkanEncoder<'static>,
81}
82
83impl WgpuTexturesEncoder {
84    /// The result is a chunk of H264 bytecode.
85    ///
86    /// If the `force_keyframe` option is set to `true`, the encoder will encode this frame as a
87    /// [keyframe](https://en.wikipedia.org/wiki/Video_compression_picture_types#Intra-coded_(I)_frames/slices_(key_frames)).
88    /// Otherwise, the encoder will decide which frames should be coded this way.
89    pub fn encode(
90        &mut self,
91        frame: InputFrame<wgpu::Texture>,
92        force_keyframe: bool,
93    ) -> Result<EncodedOutputChunk<Vec<u8>>, VulkanEncoderError> {
94        self.vulkan_encoder.encode_texture(frame, force_keyframe)
95    }
96
97    /// Retrieve encoded SPS NAL units from the video session parameters, in Annex B.
98    ///
99    /// Useful when `inline_stream_params` is `false` and the parameters need to be
100    /// sent out-of-band (e.g. in RTMP or MP4 headers).
101    pub fn sps(&self) -> Result<Vec<u8>, VulkanEncoderError> {
102        self.vulkan_encoder.stream_parameters(true, false)
103    }
104
105    /// Retrieve encoded PPS NAL units from the video session parameters, in Annex B.
106    ///
107    /// Useful when `inline_stream_params` is `false` and the parameters need to be
108    /// sent out-of-band (e.g. in RTMP or MP4 headers).
109    pub fn pps(&self) -> Result<Vec<u8>, VulkanEncoderError> {
110        self.vulkan_encoder.stream_parameters(false, true)
111    }
112}
113
114#[derive(thiserror::Error, Debug)]
115pub enum WgpuInitError {
116    #[error("Wgpu instance error: {0}")]
117    WgpuInstanceError(#[from] wgpu::hal::InstanceError),
118
119    #[error("Wgpu device error: {0}")]
120    WgpuDeviceError(#[from] wgpu::hal::DeviceError),
121
122    #[error("Wgpu request device error: {0}")]
123    WgpuRequestDeviceError(#[from] wgpu::RequestDeviceError),
124
125    #[error("Cannot create a wgpu adapter")]
126    WgpuAdapterNotCreated,
127}