1#![allow(clippy::needless_range_loop, clippy::too_many_arguments)]
2
3mod dec;
4mod enc;
5
6#[cfg(test)]
7pub(crate) mod tests;
8
9use crate::media::{self, Error, Frame, Packet, PixelFormat, Plane, Result, TimeBase};
10
11pub use dec::ColorInfo;
12pub use enc::{Preset, Profile};
13
14pub struct Decoder {
15 inner: dec::Decoder,
16 time_base: TimeBase,
17 color: Option<ColorInfo>,
18}
19
20impl Default for Decoder {
21 fn default() -> Self {
22 Self::new()
23 }
24}
25
26impl Decoder {
27 pub fn new() -> Self {
28 Self {
29 inner: dec::Decoder::new(),
30 time_base: TimeBase::new(1, 90000),
31 color: None,
32 }
33 }
34
35 pub fn with_threads(threads: usize) -> Self {
36 Self {
37 inner: dec::Decoder::with_threads(threads),
38 time_base: TimeBase::new(1, 90000),
39 color: None,
40 }
41 }
42
43 pub fn with_extradata(avcc: &[u8]) -> Result<Self> {
44 let mut d = Self::new();
45 d.inner.configure_avcc(avcc).map_err(dec_error)?;
46 Ok(d)
47 }
48
49 pub fn set_time_base(&mut self, time_base: TimeBase) {
50 self.time_base = time_base;
51 }
52
53 pub fn color(&self) -> Option<ColorInfo> {
54 self.color
55 }
56
57 pub fn threads(&self) -> usize {
58 self.inner.threads()
59 }
60
61 pub fn decode(&mut self, packet: &Packet) -> Result<Vec<Frame>> {
62 let pics = self
63 .inner
64 .decode(&packet.data, packet.pts)
65 .map_err(dec_error)?;
66 Ok(pics.into_iter().map(|p| self.frame(p)).collect())
67 }
68
69 pub fn recycle(&mut self, frame: Frame) {
70 for p in frame.planes {
71 self.inner.recycle(p.data);
72 }
73 }
74
75 pub fn flush(&mut self) -> Result<Vec<Frame>> {
76 let pics = self.inner.flush();
77 if let Some(e) = self.inner.take_error() {
78 return Err(dec_error(e));
79 }
80 Ok(pics.into_iter().map(|p| self.frame(p)).collect())
81 }
82
83 fn frame(&mut self, p: dec::Picture) -> Frame {
84 self.color = Some(p.color);
85 let (cw, ch) = (p.chroma_width as usize, p.chroma_height as usize);
86 let plane = |data: Vec<u8>, stride: usize, width: usize, height: usize| Plane {
87 data,
88 stride,
89 width,
90 height,
91 };
92 Frame {
93 format: PixelFormat::Yuv420p8,
94 width: p.width,
95 height: p.height,
96 pts: p.pts,
97 time_base: self.time_base,
98 planes: [
99 plane(p.y, p.y_stride, p.width as usize, p.height as usize),
100 plane(p.u, p.uv_stride, cw, ch),
101 plane(p.v, p.uv_stride, cw, ch),
102 ],
103 }
104 }
105}
106
107fn dec_error(e: dec::Error) -> Error {
108 e.media("h264")
109}
110
111fn enc_error(e: enc::Error) -> Error {
112 Error::Invalid(format!("h264 encoder: {e}"))
113}
114
115#[deprecated(note = "use pito_codecs::media::Rate")]
116pub type Rate = crate::media::Rate;
117
118#[deprecated(note = "use pito_codecs::media::Bitstream")]
119pub type Bitstream = crate::media::Bitstream;
120
121#[derive(Debug, Clone, PartialEq)]
122pub struct EncoderConfig {
123 pub width: u32,
124 pub height: u32,
125 pub fps: (u32, u32),
126 pub profile: Profile,
127 pub preset: Preset,
128 pub rate: media::Rate,
129 pub gop: u32,
130 pub b_frames: u8,
131 pub threads: usize,
132 pub bitstream: media::Bitstream,
133}
134
135impl Default for EncoderConfig {
136 fn default() -> Self {
137 Self {
138 width: 1920,
139 height: 1080,
140 fps: (30, 1),
141 profile: Profile::High,
142 preset: Preset::Balanced,
143 rate: media::Rate::Crf(23),
144 gop: 250,
145 b_frames: 2,
146 threads: 0,
147 bitstream: media::Bitstream::LengthPrefixed,
148 }
149 }
150}
151
152pub struct Encoder {
153 inner: enc::Encoder,
154 width: u32,
155 height: u32,
156}
157
158impl Encoder {
159 pub fn new(config: EncoderConfig) -> Result<Self> {
160 let mut cfg =
161 enc::EncoderConfig::new(config.width, config.height, config.fps.0, config.fps.1);
162 cfg.profile = config.profile;
163 cfg.preset = config.preset;
164 cfg.rate = match config.rate {
165 media::Rate::Crf(q) => enc::RateControl::Crf(f32::from(q)),
166 media::Rate::Bitrate(kbps) => enc::RateControl::Cbr { kbps },
167 media::Rate::Vbr { kbps, max_kbps } => enc::RateControl::Vbr {
168 target_kbps: kbps,
169 max_kbps,
170 },
171 media::Rate::Qp(q) => enc::RateControl::Qp(q),
172 };
173 cfg.keyint = config.gop;
174 cfg.bframes = config.b_frames;
175 cfg.threads = config.threads;
176 cfg.format = match config.bitstream {
177 media::Bitstream::LengthPrefixed => enc::PacketFormat::LengthPrefixed,
178 media::Bitstream::AnnexB => enc::PacketFormat::AnnexB,
179 };
180 let inner = enc::Encoder::new(cfg).map_err(enc_error)?;
181 Ok(Self {
182 inner,
183 width: config.width,
184 height: config.height,
185 })
186 }
187
188 pub fn avcc(&self) -> Vec<u8> {
189 self.inner.avcc()
190 }
191
192 pub fn sps(&self) -> Vec<u8> {
193 self.inner.sps_pps().0
194 }
195
196 pub fn pps(&self) -> Vec<u8> {
197 self.inner.sps_pps().1
198 }
199
200 pub fn encode(&mut self, frame: &Frame) -> Result<Vec<Packet>> {
201 if frame.format != PixelFormat::Yuv420p8 {
202 return Err(Error::Unsupported(format!(
203 "h264 encoder: {:?} input, 8-bit 4:2:0 only",
204 frame.format
205 )));
206 }
207 if (frame.width, frame.height) != (self.width, self.height) {
208 return Err(Error::Invalid(format!(
209 "h264 encoder: {}x{} frame for a {}x{} stream",
210 frame.width, frame.height, self.width, self.height
211 )));
212 }
213 let [y, u, v] = &frame.planes;
214 if u.stride != v.stride {
215 return Err(Error::Invalid(
216 "h264 encoder: chroma planes with different strides".into(),
217 ));
218 }
219 let yuv = enc::YuvFrame {
220 y: &y.data,
221 u: &u.data,
222 v: &v.data,
223 y_stride: y.stride,
224 uv_stride: u.stride,
225 };
226 let packets = self.inner.encode(&yuv, frame.pts).map_err(enc_error)?;
227 Ok(packets.into_iter().map(packet).collect())
228 }
229
230 pub fn flush(&mut self) -> Result<Vec<Packet>> {
231 let packets = self.inner.flush().map_err(enc_error)?;
232 Ok(packets.into_iter().map(packet).collect())
233 }
234}
235
236fn packet(p: enc::Packet) -> Packet {
237 Packet {
238 data: p.data,
239 pts: p.pts,
240 dts: p.dts,
241 keyframe: p.keyframe,
242 }
243}