implemented bsf as a pluggable system

Author: bwsw

Cargo.toml

@@ -22,7 +22,6 @@ crossbeam = "0.8"
 log = "0.4"
 env_logger = "0.11"
 parking_lot = "0.12"
-bitstream-io = "2.5"
 
 [dependencies.ffmpeg-next]
 version = "7"

src/lib.rs

@@ -1,12 +1,10 @@
 use anyhow::bail;
 use std::ffi::CString;
-use std::io::Cursor;
 use std::path::Path;
 use std::sync::Arc;
 use std::thread::{spawn, JoinHandle};
 use std::time::{Instant, SystemTime};
 
-use bitstream_io::{BigEndian, BitRead, BitReader};
 use crossbeam::channel::{Receiver, Sender};
 use derive_builder::Builder;
 use ffmpeg::util::frame::video::Video;
@@ -57,6 +55,27 @@ fn is_stream_key_framed(id: Id) -> Result<bool, String> {
     }
 }
 
+#[pyclass]
+#[derive(Clone, Debug)]
+pub struct BsfFilter {
+    codec: String,
+    name: String,
+    params: Vec<(String, String)>,
+}
+
+#[pymethods]
+impl BsfFilter {
+    #[new]
+    #[pyo3(signature = (codec, name, params = vec![]))]
+    fn new(codec: String, name: String, params: Vec<(String, String)>) -> Self {
+        Self {
+            codec,
+            name,
+            params,
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 #[pyclass]
 pub struct VideoFrameEnvelope {
@@ -179,6 +198,7 @@ struct HandleParams {
     autoconvert_raw_formats_to_rgb24: bool,
     block_if_queue_full: bool,
     log_level: Arc<Mutex<Option<Level>>>,
+    bsf_filters: Vec<BsfFilter>,
 }
 
 struct BitStreamFilterContext {
@@ -324,36 +344,27 @@ fn handle(params: HandleParams) -> anyhow::Result<()> {
 
     let mut video_filters = Vec::new();
 
-    match video_input.codec().id() {
-        Id::H264 => {
-            video_filters.push(init_bsf(
-                "h264_mp4toannexb",
-                &video_parameters,
-                time_base,
-                &[],
-            )?);
-            // video_filters.push(init_bsf(
-            //     "h264_metadata",
-            //     &video_parameters,
-            //     time_base,
-            //     &[("aud".into(), "insert".into())],
-            // )?);
-        }
-        Id::HEVC | Id::H265 => {
-            // video_filters.push(init_bsf(
-            //     "hevc_metadata",
-            //     &video_parameters,
-            //     time_base,
-            //     &[("aud".into(), "insert".into())],
-            // )?);
-            video_filters.push(init_bsf(
-                "hevc_mp4toannexb",
-                &video_parameters,
-                time_base,
-                &[],
-            )?);
+    let codec_name = video_input.codec().id().name();
+    for f in &params.bsf_filters {
+        if f.codec != codec_name {
+            info!(
+                "Skipping filter {} as it is not applicable to codec {}, must match {}",
+                f.name, codec_name, f.codec
+            );
+            continue;
         }
-        _ => {}
+
+        info!(
+            "Initializing filter: {} with parameters {:?}",
+            f.name, f.params
+        );
+
+        video_filters.push(init_bsf(
+            f.name.as_str(),
+            &video_parameters,
+            time_base,
+            &f.params,
+        )?);
     }
 
     let mut video_decoder =
@@ -594,7 +605,8 @@ impl FFMpegSource {
         autoconvert_raw_formats_to_rgb24 = false,
         block_if_queue_full = false,
         init_timeout_ms = 10000,
-        ffmpeg_log_level = FFmpegLogLevel::Info)
+        ffmpeg_log_level = FFmpegLogLevel::Info,
+        bsf_filters = vec![])
     )]
     pub fn new(
         uri: String,
@@ -605,6 +617,7 @@ impl FFMpegSource {
         block_if_queue_full: bool,
         init_timeout_ms: u64,
         ffmpeg_log_level: FFmpegLogLevel,
+        bsf_filters: Vec<BsfFilter>,
     ) -> PyResult<Self> {
         assert!(queue_len > 0, "Queue length must be a positive number");
 
@@ -627,6 +640,7 @@ impl FFMpegSource {
             .autoconvert_raw_formats_to_rgb24(autoconvert_raw_formats_to_rgb24)
             .block_if_queue_full(block_if_queue_full)
             .log_level(log_level.clone())
+            .bsf_filters(bsf_filters.clone())
             .build()
             .map_err(|e| {
                 error!("Unable to create handle params. Error is: {:?}", e);
@@ -690,35 +704,6 @@ impl FFMpegSource {
     }
 }
 
-// Function to decode Exp-Golomb codes from a slice and return a bit-string
-#[pyfunction]
-fn decode_exp_golomb(slice: &[u8]) -> String {
-    let mut reader = BitReader::endian(Cursor::new(slice), BigEndian);
-    let mut bit_string = String::new();
-
-    // Step 1: Count the number of leading zeros
-    let mut leading_zeros = 0;
-    while let Ok(bit) = reader.read_bit() {
-        bit_string.push(if bit { '1' } else { '0' }); // Append bit to bit-string
-        if !bit {
-            leading_zeros += 1;
-        } else {
-            break; // Stop at the first '1'
-        }
-    }
-
-    // Step 2: Read the remainder of the code
-    let mut code_value = 1; // The first '1' already encountered
-    for _ in 0..leading_zeros {
-        if let Ok(bit) = reader.read_bit() {
-            bit_string.push(if bit { '1' } else { '0' }); // Append to bit-string
-            code_value = (code_value << 1) | if bit { 1 } else { 0 };
-        }
-    }
-
-    bit_string
-}
-
 #[pymodule]
 fn ffmpeg_input(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     _ = env_logger::try_init_from_env("LOGLEVEL").map_err(|e| {
@@ -727,6 +712,6 @@ fn ffmpeg_input(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<VideoFrameEnvelope>()?;
     m.add_class::<FFMpegSource>()?;
     m.add_class::<FFmpegLogLevel>()?;
-    m.add_function(wrap_pyfunction!(decode_exp_golomb, m)?)?;
+    m.add_class::<BsfFilter>()?;
     Ok(())
 }

test.py

@@ -2,9 +2,7 @@
 
 os.environ["RUST_LOG"] = "info"
 
-from time import sleep
-
-from ffmpeg_input import FFMpegSource, FFmpegLogLevel, decode_exp_golomb
+from ffmpeg_input import FFMpegSource, FFmpegLogLevel, BsfFilter
 
 
 def bytes_to_bits_binary(byte_data):
@@ -14,14 +12,17 @@ def bytes_to_bits_binary(byte_data):
 
 if __name__ == '__main__':
     # set env LOGLEVEL=info
-    # file = "/home/ivan/Downloads/1_underground_supercut_reencode_bug_ab.mp4"
+    file = "/home/ivan/Downloads/1_underground_supercut_reencode_bug.mp4"
     # file = "/home/ivan/Downloads/1_underground_supercut.mp4"
-    file = "/home/ivan/Downloads/1_underground_supercut_reencode_bug_x265.mp4"
+    # file = "/home/ivan/Downloads/1_underground_supercut_reencode_bug_x265.mp4"
     # file = "/home/ivan/Downloads/1_underground_supercut_reencode_bug_aud.mp4"
     s = FFMpegSource(file, params=[],
                      queue_len=10, decode=False,
                      block_if_queue_full=True,
-                     ffmpeg_log_level=FFmpegLogLevel.Info)
+                     ffmpeg_log_level=FFmpegLogLevel.Info,
+                     bsf_filters=[BsfFilter("h264", "h264_mp4toannexb"),
+                                  BsfFilter("hevc", "hevc_mp4toannexb"),
+                                  BsfFilter("h265", "hevc_mp4toannexb")])
     s.log_level = FFmpegLogLevel.Info
     # counter = 0
     while True:
@@ -36,13 +37,13 @@ def bytes_to_bits_binary(byte_data):
             payload = p.payload_as_bytes()
             print("Payload length:", len(payload))
             # print 1st 3 bytes of the payload
-            bin_res = " ".join(format(x, '#010b')[2:] for x in payload[:16])
-            first_hex_res = " ".join(format(x, '02x') for x in payload[:16])
-            last_hex_res = " ".join(format(x, '02x') for x in payload[-4:])
-            code = decode_exp_golomb(payload[:16])
-            # int_val = int(code, 2)
-            print("Payload bin start:", code)
-            print("Payload hex start:", first_hex_res)
+            # bin_res = " ".join(format(x, '#010b')[2:] for x in payload[:16])
+            # first_hex_res = " ".join(format(x, '02x') for x in payload[:16])
+            # last_hex_res = " ".join(format(x, '02x') for x in payload[-4:])
+            # code = decode_exp_golomb(payload[:16])
+            # # int_val = int(code, 2)
+            # print("Payload bin start:", code)
+            # print("Payload hex start:", first_hex_res)
             # if len(payload) - 4 > int_val:
             #     first_hex_res = " ".join(format(x, '02x') for x in payload[4 + int_val:20 + int_val])
             #     print("Payload hex start:", first_hex_res)