pixelflux/webcam/
decode.rs

1//! Decoders for the client camera uplink.
2//!
3//! Clients send whatever their browser can produce: H.264 (WebCodecs or the WebRTC media track),
4//! VP8/VP9/AV1/HEVC (same sources), or MJPEG from the canvas fallback. Every codec lands in the same
5//! I420 view so the rest of the pipeline is codec-agnostic. Inter-coded codecs go through FFmpeg's
6//! avcodec (already linked for VA-API); MJPEG goes through TurboJPEG, which pixelflux already uses for
7//! its own JPEG stripes. Decoding is software here; a hardware decoder only changes which `Decoder`
8//! is constructed.
9
10use std::ffi::{c_int, CStr};
11use std::ptr;
12
13use ffmpeg_sys_next as ff;
14use turbojpeg::{Decompressor, Image, PixelFormat, Subsamp, YuvImage};
15
16use super::convert::{I420Buffer, I420View};
17
18/// Input codecs, by wire id. The ids are part of the Selkies WebSocket framing and are exported to
19/// Python as `VirtualCamera.CODEC_*`.
20#[derive(Clone, Copy, Debug, PartialEq, Eq)]
21#[repr(u32)]
22pub enum Codec {
23    Mjpeg = 0,
24    H264 = 1,
25    Vp8 = 2,
26    Vp9 = 3,
27    Av1 = 4,
28    Hevc = 5,
29}
30
31impl Codec {
32    pub fn from_id(id: u32) -> Option<Self> {
33        Some(match id {
34            0 => Codec::Mjpeg,
35            1 => Codec::H264,
36            2 => Codec::Vp8,
37            3 => Codec::Vp9,
38            4 => Codec::Av1,
39            5 => Codec::Hevc,
40            _ => return None,
41        })
42    }
43
44    pub fn name(self) -> &'static str {
45        match self {
46            Codec::Mjpeg => "mjpeg",
47            Codec::H264 => "h264",
48            Codec::Vp8 => "vp8",
49            Codec::Vp9 => "vp9",
50            Codec::Av1 => "av1",
51            Codec::Hevc => "hevc",
52        }
53    }
54
55    /// Whether frames depend on earlier ones, so a dropped frame forces a wait for a keyframe.
56    pub fn is_inter_coded(self) -> bool {
57        !matches!(self, Codec::Mjpeg)
58    }
59
60    fn av_codec_id(self) -> ff::AVCodecID {
61        match self {
62            Codec::Mjpeg => ff::AVCodecID::AV_CODEC_ID_MJPEG,
63            Codec::H264 => ff::AVCodecID::AV_CODEC_ID_H264,
64            Codec::Vp8 => ff::AVCodecID::AV_CODEC_ID_VP8,
65            Codec::Vp9 => ff::AVCodecID::AV_CODEC_ID_VP9,
66            Codec::Av1 => ff::AVCodecID::AV_CODEC_ID_AV1,
67            Codec::Hevc => ff::AVCodecID::AV_CODEC_ID_HEVC,
68        }
69    }
70}
71
72/// Cheap bitstream inspection for the codecs whose keyframes can be recognized without a full
73/// parse; `None` when the codec gives no such signal and the caller's flag must be trusted.
74pub fn sniff_keyframe(codec: Codec, data: &[u8]) -> Option<bool> {
75    match codec {
76        Codec::Mjpeg => Some(true),
77        Codec::H264 => {
78            let mut i = 0;
79            let mut idr = false;
80            while i + 3 < data.len() {
81                if data[i] == 0 && data[i + 1] == 0 && (data[i + 2] == 1 || (data[i + 2] == 0 && i + 4 <= data.len() && data[i + 3] == 1)) {
82                    let off = if data[i + 2] == 1 { 3 } else { 4 };
83                    if i + off < data.len() {
84                        let nal_type = data[i + off] & 0x1F;
85                        if nal_type == 5 {
86                            idr = true;
87                            break;
88                        }
89                    }
90                    i += off;
91                } else {
92                    i += 1;
93                }
94            }
95            Some(idr)
96        }
97        Codec::Vp8 => data.first().map(|b| b & 1 == 0),
98        Codec::Vp9 | Codec::Av1 | Codec::Hevc => None,
99    }
100}
101
102/// Padding avcodec requires past the end of any packet it parses.
103const AV_INPUT_BUFFER_PADDING: usize = 64;
104
105#[derive(Debug)]
106pub enum DecodeError {
107    /// The packet could not be decoded; the stream needs a keyframe to resynchronize.
108    Corrupt(String),
109    /// The decoder itself failed and must be recreated.
110    Fatal(String),
111}
112
113pub trait Decoder {
114    fn codec(&self) -> Codec;
115    /// Decode one encoded frame. `Ok(true)` when `frame()` now holds a new picture.
116    fn decode(&mut self, data: &[u8]) -> Result<bool, DecodeError>;
117    fn frame(&self) -> Option<I420View<'_>>;
118}
119
120pub fn new_decoder(codec: Codec) -> Result<Box<dyn Decoder>, String> {
121    match codec {
122        Codec::Mjpeg => Ok(Box::new(JpegDecoder::new()?)),
123        _ => Ok(Box::new(AvDecoder::new(codec)?)),
124    }
125}
126
127fn ff_err(code: c_int) -> String {
128    let mut buf = [0 as libc::c_char; 128];
129    unsafe {
130        ff::av_strerror(code, buf.as_mut_ptr(), buf.len());
131        CStr::from_ptr(buf.as_ptr()).to_string_lossy().into_owned()
132    }
133}
134
135/// avcodec-backed decoder for the inter-coded codecs.
136pub struct AvDecoder {
137    codec: Codec,
138    ctx: *mut ff::AVCodecContext,
139    pkt: *mut ff::AVPacket,
140    frame: *mut ff::AVFrame,
141    /// Receive target; each decoded picture is moved from here into `frame`.
142    scratch: *mut ff::AVFrame,
143    /// Encoded input with avcodec's trailing padding, reused across frames.
144    input: Vec<u8>,
145    /// Planar copy of frames whose avcodec pixel format is not I420.
146    converted: I420Buffer,
147    have_frame: bool,
148    from_converted: bool,
149}
150
151unsafe impl Send for AvDecoder {}
152
153impl AvDecoder {
154    pub fn new(codec: Codec) -> Result<Self, String> {
155        unsafe {
156            static QUIET: std::sync::Once = std::sync::Once::new();
157            QUIET.call_once(|| ff::av_log_set_level(ff::AV_LOG_ERROR));
158            let c = ff::avcodec_find_decoder(codec.av_codec_id());
159            if c.is_null() {
160                return Err(format!("no avcodec decoder for {}", codec.name()));
161            }
162            let ctx = ff::avcodec_alloc_context3(c);
163            if ctx.is_null() {
164                return Err("avcodec_alloc_context3 failed".into());
165            }
166            (*ctx).thread_count = 2;
167            (*ctx).thread_type = ff::FF_THREAD_SLICE;
168            (*ctx).flags |= ff::AV_CODEC_FLAG_LOW_DELAY as c_int;
169            let rc = ff::avcodec_open2(ctx, c, ptr::null_mut());
170            if rc < 0 {
171                let mut p = ctx;
172                ff::avcodec_free_context(&mut p);
173                return Err(format!("avcodec_open2({}) failed: {}", codec.name(), ff_err(rc)));
174            }
175            let pkt = ff::av_packet_alloc();
176            let frame = ff::av_frame_alloc();
177            let scratch = ff::av_frame_alloc();
178            if pkt.is_null() || frame.is_null() || scratch.is_null() {
179                let mut p = ctx;
180                ff::avcodec_free_context(&mut p);
181                return Err("avcodec packet/frame allocation failed".into());
182            }
183            Ok(AvDecoder {
184                codec,
185                ctx,
186                pkt,
187                frame,
188                scratch,
189                input: Vec::new(),
190                converted: I420Buffer::new(2, 2),
191                have_frame: false,
192                from_converted: false,
193            })
194        }
195    }
196
197    /// Copy a non-I420 planar/semi-planar frame into the I420 scratch. Returns false for formats
198    /// browsers never produce (they would need a full colorspace conversion).
199    fn convert_frame(&mut self) -> bool {
200        unsafe {
201            let f = &*self.frame;
202            let w = f.width as usize;
203            let h = f.height as usize;
204            let cw = w.div_ceil(2);
205            let ch = h.div_ceil(2);
206            self.converted.resize(w, h);
207            let y_len = self.converted.y_len();
208            let uv_len = self.converted.uv_len();
209            let fmt: ff::AVPixelFormat = std::mem::transmute(f.format);
210            let (yp, rest) = self.converted.data.split_at_mut(y_len);
211            let (up, vp) = rest.split_at_mut(uv_len);
212            for row in 0..h {
213                let s = f.data[0].add(row * f.linesize[0] as usize);
214                ptr::copy_nonoverlapping(s, yp.as_mut_ptr().add(row * w), w);
215            }
216            match fmt {
217                ff::AVPixelFormat::AV_PIX_FMT_NV12 | ff::AVPixelFormat::AV_PIX_FMT_NV21 => {
218                    let swap = fmt == ff::AVPixelFormat::AV_PIX_FMT_NV21;
219                    for row in 0..ch {
220                        let s = f.data[1].add(row * f.linesize[1] as usize);
221                        for x in 0..cw {
222                            let a = *s.add(2 * x);
223                            let b = *s.add(2 * x + 1);
224                            up[row * cw + x] = if swap { b } else { a };
225                            vp[row * cw + x] = if swap { a } else { b };
226                        }
227                    }
228                    true
229                }
230                ff::AVPixelFormat::AV_PIX_FMT_YUV422P | ff::AVPixelFormat::AV_PIX_FMT_YUVJ422P => {
231                    for row in 0..ch {
232                        let r0 = (2 * row).min(h - 1);
233                        let r1 = (2 * row + 1).min(h - 1);
234                        for (plane, dst) in [(1usize, &mut *up), (2usize, &mut *vp)] {
235                            let s0 = f.data[plane].add(r0 * f.linesize[plane] as usize);
236                            let s1 = f.data[plane].add(r1 * f.linesize[plane] as usize);
237                            for x in 0..cw {
238                                dst[row * cw + x] = ((*s0.add(x) as u32 + *s1.add(x) as u32 + 1) / 2) as u8;
239                            }
240                        }
241                    }
242                    true
243                }
244                ff::AVPixelFormat::AV_PIX_FMT_YUV444P | ff::AVPixelFormat::AV_PIX_FMT_YUVJ444P => {
245                    for row in 0..ch {
246                        let r0 = (2 * row).min(h - 1);
247                        let r1 = (2 * row + 1).min(h - 1);
248                        for (plane, dst) in [(1usize, &mut *up), (2usize, &mut *vp)] {
249                            let s0 = f.data[plane].add(r0 * f.linesize[plane] as usize);
250                            let s1 = f.data[plane].add(r1 * f.linesize[plane] as usize);
251                            for x in 0..cw {
252                                let x0 = 2 * x;
253                                let x1 = (2 * x + 1).min(w - 1);
254                                let sum = *s0.add(x0) as u32 + *s0.add(x1) as u32 + *s1.add(x0) as u32 + *s1.add(x1) as u32;
255                                dst[row * cw + x] = ((sum + 2) / 4) as u8;
256                            }
257                        }
258                    }
259                    true
260                }
261                _ => false,
262            }
263        }
264    }
265}
266
267impl Drop for AvDecoder {
268    fn drop(&mut self) {
269        unsafe {
270            ff::av_frame_free(&mut self.frame);
271            ff::av_frame_free(&mut self.scratch);
272            ff::av_packet_free(&mut self.pkt);
273            ff::avcodec_free_context(&mut self.ctx);
274        }
275    }
276}
277
278impl Decoder for AvDecoder {
279    fn codec(&self) -> Codec {
280        self.codec
281    }
282
283    fn decode(&mut self, data: &[u8]) -> Result<bool, DecodeError> {
284        self.have_frame = false;
285        self.input.clear();
286        self.input.extend_from_slice(data);
287        self.input.resize(data.len() + AV_INPUT_BUFFER_PADDING, 0);
288        unsafe {
289            ff::av_packet_unref(self.pkt);
290            (*self.pkt).data = self.input.as_mut_ptr();
291            (*self.pkt).size = data.len() as c_int;
292            let rc = ff::avcodec_send_packet(self.ctx, self.pkt);
293            if rc < 0 && rc != ff::AVERROR(libc::EAGAIN) {
294                return Err(DecodeError::Corrupt(ff_err(rc)));
295            }
296            // avcodec_receive_frame unreferences its destination first, so drain
297            // into the scratch frame and keep the newest picture in `frame`.
298            let mut got = false;
299            loop {
300                let rc = ff::avcodec_receive_frame(self.ctx, self.scratch);
301                if rc == ff::AVERROR(libc::EAGAIN) || rc == ff::AVERROR_EOF {
302                    break;
303                }
304                if rc < 0 {
305                    return Err(DecodeError::Corrupt(ff_err(rc)));
306                }
307                ff::av_frame_unref(self.frame);
308                ff::av_frame_move_ref(self.frame, self.scratch);
309                got = true;
310            }
311            if !got {
312                return Ok(false);
313            }
314            let fmt: ff::AVPixelFormat = std::mem::transmute((*self.frame).format);
315            self.from_converted = !matches!(fmt, ff::AVPixelFormat::AV_PIX_FMT_YUV420P | ff::AVPixelFormat::AV_PIX_FMT_YUVJ420P);
316            if self.from_converted && !self.convert_frame() {
317                return Err(DecodeError::Fatal(format!("unsupported decoded pixel format {}", (*self.frame).format)));
318            }
319            self.have_frame = true;
320            Ok(true)
321        }
322    }
323
324    fn frame(&self) -> Option<I420View<'_>> {
325        if !self.have_frame {
326            return None;
327        }
328        unsafe {
329            let f = &*self.frame;
330            let fmt: ff::AVPixelFormat = std::mem::transmute(f.format);
331            let full_range = f.color_range == ff::AVColorRange::AVCOL_RANGE_JPEG
332                || matches!(fmt, ff::AVPixelFormat::AV_PIX_FMT_YUVJ420P | ff::AVPixelFormat::AV_PIX_FMT_YUVJ422P | ff::AVPixelFormat::AV_PIX_FMT_YUVJ444P);
333            if self.from_converted {
334                return Some(self.converted.view(full_range));
335            }
336            let w = f.width as usize;
337            let h = f.height as usize;
338            let cw = w.div_ceil(2);
339            let ch = h.div_ceil(2);
340            let ys = f.linesize[0] as usize;
341            let us = f.linesize[1] as usize;
342            let vs = f.linesize[2] as usize;
343            Some(I420View {
344                width: w,
345                height: h,
346                y: std::slice::from_raw_parts(f.data[0], ys * (h - 1) + w),
347                u: std::slice::from_raw_parts(f.data[1], us * (ch - 1) + cw),
348                v: std::slice::from_raw_parts(f.data[2], vs * (ch - 1) + cw),
349                y_stride: ys,
350                uv_stride: us.max(vs).min(us),
351                full_range,
352            })
353        }
354    }
355}
356
357/// TurboJPEG decoder: 4:2:0 JPEGs decode straight into I420 planes; other subsamplings take the
358/// RGB route and are converted.
359pub struct JpegDecoder {
360    dec: Decompressor,
361    out: I420Buffer,
362    rgb: Vec<u8>,
363    have_frame: bool,
364}
365
366impl JpegDecoder {
367    pub fn new() -> Result<Self, String> {
368        Ok(JpegDecoder {
369            dec: Decompressor::new().map_err(|e| format!("turbojpeg: {}", e))?,
370            out: I420Buffer::new(2, 2),
371            rgb: Vec::new(),
372            have_frame: false,
373        })
374    }
375}
376
377impl Decoder for JpegDecoder {
378    fn codec(&self) -> Codec {
379        Codec::Mjpeg
380    }
381
382    fn decode(&mut self, data: &[u8]) -> Result<bool, DecodeError> {
383        self.have_frame = false;
384        let hdr = self.dec.read_header(data).map_err(|e| DecodeError::Corrupt(format!("jpeg header: {}", e)))?;
385        if hdr.width == 0 || hdr.height == 0 {
386            return Err(DecodeError::Corrupt("empty jpeg".into()));
387        }
388        self.out.resize(hdr.width, hdr.height);
389        if hdr.subsamp == Subsamp::Sub2x2 {
390            let img = YuvImage { pixels: &mut self.out.data[..], width: hdr.width, align: 1, height: hdr.height, subsamp: Subsamp::Sub2x2 };
391            self.dec.decompress_to_yuv(data, img).map_err(|e| DecodeError::Corrupt(format!("jpeg: {}", e)))?;
392        } else {
393            self.rgb.resize(hdr.width * hdr.height * 4, 0);
394            let img = Image { pixels: &mut self.rgb[..], width: hdr.width, pitch: hdr.width * 4, height: hdr.height, format: PixelFormat::RGBA };
395            self.dec.decompress(data, img).map_err(|e| DecodeError::Corrupt(format!("jpeg: {}", e)))?;
396            let y_len = self.out.y_len();
397            let uv_len = self.out.uv_len();
398            let cw = hdr.width.div_ceil(2);
399            let (yp, rest) = self.out.data.split_at_mut(y_len);
400            let (up, vp) = rest.split_at_mut(uv_len);
401            let mut planar = yuv::YuvPlanarImageMut {
402                y_plane: yuv::BufferStoreMut::Borrowed(yp),
403                y_stride: hdr.width as u32,
404                u_plane: yuv::BufferStoreMut::Borrowed(up),
405                u_stride: cw as u32,
406                v_plane: yuv::BufferStoreMut::Borrowed(vp),
407                v_stride: cw as u32,
408                width: hdr.width as u32,
409                height: hdr.height as u32,
410            };
411            yuv::rgba_to_yuv420(&mut planar, &self.rgb, (hdr.width * 4) as u32, yuv::YuvRange::Full, yuv::YuvStandardMatrix::Bt601, yuv::YuvConversionMode::Fast)
412                .map_err(|e| DecodeError::Corrupt(format!("rgb->i420: {:?}", e)))?;
413        }
414        self.have_frame = true;
415        Ok(true)
416    }
417
418    fn frame(&self) -> Option<I420View<'_>> {
419        if self.have_frame {
420            Some(self.out.view(true))
421        } else {
422            None
423        }
424    }
425}
426
427#[cfg(test)]
428mod tests {
429    use super::*;
430
431    #[test]
432    fn codec_ids_round_trip() {
433        for id in 0..6 {
434            let c = Codec::from_id(id).unwrap();
435            assert_eq!(c as u32, id);
436        }
437        assert!(Codec::from_id(6).is_none());
438    }
439
440    #[test]
441    fn h264_keyframe_sniff() {
442        let idr = [0, 0, 0, 1, 0x67, 0x42, 0, 0, 0, 1, 0x68, 0xCE, 0, 0, 1, 0x65, 0x88];
443        assert_eq!(sniff_keyframe(Codec::H264, &idr), Some(true));
444        let p = [0, 0, 0, 1, 0x41, 0x9A];
445        assert_eq!(sniff_keyframe(Codec::H264, &p), Some(false));
446        assert_eq!(sniff_keyframe(Codec::Vp8, &[0x10, 0, 0]), Some(true));
447        assert_eq!(sniff_keyframe(Codec::Vp8, &[0x11, 0, 0]), Some(false));
448        assert_eq!(sniff_keyframe(Codec::Vp9, &[0]), None);
449    }
450
451    #[test]
452    fn jpeg_round_trip() {
453        let w = 64;
454        let h = 48;
455        let mut img = I420Buffer::new(w, h);
456        let yl = img.y_len();
457        let ul = img.uv_len();
458        img.data[..yl].fill(200);
459        img.data[yl..yl + ul].fill(100);
460        img.data[yl + ul..].fill(150);
461        let src = YuvImage { pixels: &img.data[..], width: w, align: 1, height: h, subsamp: Subsamp::Sub2x2 };
462        let jpeg = turbojpeg::compress_yuv(src, 90).unwrap();
463        let mut dec = JpegDecoder::new().unwrap();
464        assert!(dec.decode(&jpeg).unwrap());
465        let v = dec.frame().unwrap();
466        assert_eq!((v.width, v.height), (w, h));
467        assert!(v.full_range);
468        let mid = v.y[(h / 2) * v.y_stride + w / 2];
469        assert!((mid as i32 - 200).abs() <= 3, "luma {}", mid);
470        let cu = v.u[(h / 4) * v.uv_stride + w / 4];
471        assert!((cu as i32 - 100).abs() <= 3, "cb {}", cu);
472    }
473
474    #[test]
475    fn av_decoder_constructs_for_browser_codecs() {
476        for c in [Codec::H264, Codec::Vp8, Codec::Vp9] {
477            let d = AvDecoder::new(c).unwrap();
478            assert_eq!(d.codec(), c);
479            assert!(d.frame().is_none());
480        }
481    }
482
483    #[test]
484    fn h264_round_trip_through_avcodec() {
485        use openh264::encoder::Encoder;
486        use openh264::formats::YUVBuffer;
487        let (w, h) = (64usize, 48usize);
488        let mut yuv = vec![0u8; w * h * 3 / 2];
489        yuv[..w * h].fill(145);
490        yuv[w * h..w * h + w * h / 4].fill(54);
491        yuv[w * h + w * h / 4..].fill(34);
492        let source = YUVBuffer::from_vec(yuv, w, h);
493        let mut enc = Encoder::new().unwrap();
494        let mut dec = AvDecoder::new(Codec::H264).unwrap();
495        let mut decoded = 0;
496        for _ in 0..6 {
497            let bitstream = enc.encode(&source).unwrap().to_vec();
498            assert!(!bitstream.is_empty());
499            if dec.decode(&bitstream).unwrap() {
500                decoded += 1;
501                let v = dec.frame().unwrap();
502                assert_eq!((v.width, v.height), (w, h));
503                assert!(!v.full_range);
504                let y = v.y[(h / 2) * v.y_stride + w / 2];
505                let u = v.u[(h / 4) * v.uv_stride + w / 4];
506                let vv = v.v[(h / 4) * v.uv_stride + w / 4];
507                assert!((y as i32 - 145).abs() <= 6 && (u as i32 - 54).abs() <= 6 && (vv as i32 - 34).abs() <= 6, "yuv {} {} {}", y, u, vv);
508            }
509        }
510        assert!(decoded >= 5, "decoded {} of 6 frames", decoded);
511    }
512
513    #[test]
514    fn av_decoder_rejects_garbage_without_fatal() {
515        let mut d = AvDecoder::new(Codec::H264).unwrap();
516        match d.decode(&[1, 2, 3, 4, 5, 6, 7, 8]) {
517            Ok(false) | Err(DecodeError::Corrupt(_)) => {}
518            Ok(true) => panic!("garbage decoded to a frame"),
519            Err(DecodeError::Fatal(e)) => panic!("fatal: {}", e),
520        }
521    }
522}
pixelflux/webcam/decode.rs

pixelflux/webcam/
decode.rs
