implemented automatic color conversion

Author: bwsw

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"

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<bool, String> {
@@ -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<PyObject> {
+        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<VideoFrameEnvelope>,
     signal: Arc<Mutex<bool>>,
     decode: bool,
+    autoconvert_raw_formats_to_rgb24: bool,
     log_level: Arc<Mutex<Option<Level>>>,
 ) {
     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<String, String>,
         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 {

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:

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

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: