diff --git a/src/video/codec/h265_bitstream.rs b/src/video/codec/h265_bitstream.rs new file mode 100644 index 00000000..0c0b58f7 --- /dev/null +++ b/src/video/codec/h265_bitstream.rs @@ -0,0 +1,112 @@ +const VPS_NAL_TYPE: u8 = 32; +const SPS_NAL_TYPE: u8 = 33; +const PPS_NAL_TYPE: u8 = 34; + +fn find_start_code(data: &[u8], from: usize) -> Option<(usize, usize)> { + let mut offset = from; + while offset + 3 <= data.len() { + if offset + 4 <= data.len() && data[offset..offset + 4] == [0, 0, 0, 1] { + return Some((offset, 4)); + } + if data[offset..offset + 3] == [0, 0, 1] { + return Some((offset, 3)); + } + offset += 1; + } + None +} + +fn for_each_nal(data: &[u8], mut visit: impl FnMut(u8, &[u8])) { + let mut cursor = 0; + while let Some((start, start_code_len)) = find_start_code(data, cursor) { + let nal_start = start + start_code_len; + if nal_start + 2 > data.len() { + break; + } + let next_start = find_start_code(data, nal_start) + .map(|(offset, _)| offset) + .unwrap_or(data.len()); + let mut nal_end = next_start; + while nal_end > nal_start && data[nal_end - 1] == 0 { + nal_end -= 1; + } + if nal_end >= nal_start + 2 { + visit((data[nal_start] >> 1) & 0x3f, &data[nal_start..nal_end]); + } + if next_start == data.len() { + break; + } + cursor = next_start; + } +} + +pub fn is_keyframe(data: &[u8]) -> bool { + let mut keyframe = false; + for_each_nal(data, |nal_type, _| { + if (16..=23).contains(&nal_type) { + keyframe = true; + } + }); + keyframe +} + +pub fn extract_vps_sps_pps(data: &[u8]) -> (Option>, Option>, Option>) { + let mut vps = None; + let mut sps = None; + let mut pps = None; + for_each_nal(data, |nal_type, nal| match nal_type { + VPS_NAL_TYPE => vps = Some(nal.to_vec()), + SPS_NAL_TYPE => sps = Some(nal.to_vec()), + PPS_NAL_TYPE => pps = Some(nal.to_vec()), + _ => {} + }); + (vps, sps, pps) +} + +pub fn has_vps_sps_pps(data: &[u8]) -> bool { + let (vps, sps, pps) = extract_vps_sps_pps(data); + vps.is_some() && sps.is_some() && pps.is_some() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn detects_irap_but_not_trail_frame() { + assert!(is_keyframe(&[0, 0, 0, 1, 19 << 1, 1, 0xaa])); + assert!(is_keyframe(&[0, 0, 1, 21 << 1, 1, 0xbb])); + assert!(!is_keyframe(&[0, 0, 0, 1, 1 << 1, 1, 0xcc])); + } + + #[test] + fn extracts_parameter_sets() { + let data = [ + 0, + 0, + 0, + 1, + 32 << 1, + 1, + 0xaa, + 0, + 0, + 1, + 33 << 1, + 1, + 0xbb, + 0, + 0, + 0, + 1, + 34 << 1, + 1, + 0xcc, + ]; + let (vps, sps, pps) = extract_vps_sps_pps(&data); + assert_eq!(vps.unwrap(), [32 << 1, 1, 0xaa]); + assert_eq!(sps.unwrap(), [33 << 1, 1, 0xbb]); + assert_eq!(pps.unwrap(), [34 << 1, 1, 0xcc]); + assert!(has_vps_sps_pps(&data)); + } +} diff --git a/src/video/codec/mod.rs b/src/video/codec/mod.rs index ba0ab378..cfa663a3 100644 --- a/src/video/codec/mod.rs +++ b/src/video/codec/mod.rs @@ -8,6 +8,7 @@ pub mod convert; pub mod h264; pub mod h264_bitstream; pub mod h265; +pub mod h265_bitstream; pub mod jpeg; pub mod registry; pub mod self_check; diff --git a/src/video/pipeline/shared.rs b/src/video/pipeline/shared.rs index 6cacd308..ddfe33a0 100644 --- a/src/video/pipeline/shared.rs +++ b/src/video/pipeline/shared.rs @@ -48,9 +48,9 @@ use crate::video::capture::status::{ signal_status_from_capture_kind, CaptureIoErrorKind, }; use crate::video::capture::{BridgeContext, CaptureReadError, CaptureStream}; -use crate::video::codec::h264_bitstream; use crate::video::codec::registry::{EncoderBackend, VideoEncoderType}; use crate::video::codec::MjpegToNv12Decoder; +use crate::video::codec::{h264_bitstream, h265_bitstream}; use crate::video::device::parse_bridge_kind; use crate::video::device::VideoControlMode; use crate::video::format::{PixelFormat, Resolution}; @@ -276,6 +276,15 @@ pub struct SharedVideoPipelineStats { pub current_fps: f32, } +#[derive(Default)] +struct CachedH26xParameterSets { + h264_sps: Option>, + h264_pps: Option>, + h265_vps: Option>, + h265_sps: Option>, + h265_pps: Option>, +} + /// Universal shared video pipeline pub struct SharedVideoPipeline { config: RwLock, @@ -296,6 +305,7 @@ pub struct SharedVideoPipeline { sequence: AtomicU64, /// Atomic flag for keyframe request (avoids lock contention) keyframe_requested: AtomicBool, + parameter_sets: ParkingMutex, /// Pipeline start time for monotonic PTS calculation (microseconds from process start). /// Uses AtomicI64 instead of Mutex for lock-free access. pipeline_start_time_us: AtomicI64, @@ -335,6 +345,7 @@ impl SharedVideoPipeline { running_flag: AtomicBool::new(false), sequence: AtomicU64::new(0), keyframe_requested: AtomicBool::new(false), + parameter_sets: ParkingMutex::new(CachedH26xParameterSets::default()), pipeline_start_time_us: AtomicI64::new(0), pending_sync_geometry: ParkingMutex::new(None), device_lost_reason: ParkingMutex::new(None), @@ -507,6 +518,86 @@ impl SharedVideoPipeline { let _ = self.h264_profile_level_id.send(Some(profile_level_id)); } + fn inspect_and_parameterize_packet( + &self, + codec: VideoEncoderType, + data: Bytes, + ffmpeg_keyframe: bool, + ) -> (Bytes, bool) { + match codec { + VideoEncoderType::H264 => { + let (sps, pps) = h264_bitstream::extract_sps_pps(data.as_ref()); + let is_keyframe = ffmpeg_keyframe || h264_bitstream::is_keyframe(data.as_ref()); + let mut cache = self.parameter_sets.lock(); + if let Some(sps) = sps.as_ref() { + cache.h264_sps = Some(sps.clone()); + } + if let Some(pps) = pps.as_ref() { + cache.h264_pps = Some(pps.clone()); + } + + if !is_keyframe || (sps.is_some() && pps.is_some()) { + return (data, is_keyframe); + } + + match (&cache.h264_sps, &cache.h264_pps) { + (Some(cached_sps), Some(cached_pps)) => { + let mut output = Vec::with_capacity( + data.len() + cached_sps.len() + cached_pps.len() + 8, + ); + output.extend_from_slice(&[0, 0, 0, 1]); + output.extend_from_slice(cached_sps); + output.extend_from_slice(&[0, 0, 0, 1]); + output.extend_from_slice(cached_pps); + output.extend_from_slice(data.as_ref()); + debug!("[Pipeline] Prepended cached SPS/PPS to H264 IDR"); + (Bytes::from(output), true) + } + _ => (data, true), + } + } + VideoEncoderType::H265 => { + let (vps, sps, pps) = h265_bitstream::extract_vps_sps_pps(data.as_ref()); + let is_keyframe = ffmpeg_keyframe || h265_bitstream::is_keyframe(data.as_ref()); + let mut cache = self.parameter_sets.lock(); + if let Some(vps) = vps.as_ref() { + cache.h265_vps = Some(vps.clone()); + } + if let Some(sps) = sps.as_ref() { + cache.h265_sps = Some(sps.clone()); + } + if let Some(pps) = pps.as_ref() { + cache.h265_pps = Some(pps.clone()); + } + + if !is_keyframe || (vps.is_some() && sps.is_some() && pps.is_some()) { + return (data, is_keyframe); + } + + match (&cache.h265_vps, &cache.h265_sps, &cache.h265_pps) { + (Some(cached_vps), Some(cached_sps), Some(cached_pps)) => { + let mut output = Vec::with_capacity( + data.len() + + cached_vps.len() + + cached_sps.len() + + cached_pps.len() + + 12, + ); + for parameter_set in [cached_vps, cached_sps, cached_pps] { + output.extend_from_slice(&[0, 0, 0, 1]); + output.extend_from_slice(parameter_set); + } + output.extend_from_slice(data.as_ref()); + debug!("[Pipeline] Prepended cached VPS/SPS/PPS to H265 IRAP"); + (Bytes::from(output), true) + } + _ => (data, true), + } + } + _ => (data, ffmpeg_keyframe), + } + } + fn broadcast_encoded(&self, frame: Arc) { let subscribers = { let guard = self.subscribers.read(); @@ -1193,9 +1284,11 @@ impl SharedVideoPipeline { })?; if let Some((data, is_keyframe)) = packet { + let (data, is_keyframe) = + self.inspect_and_parameterize_packet(codec, Bytes::from(data), is_keyframe); let sequence = self.sequence.fetch_add(1, Ordering::Relaxed) + 1; return Ok(vec![EncodedVideoFrame { - data: Bytes::from(data), + data, pts_ms, is_keyframe, sequence, @@ -1275,14 +1368,15 @@ impl SharedVideoPipeline { let mut encoded_frames = Vec::with_capacity(frames.len()); for encoded in frames { - let is_keyframe = encoded.key == 1; + let (data, is_keyframe) = + self.inspect_and_parameterize_packet(codec, encoded.data, encoded.key == 1); let sequence = self.sequence.fetch_add(1, Ordering::Relaxed) + 1; if codec == VideoEncoderType::H264 { - self.update_h264_profile_level_id(&encoded.data); + self.update_h264_profile_level_id(&data); } encoded_frames.push(EncodedVideoFrame { - data: encoded.data, + data, pts_ms, is_keyframe, sequence, diff --git a/src/webrtc/universal_session.rs b/src/webrtc/universal_session.rs index 83b50a3b..6b4a8059 100644 --- a/src/webrtc/universal_session.rs +++ b/src/webrtc/universal_session.rs @@ -36,7 +36,7 @@ use crate::audio::OpusFrame; use crate::error::{AppError, Result}; use crate::hid::datachannel::{parse_hid_message, HidChannelEvent}; use crate::hid::HidController; -use crate::video::codec::h264_bitstream; +use crate::video::codec::{h264_bitstream, h265_bitstream}; use crate::video::types::{ BitratePreset, EncodedVideoFrame, PixelFormat, Resolution, VideoEncoderType, }; @@ -638,11 +638,16 @@ impl UniversalSession { next_keyframe_retry = Instant::now(); } - // Some H264 encoders output SPS/PPS in a separate non-keyframe AU - // before IDR. Keep this frame so browser can decode the next IDR. - let forward_h264_parameter_frame = waiting_for_keyframe - && expected_codec == VideoEncoderType::H264 - && h264_bitstream::has_sps_pps(encoded_frame.data.as_ref()); + let forward_parameter_frame = waiting_for_keyframe + && match expected_codec { + VideoEncoderType::H264 => h264_bitstream::has_sps_pps( + encoded_frame.data.as_ref(), + ), + VideoEncoderType::H265 => h265_bitstream::has_vps_sps_pps( + encoded_frame.data.as_ref(), + ), + _ => false, + }; let now = Instant::now(); if keyframe_requests < KEYFRAME_RETRY_LIMIT @@ -659,7 +664,7 @@ impl UniversalSession { ); } } - if !forward_h264_parameter_frame { + if !forward_parameter_frame { continue; } }