diff --git a/Cargo.toml b/Cargo.toml index ec1c262..5439ca7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,7 +7,7 @@ repository = "https://github.com/insight-platform/FFmpeg-Input" readme = "README.md" keywords = ["FFmpeg", "Video"] categories = ["computer-vision"] -version = "0.1.22" +version = "0.1.23" edition = "2021" license="Apache-2.0" rust-version = "1.62" diff --git a/src/lib.rs b/src/lib.rs index b5045b9..8a44936 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,9 @@ +use std::collections::HashMap; +use std::path::Path; +use std::sync::Arc; +use std::thread::{spawn, JoinHandle}; +use std::time::SystemTime; + use crossbeam::channel::{Receiver, Sender}; use ffmpeg::util::frame::video::Video; use ffmpeg_next as ffmpeg; @@ -10,13 +16,8 @@ use parking_lot::Mutex; use pyo3::exceptions::PyBrokenPipeError; use pyo3::prelude::*; use pyo3::types::PyBytes; -use std::collections::HashMap; -use std::path::Path; -use std::sync::Arc; -use std::thread::{spawn, JoinHandle}; -use std::time::SystemTime; -const DECODING_FORMAT: Pixel = Pixel::BGR24; +const DECODING_FORMAT: Pixel = Pixel::RGB24; const DECODED_PIX_BYTES: u32 = 3; fn is_stream_key_framed(id: Id) -> Result { @@ -108,8 +109,12 @@ impl VideoFrameEnvelope { self.__repr__() } - fn payload_as_bytes(&self, py: Python) -> PyObject { - PyBytes::new(py, &self.payload).into() + fn payload_as_bytes(&self, py: Python) -> PyResult { + let res = PyBytes::new_with(py, self.payload.len(), |b: &mut [u8]| { + b.copy_from_slice(&self.payload); + Ok(()) + })?; + Ok(res.into()) } } @@ -139,6 +144,7 @@ fn handle( tx: Sender, signal: Arc>, decode: bool, + autoconvert_raw_formats_to_rgb24: bool, log_level: Arc>>, ) { let mut queue_full_skipped_count = 0; @@ -181,7 +187,7 @@ fn handle( // video_decoder.format(), // video_decoder.width(), // video_decoder.height(), - // Pixel::BGR24, + // Pixel::rgb24, // video_decoder.width(), // video_decoder.height(), // Flags::FAST_BILINEAR, @@ -256,6 +262,10 @@ fn handle( continue; } + let decode = decode + || (autoconvert_raw_formats_to_rgb24 + && video_decoder.codec().map(|c| c.id()) == Some(Id::RAWVIDEO)); + let raw_frames = if decode { let mut raw_frames = Vec::new(); video_decoder @@ -364,12 +374,18 @@ fn assign_log_level(ffmpeg_log_level: FFmpegLogLevel) -> Level { #[pymethods] impl FFMpegSource { #[new] - #[pyo3(signature = (uri, params, queue_len = 32, decode = false, ffmpeg_log_level = FFmpegLogLevel::Info))] + #[pyo3(signature = (uri, params, + queue_len = 32, + decode = false, + autoconvert_raw_formats_to_rgb24 = false, + ffmpeg_log_level = FFmpegLogLevel::Info) + )] pub fn new( uri: String, params: HashMap, queue_len: i64, decode: bool, + autoconvert_raw_formats_to_rgb24: bool, ffmpeg_log_level: FFmpegLogLevel, ) -> Self { assert!(queue_len > 0, "Queue length must be a positive number"); @@ -383,7 +399,15 @@ impl FFMpegSource { let thread_exit_signal = exit_signal.clone(); let thread_ll = log_level.clone(); let thread = Some(spawn(move || { - handle(uri, params, tx, thread_exit_signal, decode, thread_ll) + handle( + uri, + params, + tx, + thread_exit_signal, + decode, + autoconvert_raw_formats_to_rgb24, + thread_ll, + ) })); Self { diff --git a/test.py b/test.py index 910395d..655a380 100644 --- a/test.py +++ b/test.py @@ -2,7 +2,7 @@ if __name__ == '__main__': s = FFMpegSource("/home/ivan/road-traffic-processed.mp4", params={"fflags": "+genpts"}, - queue_len=100, decode=True, + queue_len=100, decode=False, ffmpeg_log_level=FFmpegLogLevel.Debug) s.log_level = FFmpegLogLevel.Trace while True: diff --git a/test_autoconvert.py b/test_autoconvert.py new file mode 100644 index 0000000..2e24234 --- /dev/null +++ b/test_autoconvert.py @@ -0,0 +1,29 @@ +from ffmpeg_input import FFMpegSource, FFmpegLogLevel +import numpy as np +import cv2 +import time + +if __name__ == '__main__': + s = FFMpegSource("/dev/video0", + params={"video_size": "1920x1080", "c:v": "v4l2m2m", "input_format": "yuyv422"}, + queue_len=100, + autoconvert_raw_formats_to_rgb24=True, + ffmpeg_log_level=FFmpegLogLevel.Info) + s.log_level = FFmpegLogLevel.Panic + while True: + try: + p = s.video_frame() + res = p.payload_as_bytes() + # 1944 2592 + #print(p.frame_height, p.frame_width) + #res = np.frombuffer(res, dtype=np.uint8) + #res = np.reshape(res, (p.frame_height, p.frame_width, 3)) + end = time.time() + print(p.codec, p.pixel_format, p.queue_len, "all_time={}".format(int(end * 1000 - p.frame_received_ts)), + "python_time={}".format(int(end * 1000 - p.frame_processed_ts))) + # cv2.imshow('Image', res) + #if cv2.waitKey(1) & 0xFF == ord('q'): + # break + except BrokenPipeError: + print("EOS") + break diff --git a/test_decode.py b/test_decode.py index 32d59df..8c570b9 100644 --- a/test_decode.py +++ b/test_decode.py @@ -7,7 +7,7 @@ s = FFMpegSource("/dev/video0", params={"video_size": "1920x1080", "c:v": "v4l2m2m", "input_format": "mjpeg"}, queue_len=100, - decode=False, + decode=True, ffmpeg_log_level=FFmpegLogLevel.Info) s.log_level = FFmpegLogLevel.Panic while True: