image_av1/encoder/
encoding.rs

1use std::io::Write;
2
3use image::{
4    error::{ParameterError, ParameterErrorKind},
5    imageops::FilterType,
6    DynamicImage, ImageError, Rgba, RgbaImage,
7};
8use rav1e::prelude::*;
9
10use crate::{utils::Result, Av1Encoder};
11
12impl Av1Encoder {
13    pub fn resize_image(&mut self, image: DynamicImage, filter: FilterType) -> RgbaImage {
14        image.resize_exact(self.config.width as u32, self.config.height as u32, filter).to_rgba8()
15    }
16
17    pub fn write_image(&mut self, image: &RgbaImage) -> Result<usize> {
18        self.check_size(image)?;
19        let mut size = 0;
20        let mut ctx = self.encode_context()?;
21        let mut frame = ctx.new_frame();
22        self.init_frame_3(&image, &mut frame)?;
23        ctx.send_frame(frame).unwrap();
24        ctx.flush();
25        loop {
26            match ctx.receive_packet() {
27                Ok(packet) => {
28                    size = self.output.write(&packet.data)?;
29                    continue;
30                }
31                Err(EncoderStatus::Encoded) => continue,
32                Err(EncoderStatus::LimitReached) => break,
33                Err(err) => Err(err).unwrap(),
34            }
35        }
36        Ok(size)
37    }
38    pub fn write_image_repeats(&mut self, image: RgbaImage, count: usize) -> Result<usize> {
39        let mut ctx = self.encode_context()?;
40        let frame = ctx.new_frame();
41        let mut size = 0;
42        for _ in 0..count {
43            ctx.send_frame(frame.clone()).unwrap();
44        }
45        ctx.flush();
46        loop {
47            match ctx.receive_packet() {
48                Ok(packet) => {
49                    size = self.output.write(&packet.data)?;
50                    continue;
51                }
52                Err(EncoderStatus::Encoded) => continue,
53                Err(EncoderStatus::LimitReached) => break,
54                Err(err) => Err(err).unwrap(),
55            }
56        }
57        Ok(size)
58    }
59    pub fn check_size(&self, image: &RgbaImage) -> Result<()> {
60        if self.config.width != image.width() as usize {
61            Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))?
62        }
63        if self.config.height != image.height() as usize {
64            Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))?
65        }
66        Ok(())
67    }
68
69    fn encode_context(&self) -> Result<Context<u8>> {
70        if self.config.chroma_sampling != ChromaSampling::Cs444 {
71            panic!("Only 444 chroma sampling is supported")
72        }
73        let config = Config::new().with_encoder_config(self.config.clone()).with_rate_control(self.rate_control.clone());
74        match config.new_context::<u8>() {
75            Ok(o) => Ok(o),
76            Err(e) => {
77                panic!("Error creating context: {}", e)
78            }
79        }
80    }
81    fn init_frame_3(&self, image: &RgbaImage, frame: &mut Frame<u8>) -> Result<()> {
82        let width = self.config.width;
83        let height = self.config.height;
84        let mut f = frame.planes.iter_mut();
85        let mut planes = image.pixels();
86
87        // it doesn't seem to be necessary to fill padding area
88        let mut y = f.next().unwrap().mut_slice(Default::default());
89        let mut u = f.next().unwrap().mut_slice(Default::default());
90        let mut v = f.next().unwrap().mut_slice(Default::default());
91
92        for ((y, u), v) in y.rows_iter_mut().zip(u.rows_iter_mut()).zip(v.rows_iter_mut()).take(height) {
93            let y = &mut y[..width];
94            let u = &mut u[..width];
95            let v = &mut v[..width];
96            for ((y, u), v) in y.iter_mut().zip(u).zip(v) {
97                let px = planes.next().expect("Too few pixels");
98                let px = rgba_to_yuv(px);
99                *y = px[0];
100                *u = px[1];
101                *v = px[2];
102            }
103        }
104        Ok(())
105    }
106
107    fn init_frame_1(width: usize, height: usize, planes: &[u8], frame: &mut Frame<u8>) -> Result<()> {
108        let mut y = frame.planes[0].mut_slice(Default::default());
109        let mut planes = planes.into_iter();
110        for y in y.rows_iter_mut().take(height) {
111            let y = &mut y[..width];
112            for y in y.iter_mut() {
113                *y = *planes.next().expect("Too few pixels");
114            }
115        }
116        Ok(())
117    }
118}
119
120// ## RGB to YUV
121// Y = (( 66 * R + 129 * G +  25 * B + 128) >> 8) +  16
122// U = ((-38 * R -  74 * G + 112 * B + 128) >> 8) + 128
123// V = ((112 * R -  94 * G -  18 * B + 128) >> 8) + 128
124fn rgba_to_yuv(rgba: &Rgba<u8>) -> [u8; 3] {
125    let r = rgba[0] as i32;
126    let g = rgba[1] as i32;
127    let b = rgba[2] as i32;
128    let y = ((66 * r + 129 * g + 25 * b + 128) >> 8) + 16;
129    let u = ((-38 * r - 74 * g + 112 * b + 128) >> 8) + 128;
130    let v = ((112 * r - 94 * g - 18 * b + 128) >> 8) + 128;
131    [y as u8, u as u8, v as u8]
132}
image_av1/encoder/encoding.rs

image_av1/encoder/
encoding.rs
