python/st20p_rx: add more format support (#648)

move all cv2 code to cv2_util.py

Signed-off-by: Frank Du <[email protected]>

.gitignore

@@ -75,3 +75,4 @@ assets/
 *.gz
 *.tar
 mtl_system_status_*
+*.pyc

README.md

@@ -25,6 +25,7 @@ If you find value in our project, please consider giving it a star. Your support
 * Virtualization support by SR-IOV.
 * Built-in PTP protocol with hardware timestamp offload.
 * FFMPEG plugin, OBS(Open Broadcaster Software) plugin, and Intel® Media SDK support.
+* In addition to the native C/C++ API, it also offers bindings for [Python](python/README.md) and [Rust](rust/README.md).
 
 #### 1.1.1 ST2110 features
 

format-coding.sh

@@ -9,4 +9,7 @@
 set -e
 
 echo "clang-format check"
-find . -regex '.*\.\(cpp\|hpp\|cc\|c\|h\)' -exec clang-format --verbose -i {} \;
+find . -path ./build -prune -o -regex '.*\.\(cpp\|hpp\|cc\|c\|h\)' ! -name 'pymtl_wrap.c' -exec clang-format --verbose -i {} +
+
+black python/
+isort python/

include/st_pipeline_api.h

@@ -211,18 +211,21 @@ enum st_frame_fmt {
 /** ST format cap of ST_FRAME_FMT_H264_CBR_CODESTREAM, used in the st22_plugin caps */
 #define ST_FMT_CAP_H264_CBR_CODESTREAM (MTL_BIT64(ST_FRAME_FMT_H264_CBR_CODESTREAM))
 
-/**
- * Flag bit in flags of struct st_frame.
- * Frame has external buffer attached.
- */
-#define ST_FRAME_FLAG_EXT_BUF (MTL_BIT32(0))
+/** Flag bit in flags of struct st_frame. */
+enum st_frame_flag {
+  /** Frame has external buffer attached. */
+  ST_FRAME_FLAG_EXT_BUF = (MTL_BIT32(0)),
+  /** Frame planes data by single malloc */
+  ST_FRAME_FLAG_SINGLE_MALLOC = (MTL_BIT32(1)),
+  /** Frame planes data by rte_malloc */
+  ST_FRAME_FLAG_RTE_MALLOC = (MTL_BIT32(2)),
+};
 
 /** Max planes number for one frame */
 #define ST_MAX_PLANES (4)
 
 /** The structure info for external frame */
-struct st_ext_frame {
-  /** Each plane's virtual address of external frame */
+struct st_ext_frame { /** Each plane's virtual address of external frame */
   void* addr[ST_MAX_PLANES];
   /** Each plane's IOVA of external frame */
   mtl_iova_t iova[ST_MAX_PLANES];
@@ -1940,6 +1943,12 @@ static inline mtl_iova_t st_frame_iova(struct st_frame* frame, uint8_t plane) {
   return frame->iova[plane];
 }
 
+/** request to create a plained memory by rte malloc to hold the frame buffer */
+struct st_frame* st_frame_create(mtl_handle mt, enum st_frame_fmt fmt, uint32_t w,
+                                 uint32_t h, bool interlaced);
+/** free the frame created by st_frame_create */
+int st_frame_free(struct st_frame* frame);
+
 #if defined(__cplusplus)
 }
 #endif

lib/src/st2110/st_fmt.c

@@ -1061,6 +1061,48 @@ void st_frame_init_plane_single_src(struct st_frame* frame, void* addr, mtl_iova
   }
 }
 
+struct st_frame* st_frame_create(mtl_handle mt, enum st_frame_fmt fmt, uint32_t w,
+                                 uint32_t h, bool interlaced) {
+  struct mtl_main_impl* impl = mt;
+  int soc_id = mt_socket_id(impl, MTL_PORT_P);
+  struct st_frame* frame = mt_rte_zmalloc_socket(sizeof(*frame), soc_id);
+  if (!frame) {
+    err("%s, frame malloc fail\n", __func__);
+    return NULL;
+  }
+  frame->fmt = fmt;
+  frame->interlaced = interlaced;
+  frame->width = w;
+  frame->height = h;
+  frame->flags = ST_FRAME_FLAG_SINGLE_MALLOC | ST_FRAME_FLAG_RTE_MALLOC;
+
+  size_t data_sz = st_frame_size(fmt, w, h, interlaced);
+  void* data = mt_rte_zmalloc_socket(data_sz, soc_id);
+  if (!data) {
+    err("%s, data malloc fail, size %" PRIu64 "\n", __func__, data_sz);
+    st_frame_free(frame);
+    return NULL;
+  }
+  frame->buffer_size = data_sz;
+  frame->data_size = data_sz;
+  /* init plane */
+  st_frame_init_plane_single_src(frame, data, mtl_hp_virt2iova(impl, data));
+  return frame;
+}
+
+int st_frame_free(struct st_frame* frame) {
+  uint32_t reqiured_flags = ST_FRAME_FLAG_SINGLE_MALLOC | ST_FRAME_FLAG_RTE_MALLOC;
+  if (reqiured_flags != (frame->flags & reqiured_flags)) {
+    err("%s, frame %p is not created by st_frame_create\n", __func__, frame);
+    return -EINVAL;
+  }
+  if (frame->addr[0]) {
+    mt_rte_free(frame->addr[0]);
+  }
+  mt_rte_free(frame);
+  return 0;
+}
+
 /* the reference rl pad interval table for CVL NIC */
 struct cvl_pad_table {
   enum st20_fmt fmt;

python/README.md

@@ -2,6 +2,8 @@
 
 IMTL leverage SWIG, found at <https://github.com/swig/swig/tree/master>, to transform C APIs into a binding layer that Python can utilize.
 
+Before using the Python binding, please ensure that IMTL is [built](../doc/build.md) and the NIC is [set up](../doc/run.md) correctly.
+
 ## 1. Build and install swig
 
 Following the build and installation guide from the SWIG GitHub repository at <https://github.com/swig/swig/tree/master>.
@@ -58,6 +60,8 @@ Extracting pymtl-0.1-py3.10-linux-x86_64.egg to /usr/local/lib/python3.10/dist-p
 ## 3. Run python example code
 
 ```bash
-cd $imtl_source_code/python/example/
-python3 version.py
+cd $imtl_source_code/
+python3 python/example/version.py
+# Customize the port and IP in the code before using
+python3 python/example/st20p_rx.py
 ```

python/example/cv2_util.py

@@ -0,0 +1,108 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright 2023 Intel Corporation
+# opencv2 utils
+
+import ctypes
+
+import cv2
+import numpy as np
+import pymtl as mtl
+
+
+def yuv422p10le_to_yuv422(ptr, width, height):
+    y_size = width * height
+    u_size = width * height // 2
+
+    frame = np.frombuffer(ptr, dtype=np.uint16)
+
+    y = frame[:y_size].reshape((height, width))
+    u = frame[y_size : y_size + u_size].reshape((height, width // 2))
+    v = frame[y_size + u_size :].reshape((height, width // 2))
+
+    y = np.right_shift(y, 2).astype(np.uint8)
+    u = np.right_shift(u, 2).astype(np.uint8)
+    v = np.right_shift(v, 2).astype(np.uint8)
+
+    return y, u, v
+
+
+def downscale_yuv422(y, u, v, scale_factor):
+    height, width = y.shape
+
+    y_downscaled = cv2.resize(
+        y,
+        (width // scale_factor, height // scale_factor),
+        interpolation=cv2.INTER_LINEAR,
+    )
+    u_downscaled = cv2.resize(
+        u,
+        (width // (2 * scale_factor), height // scale_factor),
+        interpolation=cv2.INTER_LINEAR,
+    )
+    v_downscaled = cv2.resize(
+        v,
+        (width // (2 * scale_factor), height // scale_factor),
+        interpolation=cv2.INTER_LINEAR,
+    )
+
+    return y_downscaled, u_downscaled, v_downscaled
+
+
+def display_yuv422(y, u, v):
+    u = cv2.resize(u, (y.shape[1], y.shape[0]), interpolation=cv2.INTER_LINEAR)
+    v = cv2.resize(v, (y.shape[1], y.shape[0]), interpolation=cv2.INTER_LINEAR)
+
+    yuv = cv2.merge([y, u, v])
+    display_img = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR)
+    cv2.imshow("st20p_rx_python", display_img)
+    cv2.waitKey(1)
+
+
+def destroy():
+    cv2.destroyAllWindows()
+
+
+def frame_display_yuv422p10le(frame, display_scale_factor):
+    # Pack frame pointer from frame addr
+    ptr = (ctypes.c_ubyte * (frame.data_size)).from_address(mtl.st_frame_addr(frame, 0))
+    width = frame.width
+    height = frame.height
+
+    # It's slow since it include many Converting steps, just demo the usage
+    # Convert to yuv422 planar 8
+    y, u, v = yuv422p10le_to_yuv422(ptr, width, height)
+    # Downscale
+    y_scale, u_scale, v_scale = downscale_yuv422(y, u, v, display_scale_factor)
+    # Convert to RGB and display
+    display_yuv422(y_scale, u_scale, v_scale)
+
+
+def frame_display_uyvy(frame, display_scale_factor):
+    # Pack frame pointer from frame addr
+    ptr = (ctypes.c_ubyte * (frame.data_size)).from_address(mtl.st_frame_addr(frame, 0))
+    width = frame.width
+    height = frame.height
+    downscaled_width = width // display_scale_factor
+    downscaled_height = height // display_scale_factor
+
+    uyvy = np.frombuffer(ptr, dtype=np.uint8).reshape((height, width, 2))
+    rgb = cv2.cvtColor(uyvy, cv2.COLOR_YUV2BGR_UYVY)
+    downscaled_rgb = cv2.resize(
+        rgb, (downscaled_width, downscaled_height), interpolation=cv2.INTER_AREA
+    )
+    cv2.imshow("st20p_rx_python", downscaled_rgb)
+    cv2.waitKey(1)
+
+
+def frame_display_rfc4175be10(mtl_handle, frame, display_scale_factor):
+    yuv422p10le_frame = mtl.st_frame_create(
+        mtl_handle,
+        mtl.ST_FRAME_FMT_YUV422PLANAR10LE,
+        frame.width,
+        frame.height,
+        frame.interlaced,
+    )
+    if yuv422p10le_frame:
+        mtl.st_frame_convert(frame, yuv422p10le_frame)
+        frame_display_yuv422p10le(yuv422p10le_frame, display_scale_factor)
+        mtl.st_frame_free(yuv422p10le_frame)

python/example/st20p_rx.py

@@ -1,64 +1,17 @@
 # SPDX-License-Identifier: BSD-3-Clause
 # Copyright 2023 Intel Corporation
 
-import ctypes
+from pickle import TRUE
 
-import cv2
-import numpy as np
+import cv2_util
 import pymtl as mtl
 
 
-def yuv422p10le_to_yuv422(ptr, width, height):
-    y_size = width * height
-    u_size = width * height // 2
-
-    frame = np.frombuffer(ptr, dtype=np.uint16)
-
-    y = frame[:y_size].reshape((height, width))
-    u = frame[y_size : y_size + u_size].reshape((height, width // 2))
-    v = frame[y_size + u_size :].reshape((height, width // 2))
-
-    y = np.right_shift(y, 2).astype(np.uint8)
-    u = np.right_shift(u, 2).astype(np.uint8)
-    v = np.right_shift(v, 2).astype(np.uint8)
-
-    return y, u, v
-
-
-def downscale_yuv422(y, u, v, scale_factor):
-    height, width = y.shape
-
-    y_downscaled = cv2.resize(
-        y,
-        (width // scale_factor, height // scale_factor),
-        interpolation=cv2.INTER_LINEAR,
-    )
-    u_downscaled = cv2.resize(
-        u,
-        (width // (2 * scale_factor), height // scale_factor),
-        interpolation=cv2.INTER_LINEAR,
-    )
-    v_downscaled = cv2.resize(
-        v,
-        (width // (2 * scale_factor), height // scale_factor),
-        interpolation=cv2.INTER_LINEAR,
-    )
-
-    return y_downscaled, u_downscaled, v_downscaled
-
-
-def display_yuv422(y, u, v):
-    u = cv2.resize(u, (y.shape[1], y.shape[0]), interpolation=cv2.INTER_LINEAR)
-    v = cv2.resize(v, (y.shape[1], y.shape[0]), interpolation=cv2.INTER_LINEAR)
-
-    yuv = cv2.merge([y, u, v])
-    display_img = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR)
-    cv2.imshow("st20p_rx_python", display_img)
-    cv2.waitKey(1)
-
-
 def main():
+    display = TRUE
     display_scale_factor = 2
+    # mtl.ST_FRAME_FMT_YUV422PLANAR10LE or mtl.ST_FRAME_FMT_UYVY or mtl.ST_FRAME_FMT_YUV422RFC4175PG2BE10
+    output_fmt = mtl.ST_FRAME_FMT_YUV422PLANAR10LE
 
     # Init para
     init_para = mtl.mtl_init_params()
@@ -80,7 +33,7 @@ def main():
     rx_para.fps = mtl.ST_FPS_P59_94
     rx_para.framebuff_cnt = 3
     rx_para.transport_fmt = mtl.ST20_FMT_YUV_422_10BIT
-    rx_para.output_fmt = mtl.ST_FRAME_FMT_YUV422PLANAR10LE
+    rx_para.output_fmt = output_fmt
     # rx port
     rx_port = mtl.st_rx_port()
     mtl.st_rxp_para_port_set(rx_port, mtl.MTL_SESSION_PORT_P, "0000:af:01.0")
@@ -100,31 +53,26 @@ def main():
                 # print(f"frame addr: {hex(mtl.st_frame_addr(frame, 0))}")
                 # print(f"frame iova: {hex(mtl.st_frame_iova(frame, 0))}")
                 # print(f"pkts_total: {frame.pkts_total}")
-
-                # Pack frame pointer from frame addr
-                ptr = (ctypes.c_ubyte * (frame.data_size)).from_address(
-                    mtl.st_frame_addr(frame, 0)
-                )
-                width = frame.width
-                height = frame.height
-
-                # It's slow since it include many Converting steps, just demo the usage
-                # Convert to yuv422 planar 8
-                y, u, v = yuv422p10le_to_yuv422(ptr, width, height)
-                # Downscale
-                y_scale, u_scale, v_scale = downscale_yuv422(
-                    y, u, v, display_scale_factor
-                )
-                # Convert to RGB and display
-                display_yuv422(y_scale, u_scale, v_scale)
-
+                if display:
+                    if output_fmt == mtl.ST_FRAME_FMT_YUV422PLANAR10LE:
+                        cv2_util.frame_display_yuv422p10le(frame, display_scale_factor)
+                    elif output_fmt == mtl.ST_FRAME_FMT_UYVY:
+                        cv2_util.frame_display_uyvy(frame, display_scale_factor)
+                    elif output_fmt == mtl.ST_FRAME_FMT_YUV422RFC4175PG2BE10:
+                        cv2_util.frame_display_rfc4175be10(
+                            mtl_handle, frame, display_scale_factor
+                        )
+                    else:
+                        print(
+                            f"Unknown output_fmt: {mtl.st_frame_fmt_name(output_fmt)}"
+                        )
                 # return the frame
                 mtl.st20p_rx_put_frame(st20p_rx, frame)
 
     except KeyboardInterrupt:
         print("KeyboardInterrupt")
 
-    cv2.destroyAllWindows()
+    cv2_util.destroy()
 
     # Free st20p_rx session
     mtl.st20p_rx_free(st20p_rx)