Skip to main content

less_avc_wrapper/
lib.rs

1// Copyright (C) The Strand-Braid Authors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4#![deny(unsafe_code)]
5use std::io::Write;
6
7use strand_dynamic_frame::DynamicFrame;
8
9use less_avc::ycbcr_image::*;
10
11/// An H.264 encoding error.
12#[derive(thiserror::Error, Debug)]
13pub enum Error {
14    #[error("LessAVC error: {source}")]
15    LessAvcError {
16        #[from]
17        source: less_avc::Error,
18    },
19    #[error("convert image error: {source}")]
20    ConvertImageError {
21        #[from]
22        source: convert_image::Error,
23    },
24    #[error("y4m writer error: {0}")]
25    Y4mError(#[from] y4m_writer::Error),
26}
27
28type Result<T> = std::result::Result<T, Error>;
29
30fn convert_dynamic_to_y4m(frame: &DynamicFrame) -> Result<y4m_writer::Y4MFrame> {
31    let out_colorspace = y4m::Colorspace::C420paldv;
32    let forced_block_size = Some(16);
33    let y4m = y4m_writer::encode_y4m_dynamic_frame(frame, out_colorspace, forced_block_size)?;
34    Ok(y4m)
35}
36
37fn gen_y4m_ref(y4m: &y4m_writer::Y4MFrame) -> Result<YCbCrImage<'_>> {
38    let width = y4m.width.try_into().unwrap();
39    let height = y4m.height.try_into().unwrap();
40
41    let y_plane = DataPlane {
42        data: y4m.y_plane_data(),
43        stride: y4m.y_stride(),
44        bit_depth: less_avc::BitDepth::Depth8,
45    };
46
47    let planes = if y4m.is_known_mono_only() {
48        Planes::Mono(y_plane)
49    } else {
50        let u_plane = DataPlane {
51            data: y4m.u_plane_data(),
52            stride: y4m.u_stride(),
53            bit_depth: less_avc::BitDepth::Depth8,
54        };
55        let v_plane = DataPlane {
56            data: y4m.v_plane_data(),
57            stride: y4m.v_stride(),
58            bit_depth: less_avc::BitDepth::Depth8,
59        };
60        Planes::YCbCr((y_plane, u_plane, v_plane))
61    };
62    let im = YCbCrImage {
63        planes,
64        width,
65        height,
66    };
67
68    Ok(im)
69}
70
71#[derive(Default)]
72pub struct WrappedLessEncoder {
73    inner: Option<less_avc::LessEncoder>,
74}
75
76impl WrappedLessEncoder {
77    pub fn encode_dynamic_to_nal_units(
78        &mut self,
79        frame: &strand_dynamic_frame::DynamicFrame,
80    ) -> Result<Vec<Vec<u8>>> {
81        let y4m = convert_dynamic_to_y4m(frame)?;
82        let y4m_ref = gen_y4m_ref(&y4m)?;
83
84        let (nals, encoder) = match self.inner.take() {
85            None => {
86                let (nal_units, encoder) = less_avc::LessEncoder::new(&y4m_ref)?;
87                let nals: Vec<Vec<u8>> = nal_units
88                    .into_iter()
89                    .map(|nal_unit| nal_unit.to_nal_unit())
90                    .collect();
91                (nals, encoder)
92            }
93            Some(mut encoder) => {
94                let nal_unit = encoder.encode(&y4m_ref)?;
95                (vec![nal_unit.to_nal_unit()], encoder)
96            }
97        };
98
99        self.inner = Some(encoder);
100
101        Ok(nals)
102    }
103}
104
105pub struct H264WriterWrapper<W> {
106    inner: less_avc::H264Writer<W>,
107}
108
109impl<W: Write> H264WriterWrapper<W> {
110    pub fn new(wtr: W) -> Result<Self> {
111        Ok(Self {
112            inner: less_avc::H264Writer::new(wtr)?,
113        })
114    }
115
116    pub fn write_dynamic(&mut self, frame: &DynamicFrame) -> Result<()> {
117        let y4m = convert_dynamic_to_y4m(frame)?;
118        let y4m_ref = gen_y4m_ref(&y4m)?;
119        self.inner.write(&y4m_ref)?;
120        Ok(())
121    }
122}