corevm_codec/
video.rs

1//! Video codec.
2
3use crate::{errors, rans, timer_finish, timer_start, Input, Output};
4use alloc::{vec, vec::Vec};
5use core::num::NonZero;
6use jam_codec::{Compact, Decode, Encode};
7
8mod delta;
9mod haar;
10mod quant;
11mod signmag;
12mod stats;
13mod yuv;
14
15pub use self::stats::*;
16use self::{quant::*, yuv::*};
17
18/// Video encoder configuration.
19#[derive(Debug)]
20#[non_exhaustive]
21pub struct Config {
22    /// Quantization level.
23    ///
24    /// Typical value is 3-4, maximum is 14.
25    /// Low values might increase the input size instead of decreasing.
26    pub quantization_level: u8,
27    /// Enable 4:2:0 chroma subsampling.
28    pub chroma_subsampling: bool,
29    /// Enable raw mode.
30    ///
31    /// In this mode no encoding is done, input frames are simply copied into
32    /// the output instead.
33    ///
34    /// This mode is useful running the code under a RISCV interpreter.
35    ///
36    /// Other configuration options have no effect when raw mode is enabled.
37    pub raw: bool,
38}
39
40impl Config {
41    /// Returns default lossless configuration.
42    pub fn default_lossless() -> Self {
43        Self {
44            quantization_level: 0,
45            chroma_subsampling: false,
46            raw: false,
47        }
48    }
49
50    /// Returns default raw mode configuration.
51    pub fn default_raw() -> Self {
52        Self {
53            quantization_level: 0,
54            chroma_subsampling: false,
55            raw: true,
56        }
57    }
58}
59
60impl Default for Config {
61    fn default() -> Self {
62        Self {
63            quantization_level: 4,
64            chroma_subsampling: true,
65            raw: false,
66        }
67    }
68}
69
70#[derive(Debug, Encode, Decode)]
71enum RawVideoFrameFormat {
72    Rgb888 = 0,
73    Rgb888Indexed8 = 1,
74}
75
76/// Video encoder.
77///
78/// This encoder uses semi-standard pipeline with the simplest and fastest
79/// implementation for each stage. Currently the following stages are
80/// implemented.
81///
82/// | Stage | How it is implemented |
83/// |-------|-----------------------|
84/// | RGB to YUV conversion | PGF (no division). |
85/// | Transform | Haar transform (no division). |
86/// | Quantization | Scalar quantization without a deadzone. |
87/// | Motion estimation | Simple delta encoding. |
88/// | Bit encoder | Sign-magniude conversion and rANS encoder. |
89///
90/// # Example
91///
92/// ```
93/// use core::num::NonZero;
94/// use corevm_codec::video;
95///
96/// let rgb_frame = [
97///     0,   0, 0,
98///     100, 0, 0,
99///     0, 100, 0,
100///     0, 0, 100,
101/// ];
102/// let width = NonZero::new(2).unwrap();
103/// let height = width;
104/// let mut output = Vec::new();
105/// let config = video::Config::default();
106/// let mut encoder = video::Encoder::new(width, height, config);
107/// encoder.start(&mut output);
108/// encoder.write_rgb888_frame(&rgb_frame, &mut output);
109/// encoder.finish(&mut output);
110/// ```
111pub struct Encoder {
112    width: NonZero<u16>,
113    height: NonZero<u16>,
114    quant: u8,
115    raw: bool,
116    // Current frame.
117    frame: YuvFrame,
118    // Previous frame.
119    prev_frame: YuvFrame,
120    // Output buffer.
121    buf: Vec<u8>,
122}
123
124impl Encoder {
125    /// Creates new encoder with the provided width, height, and configuration.
126    pub fn new(width: NonZero<u16>, height: NonZero<u16>, config: Config) -> Self {
127        let quant = config.quantization_level.min(MAX_QUANTIZATION_LEVEL);
128        let frame = YuvFrame::new(width, height, config.chroma_subsampling);
129        let prev_frame = frame.clone();
130        let buf = Vec::with_capacity(32 * 1024);
131        let raw = config.raw;
132        Self {
133            width,
134            height,
135            quant,
136            frame,
137            prev_frame,
138            buf,
139            raw,
140        }
141    }
142
143    /// Encode RGB888 frame and append the resulting bytes to the output.
144    pub fn write_rgb888_frame(&mut self, rgb_frame: &[u8], output: &mut impl Output) -> Stats {
145        assert_eq!(
146            usize::from(self.width.get()) * usize::from(self.height.get()) * 3,
147            rgb_frame.len()
148        );
149        if self.raw {
150            RawVideoFrameFormat::Rgb888.encode_to(output);
151            output.write(rgb_frame);
152            return Stats::default();
153        }
154        match self.frame {
155            YuvFrame::Yuv420p(ref mut frame) => {
156                let (y, u, v) = frame.as_mut_slices();
157                rgb888_to_yuv420p(rgb_frame, self.width, y, u, v);
158            }
159            YuvFrame::Yuv444p(ref mut frame) => {
160                let (y, u, v) = frame.as_mut_slices();
161                rgb888_to_yuv444p(rgb_frame, self.width, y, u, v);
162            }
163        }
164        self.write_frame(output)
165    }
166
167    /// Encode indexed RGB888 frame and append the resulting bytes to the
168    /// output.
169    ///
170    /// The frame consists of a 256-color palette and an array of 8-bit indices
171    /// into the palette. Each palette element is encoded as RGB888.
172    pub fn write_rgb888_indexed8_frame(
173        &mut self,
174        indexed_rgb_frame: &[u8],
175        output: &mut impl Output,
176    ) -> Stats {
177        assert_eq!(
178            usize::from(self.width.get()) * usize::from(self.height.get()) + 3 * 256,
179            indexed_rgb_frame.len()
180        );
181        if self.raw {
182            RawVideoFrameFormat::Rgb888Indexed8.encode_to(output);
183            output.write(indexed_rgb_frame);
184            return Stats::default();
185        }
186        match self.frame {
187            YuvFrame::Yuv420p(ref mut frame) => {
188                let (y, u, v) = frame.as_mut_slices();
189                rgb888_indexed8_to_yuv420p(indexed_rgb_frame, self.width, y, u, v);
190            }
191            YuvFrame::Yuv444p(ref mut frame) => {
192                let (y, u, v) = frame.as_mut_slices();
193                rgb888_indexed8_to_yuv444p(indexed_rgb_frame, self.width, y, u, v);
194            }
195        }
196        self.write_frame(output)
197    }
198
199    fn write_frame(&mut self, output: &mut impl Output) -> Stats {
200        let (uv_width, uv_height) = (self.frame.uv_width(), self.frame.uv_height());
201        let (y, u, v) = self.frame.as_mut_slices();
202        let t_transform = timer_start!();
203        haar::forward(y, self.width, self.height, self.quant);
204        haar::forward(u, uv_width, uv_height, self.quant);
205        haar::forward(v, uv_width, uv_height, self.quant);
206        timer_finish!(t_transform);
207        let t_delta = timer_start!();
208        let (y_prev, u_prev, v_prev) = self.prev_frame.as_mut_slices();
209        delta::encode(y, y_prev);
210        delta::encode(u, u_prev);
211        delta::encode(v, v_prev);
212        timer_finish!(t_delta);
213        let t_signmag = timer_start!();
214        self.buf.clear();
215        let y_stats = signmag::encode(y, &mut self.buf);
216        let u_stats = signmag::encode(u, &mut self.buf);
217        let v_stats = signmag::encode(v, &mut self.buf);
218        // Reverse to be able to read frames in normal order.
219        self.buf.reverse();
220        timer_finish!(t_signmag);
221        let stats = Stats {
222            y: y_stats,
223            u: u_stats,
224            v: v_stats,
225            #[cfg(all(feature = "stats", feature = "std"))]
226            time: TimeStats {
227                delta: t_delta.into_duration(),
228                transform: t_transform.into_duration(),
229                signmag: t_signmag.into_duration(),
230                count: 1,
231            },
232            count: 1,
233        };
234        output.write(&self.buf[..]);
235        stats
236    }
237
238    /// Write stream header.
239    pub fn start(&mut self, output: &mut impl Output) {
240        Compact(self.width.get()).encode_to(output);
241        Compact(self.height.get()).encode_to(output);
242        let chroma_subsampling = match self.frame {
243            YuvFrame::Yuv420p(..) => 1_u8,
244            YuvFrame::Yuv444p(..) => 0_u8,
245        };
246        let raw = match self.raw {
247            true => 1_u8,
248            false => 0_u8,
249        };
250        let config = self.quant
251            | (chroma_subsampling << QUANTIZATION_LEVEL_BITS)
252            | (raw << (QUANTIZATION_LEVEL_BITS + 1));
253        config.encode_to(output);
254    }
255
256    /// Write stream footer.
257    ///
258    /// This is currently empty but might change in the future.
259    pub fn finish(self, _output: &mut impl Output) {}
260}
261
262errors! {
263    (InvalidVideoStream "Invalid video stream")
264}
265
266#[doc(hidden)]
267impl From<rans::InvalidRansStream> for InvalidVideoStream {
268    fn from(_: rans::InvalidRansStream) -> Self {
269        InvalidVideoStream
270    }
271}
272
273#[doc(hidden)]
274impl From<signmag::InvalidSignMagStream> for InvalidVideoStream {
275    fn from(_: signmag::InvalidSignMagStream) -> Self {
276        InvalidVideoStream
277    }
278}
279
280#[doc(hidden)]
281impl From<jam_codec::Error> for InvalidVideoStream {
282    fn from(_: jam_codec::Error) -> Self {
283        InvalidVideoStream
284    }
285}
286
287/// Video decoder.
288pub struct Decoder {
289    width: NonZero<u16>,
290    height: NonZero<u16>,
291    quant: u8,
292    raw: bool,
293    frame: YuvFrame,
294    prev_frame: YuvFrame,
295}
296
297impl Decoder {
298    /// Create new decoder from the provided input.
299    pub fn new(input: &mut impl Input) -> Result<Self, InvalidVideoStream> {
300        let Compact(width) = Compact::<u16>::decode(input).map_err(|_| InvalidVideoStream)?;
301        let Compact(height) = Compact::<u16>::decode(input).map_err(|_| InvalidVideoStream)?;
302        let config = u8::decode(input).map_err(|_| InvalidVideoStream)?;
303        let quant = config & 0b1111;
304        if quant > MAX_QUANTIZATION_LEVEL {
305            return Err(InvalidVideoStream);
306        }
307        let chroma_subsampling = match (config >> QUANTIZATION_LEVEL_BITS) & 1 {
308            0 => false,
309            1 => true,
310            _ => return Err(InvalidVideoStream),
311        };
312        let raw = match (config >> (QUANTIZATION_LEVEL_BITS + 1)) & 1 {
313            0 => false,
314            1 => true,
315            _ => return Err(InvalidVideoStream),
316        };
317        let width = NonZero::new(width).ok_or(InvalidVideoStream)?;
318        let height = NonZero::new(height).ok_or(InvalidVideoStream)?;
319        let frame = YuvFrame::new(width, height, chroma_subsampling);
320        let prev_frame = frame.clone();
321        Ok(Self {
322            width,
323            height,
324            quant,
325            raw,
326            frame,
327            prev_frame,
328        })
329    }
330
331    /// Returns video frame width.
332    pub fn width(&self) -> NonZero<u16> {
333        self.width
334    }
335
336    /// Returns video frame height.
337    pub fn height(&self) -> NonZero<u16> {
338        self.height
339    }
340
341    /// Decode next frame as RGB888.
342    pub fn read_rgb888_frame(
343        &mut self,
344        input: &mut impl Input,
345        rgb_frame: &mut [u8],
346    ) -> Result<(), InvalidVideoStream> {
347        assert_eq!(
348            usize::from(self.width.get()) * usize::from(self.height.get()) * 3,
349            rgb_frame.len()
350        );
351        if self.raw {
352            return self.read_rgb888_frame_raw(input, rgb_frame);
353        }
354        self.read_yuv420p_frame(input)?;
355        match self.frame {
356            YuvFrame::Yuv420p(ref frame) => {
357                let (y, u, v) = frame.as_slices();
358                yuv420p_to_rgb888(y, u, v, self.width, rgb_frame);
359            }
360            YuvFrame::Yuv444p(ref frame) => {
361                let (y, u, v) = frame.as_slices();
362                yuv444p_to_rgb888(y, u, v, self.width, rgb_frame);
363            }
364        }
365        Ok(())
366    }
367
368    fn read_rgb888_frame_raw(
369        &mut self,
370        input: &mut impl Input,
371        rgb_frame: &mut [u8],
372    ) -> Result<(), InvalidVideoStream> {
373        let format = RawVideoFrameFormat::decode(input)?;
374        match format {
375            RawVideoFrameFormat::Rgb888Indexed8 => {
376                let palette_len = 256 * 3;
377                let indices_len = usize::from(self.width.get()) * usize::from(self.height.get());
378                let mut indexed_rgb = vec![0_u8; palette_len + indices_len];
379                input.read(&mut indexed_rgb[..])?;
380                let (palette, indices) = indexed_rgb.split_at(palette_len);
381                for (i, rgb) in indices.iter().copied().zip(rgb_frame.chunks_exact_mut(3)) {
382                    let j = 3 * i as usize;
383                    rgb.copy_from_slice(&palette[j..j + 3]);
384                }
385            }
386            RawVideoFrameFormat::Rgb888 => input.read(rgb_frame)?,
387        }
388        Ok(())
389    }
390
391    /// Decode next frame as RGBA8888.
392    ///
393    /// All pixels are fully opaque.
394    pub fn read_rgba8888_frame(
395        &mut self,
396        input: &mut impl Input,
397        rgba_frame: &mut [u8],
398    ) -> Result<(), InvalidVideoStream> {
399        assert_eq!(
400            usize::from(self.width.get()) * usize::from(self.height.get()) * 4,
401            rgba_frame.len()
402        );
403        if self.raw {
404            return self.read_rgba8888_frame_raw(input, rgba_frame);
405        }
406        self.read_yuv420p_frame(input)?;
407        match self.frame {
408            YuvFrame::Yuv420p(ref frame) => {
409                let (y, u, v) = frame.as_slices();
410                yuv420p_to_rgba8888(y, u, v, self.width, rgba_frame);
411            }
412            YuvFrame::Yuv444p(ref frame) => {
413                let (y, u, v) = frame.as_slices();
414                yuv444p_to_rgba8888(y, u, v, self.width, rgba_frame);
415            }
416        }
417        Ok(())
418    }
419
420    fn read_rgba8888_frame_raw(
421        &mut self,
422        input: &mut impl Input,
423        rgba_frame: &mut [u8],
424    ) -> Result<(), InvalidVideoStream> {
425        let format = RawVideoFrameFormat::decode(input)?;
426        match format {
427            RawVideoFrameFormat::Rgb888Indexed8 => {
428                let palette_len = 256 * 3;
429                let indices_len = usize::from(self.width.get()) * usize::from(self.height.get());
430                let mut indexed_rgb = vec![0_u8; palette_len + indices_len];
431                input.read(&mut indexed_rgb[..])?;
432                let (palette, indices) = indexed_rgb.split_at(palette_len);
433                for (i, rgba) in indices.iter().copied().zip(rgba_frame.chunks_exact_mut(4)) {
434                    let j = 3 * i as usize;
435                    rgba[..3].copy_from_slice(&palette[j..j + 3]);
436                    rgba[3] = u8::MAX;
437                }
438            }
439            RawVideoFrameFormat::Rgb888 => {
440                let rgb_len = 3 * usize::from(self.width.get()) * usize::from(self.height.get());
441                let mut rgb_frame = vec![0_u8; rgb_len];
442                input.read(&mut rgb_frame[..])?;
443                for (rgb, rgba) in rgb_frame
444                    .chunks_exact(3)
445                    .zip(rgba_frame.chunks_exact_mut(4))
446                {
447                    rgba[..3].copy_from_slice(rgb);
448                    rgba[3] = u8::MAX;
449                }
450            }
451        }
452        Ok(())
453    }
454
455    fn read_yuv420p_frame(&mut self, input: &mut impl Input) -> Result<(), InvalidVideoStream> {
456        let (uv_width, uv_height) = (self.frame.uv_width(), self.frame.uv_height());
457        let (y, u, v) = self.frame.as_mut_slices();
458        let (y_prev, u_prev, v_prev) = self.prev_frame.as_mut_slices();
459        // Read in reverse order: v, u, y.
460        let mut ans_input = rans::JamInputAsAnsInput(input);
461        // v
462        signmag::decode(&mut ans_input, v)?;
463        delta::decode(v, v_prev);
464        haar::backward(v, uv_width, uv_height, self.quant);
465        // u
466        signmag::decode(&mut ans_input, u)?;
467        delta::decode(u, u_prev);
468        haar::backward(u, uv_width, uv_height, self.quant);
469        // y
470        signmag::decode(&mut ans_input, y)?;
471        delta::decode(y, y_prev);
472        haar::backward(y, self.width, self.height, self.quant);
473        Ok(())
474    }
475}
476
477#[cfg(test)]
478mod tests {
479    use super::*;
480    use crate::ToUsize;
481    use alloc::vec;
482    use core::ops::RangeInclusive;
483    use proptest::{collection, prelude::*};
484
485    impl Config {
486        fn no_quantization() -> Self {
487            Self {
488                quantization_level: 0,
489                ..Default::default()
490            }
491        }
492    }
493
494    impl Encoder {
495        fn write_yuv420p_frame(
496            &mut self,
497            y: &[i16],
498            u: &[i16],
499            v: &[i16],
500            output: &mut impl Output,
501        ) -> Stats {
502            let (y0, u0, v0) = self.frame.as_mut_slices();
503            y0.copy_from_slice(y);
504            u0.copy_from_slice(u);
505            v0.copy_from_slice(v);
506            self.write_frame(output)
507        }
508    }
509
510    #[test]
511    fn encoder_simple_works() {
512        let frames = [
513            [
514                [1, 2, 3, 4].as_slice(), // y
515                [5].as_slice(),          // u
516                [6].as_slice(),          // v
517            ],
518            [
519                [7, 8, 9, 10].as_slice(), // y
520                [11].as_slice(),          // u
521                [12].as_slice(),          // v
522            ],
523        ];
524        let mut output = Vec::new();
525        let mut encoder = Encoder::new(
526            NonZero::new(2).unwrap(),
527            NonZero::new(2).unwrap(),
528            Config::no_quantization(),
529        );
530        encoder.start(&mut output);
531        for [y, u, v] in frames {
532            encoder.write_yuv420p_frame(y, u, v, &mut output);
533        }
534        encoder.finish(&mut output);
535        let mut decoder_input = &output[..];
536        let mut decoder = Decoder::new(&mut decoder_input).unwrap();
537        for [y, u, v] in frames {
538            decoder.read_yuv420p_frame(&mut decoder_input).unwrap();
539            let (decoded_y, decoded_u, decoded_v) = decoder.frame.as_mut_slices();
540            assert_eq!(y, decoded_y);
541            assert_eq!(u, decoded_u);
542            assert_eq!(v, decoded_v);
543        }
544    }
545
546    fn frames(
547        width: RangeInclusive<u16>,
548        height: RangeInclusive<u16>,
549        num_frames: RangeInclusive<u32>,
550    ) -> impl Strategy<Value = (u16, u16, Vec<Vec<u8>>)> {
551        (width, height, num_frames).prop_flat_map(|(width, height, num_frames)| {
552            let (y_len, uv_len, ..) =
553                yuv420p_dimensions(NonZero::new(width).unwrap(), NonZero::new(height).unwrap());
554            (
555                width..=width,
556                height..=height,
557                vec![
558                    collection::vec(any::<u8>(), y_len.to_usize() + 2 * uv_len.to_usize());
559                    num_frames.to_usize()
560                ],
561            )
562        })
563    }
564
565    #[test]
566    fn encoder_works() {
567        proptest!(|((width, height, frames) in frames(1..=10, 1..=10, 1..=3))| {
568            let width = NonZero::new(width).unwrap();
569            let height = NonZero::new(height).unwrap();
570            let (y_len, uv_len, ..) = yuv420p_dimensions(width, height);
571            let mut output = Vec::new();
572            let mut encoder = Encoder::new(width, height, Config::no_quantization());
573            encoder.start(&mut output);
574            for frame in frames.iter() {
575                let frame: Vec<i16> = frame.iter().copied().map(i16::from).collect();
576                let (y, uv) = frame.split_at(y_len.to_usize());
577                let (u, v) = uv.split_at(uv_len.to_usize());
578                encoder.write_yuv420p_frame(y, u, v, &mut output);
579            }
580            encoder.finish(&mut output);
581            let mut decoder_input = &output[..];
582            let mut decoder = Decoder::new(&mut decoder_input).unwrap();
583            for frame in frames {
584                let frame: Vec<i16> = frame.iter().copied().map(i16::from).collect();
585                let (y, uv) = frame.split_at(y_len.to_usize());
586                let (u, v) = uv.split_at(uv_len.to_usize());
587                decoder.read_yuv420p_frame(&mut decoder_input).unwrap();
588                let (decoded_y, decoded_u, decoded_v) = decoder.frame.as_mut_slices();
589                assert_eq!(y, decoded_y);
590                assert_eq!(u, decoded_u);
591                assert_eq!(v, decoded_v);
592            }
593        });
594    }
595
596    fn rgb_frames(
597        width: RangeInclusive<u16>,
598        height: RangeInclusive<u16>,
599        num_frames: RangeInclusive<u32>,
600    ) -> impl Strategy<Value = (u16, u16, Vec<Vec<u8>>)> {
601        (width, height, num_frames).prop_flat_map(|(width, height, num_frames)| {
602            (
603                width..=width,
604                height..=height,
605                vec![
606                    collection::vec(any::<u8>(), usize::from(width) * usize::from(height) * 3);
607                    num_frames.to_usize()
608                ],
609            )
610        })
611    }
612
613    #[test]
614    fn encoder_lossless_works() {
615        proptest!(|((width, height, frames) in rgb_frames(1..=10, 1..=10, 1..=3))| {
616            let width = NonZero::new(width).unwrap();
617            let height = NonZero::new(height).unwrap();
618            let mut output = Vec::new();
619            let mut encoder = Encoder::new(width, height, Config::default_lossless());
620            encoder.start(&mut output);
621            for frame in frames.iter() {
622                encoder.write_rgb888_frame(frame, &mut output);
623            }
624            encoder.finish(&mut output);
625            let mut decoder_input = &output[..];
626            let mut decoder = Decoder::new(&mut decoder_input).unwrap();
627            for frame in frames {
628                let mut decoded_frame = vec![0_u8; usize::from(width.get()) * usize::from(height.get()) * 3];
629                decoder.read_rgb888_frame(&mut decoder_input, &mut decoded_frame).unwrap();
630                assert_eq!(frame, decoded_frame);
631            }
632        });
633    }
634
635    #[test]
636    fn encoder_raw_works() {
637        proptest!(|((width, height, frames) in rgb_frames(1..=10, 1..=10, 1..=3))| {
638            let width = NonZero::new(width).unwrap();
639            let height = NonZero::new(height).unwrap();
640            let mut output = Vec::new();
641            let mut encoder = Encoder::new(width, height, Config::default_raw());
642            encoder.start(&mut output);
643            for frame in frames.iter() {
644                encoder.write_rgb888_frame(frame, &mut output);
645            }
646            encoder.finish(&mut output);
647            let mut decoder_input = &output[..];
648            let mut decoder = Decoder::new(&mut decoder_input).unwrap();
649            for frame in frames {
650                let mut decoded_frame = vec![0_u8; usize::from(width.get()) * usize::from(height.get()) * 3];
651                decoder.read_rgb888_frame(&mut decoder_input, &mut decoded_frame).unwrap();
652                assert_eq!(frame, decoded_frame);
653            }
654        });
655    }
656}
corevm_codec/video.rs

corevm_codec/
video.rs
