python/example: add encode/decode support for rx/tx (#672)

python/example: add encode/decode support for rx/tx

test with:
python3 python/example/st20p_tx.py
python3 python/example/st20p_rx_encode.py

TX decode YUV video from an encoded file and transmit it as a ST2110
ST_FRAME_FMT_YUV422RFC4175PG2BE10 stream across the network.

RX receive a ST2110 ST_FRAME_FMT_YUV422RFC4175PG2BE10 stream and encode
the video stream to a `.mp4` encoder file.

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

.gitignore

@@ -76,3 +76,4 @@ assets/
 *.tar
 mtl_system_status_*
 *.pyc
+*.mp4

python/README.md

@@ -75,21 +75,44 @@ Extracting pymtl-0.1-py3.10-linux-x86_64.egg to /usr/local/lib/python3.10/dist-p
 Install `opencv-python` dependency:
 
 ```bash
+# for yuv display
 sudo pip3 install opencv-python
+# PyAv for video decode/encode
+sudo pip3 install av
 ```
 
-Run example: st20p_rx.py
+### 3.1 st20p_rx.py
+
+Execute the `st20p_rx.py` to receive a ST2110 ST_FRAME_FMT_YUV422RFC4175PG2BE10 stream and display it.
 
 ```bash
-cd $imtl_source_code/
 # Customize the port, IP and display option in the code before using
 python3 python/example/st20p_rx.py
 ```
 
-Run example: st20p_tx.py
+### 3.2 st20p_tx.py
+
+Run the `st20p_tx.py`, which reads YUV video data from a file and transmits it over the network as a ST2110 ST_FRAME_FMT_YUV422RFC4175PG2BE10 stream.
 
 ```bash
-cd $imtl_source_code/
 # Customize the port, IP and display option in the code before using
 python3 python/example/st20p_tx.py
 ```
+
+### 3.3 st20p_rx_encode.py
+
+Run the `st20p_rx_encode.py` to receive a ST2110 ST_FRAME_FMT_YUV422RFC4175PG2BE10 stream and encode it to a `.mp4` encoder file.
+
+```bash
+# Customize the port, IP and display option in the code before using
+python3 python/example/st20p_rx_encode.py
+```
+
+### 3.4 st20p_tx_decode.py
+
+Use `st20p_tx_decode.py` to decode YUV video from an encoded file `jellyfish-3-mbps-hd-hevc-10bit.mkv` and transmit it as a ST2110 ST_FRAME_FMT_YUV422RFC4175PG2BE10 stream across the network.
+
+```bash
+# Customize the port, IP and display option in the code before using
+python3 python/example/st20p_tx_decode.py
+```

python/example/st20p_rx_encode.py

@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright 2023 Intel Corporation
+
+import ctypes
+import sys
+from datetime import datetime
+
+import av
+import numpy as np
+import pymtl as mtl
+
+
+def main():
+    width = 1920
+    height = 1080
+    mtl_output_fmt = mtl.ST_FRAME_FMT_YUV422PLANAR8
+    av_pixel_input_format = "yuv422p"
+
+    current_datetime = datetime.now()
+    output_file = current_datetime.strftime("%Y_%m_%d_%H_%M_%S") + "_output.mp4"
+
+    # Init para
+    init_para = mtl.mtl_init_params()
+    mtl.mtl_para_port_set(init_para, mtl.MTL_PORT_P, "0000:af:01.0")
+    init_para.num_ports = 1
+    mtl.mtl_para_sip_set(init_para, mtl.MTL_PORT_P, "192.168.108.102")
+    init_para.flags = mtl.MTL_FLAG_BIND_NUMA | mtl.MTL_FLAG_DEV_AUTO_START_STOP
+    mtl.mtl_para_tx_queues_cnt_set(init_para, mtl.MTL_PORT_P, 0)
+    mtl.mtl_para_rx_queues_cnt_set(init_para, mtl.MTL_PORT_P, 1)
+
+    # Create MTL instance
+    mtl_handle = mtl.mtl_init(init_para)
+    if not mtl_handle:
+        print("mtl_init fail")
+        sys.exit(1)
+
+    # Create st20p rx session
+    rx_para = mtl.st20p_rx_ops()
+    rx_para.name = "st20p_rx_python"
+    rx_para.width = width
+    rx_para.height = height
+    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_output_fmt
+    # rx port
+    rx_port = mtl.st_rx_port()
+    mtl.st_rxp_para_port_set(
+        rx_port,
+        mtl.MTL_SESSION_PORT_P,
+        mtl.mtl_para_port_get(init_para, mtl.MTL_SESSION_PORT_P),
+    )
+    rx_port.num_port = 1
+    mtl.st_rxp_para_sip_set(rx_port, mtl.MTL_SESSION_PORT_P, "239.168.85.20")
+    mtl.st_rxp_para_udp_port_set(rx_port, mtl.MTL_SESSION_PORT_P, 20000)
+    rx_port.payload_type = 112
+    rx_para.port = rx_port
+    # enable block get mode
+    rx_para.flags = mtl.ST20P_RX_FLAG_BLOCK_GET
+    # create st20p_rx session
+    st20p_rx = mtl.st20p_rx_create(mtl_handle, rx_para)
+    if not st20p_rx:
+        print("st20p_rx_create fail")
+        sys.exit(1)
+
+    # create h264_stream
+    h264 = av.open(output_file, mode="w")
+    h264_stream = h264.add_stream("libx264", rate=60)
+    h264_stream.width = width
+    h264_stream.height = height
+    h264_stream.pix_fmt = "yuv420p"  # h264 use yuv420p
+
+    # loop until ctrl-c
+    try:
+        while True:
+            frame = mtl.st20p_rx_get_frame(st20p_rx)
+            if not frame:
+                continue
+
+            video_frame = av.VideoFrame(width, height, av_pixel_input_format)
+            for plane in range(mtl.st_frame_fmt_planes(frame.fmt)):
+                p_size = mtl.st_frame_plane_size(frame, plane)
+                # print(f"plane: {plane} size: {p_size}")
+                ptr = (ctypes.c_ubyte * p_size).from_address(
+                    mtl.st_frame_addr_cpuva(frame, plane)
+                )
+                y = np.ctypeslib.as_array(ptr, (p_size,))
+                video_frame.planes[plane].update(y)
+
+            for packet in h264_stream.encode(video_frame):
+                h264.mux(packet)
+
+            # return the frame
+            mtl.st20p_rx_put_frame(st20p_rx, frame)
+
+    except KeyboardInterrupt:
+        print("KeyboardInterrupt")
+
+    # write eof stream packet
+    for packet in h264_stream.encode():
+        h264.mux(packet)
+
+    # Free st20p_rx session
+    mtl.st20p_rx_free(st20p_rx)
+
+    # Free MTL instance
+    mtl.mtl_uninit(mtl_handle)
+
+    print(f"Everything fine with {output_file}, bye")
+
+
+if __name__ == "__main__":
+    main()

python/example/st20p_tx.py

@@ -16,6 +16,8 @@ def main():
     # yuv422p10le_1080p.yuv, yuv422rfc4175be10_1080p.yuv
     # yuv422p10le_1080i.yuv, yuv422rfc4175be10_1080i.yuv
     yuv_file_path = "yuv422p10le_1080p.yuv"
+    width = 1920
+    height = 1080
     interlaced = False
 
     yuv_file = open(yuv_file_path, "rb")
@@ -45,8 +47,8 @@ def main():
     # Create st20p tx session
     tx_para = mtl.st20p_tx_ops()
     tx_para.name = "st20p_tx_python"
-    tx_para.width = 1920
-    tx_para.height = 1080
+    tx_para.width = width
+    tx_para.height = height
     tx_para.fps = mtl.ST_FPS_P59_94
     tx_para.interlaced = interlaced
     tx_para.framebuff_cnt = 3

python/example/st20p_tx_decode.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright 2023 Intel Corporation
+
+import ctypes
+import sys
+
+import av
+import numpy as np
+import pymtl as mtl
+
+
+def process_frame(mtl_handle, st20p_tx, frame):
+    # print(f"{frame.format.name}");
+    # convert to yuv422p10le from yuv420p10le
+    yuv_frame = frame.reformat(format="yuv422p10le")
+    tx_frame = mtl.st20p_tx_get_frame(st20p_tx)
+    if not tx_frame:
+        print("st20p_tx_get_frame fail, skip this video frame")
+        pass
+
+    width = frame.width
+    height = frame.height
+    y_size = width * height
+    u_size = y_size // 2
+    v_size = u_size
+    yuv_array = np.empty(y_size + u_size + v_size, dtype=np.uint16)
+
+    # pyav yuv422p10le not support Conversion to numpy array, use copy mode
+    yuv_array[:y_size] = np.frombuffer(yuv_frame.planes[0], np.uint16)
+    yuv_array[y_size : y_size + u_size] = np.frombuffer(yuv_frame.planes[1], np.uint16)
+    yuv_array[y_size + u_size :] = np.frombuffer(yuv_frame.planes[2], np.uint16)
+    src_p = ctypes.c_char_p(yuv_array.ctypes.data)
+    src_address = ctypes.cast(src_p, ctypes.c_void_p).value
+    src_address_uint64 = ctypes.c_uint64(src_address).value
+
+    memcpy_ops = mtl.mtl_memcpy_ops()
+    memcpy_ops.dst = mtl.st_frame_addr_cpuva(tx_frame, 0)
+    memcpy_ops.src = src_address_uint64
+    memcpy_ops.sz = tx_frame.data_size
+    mtl.mtl_memcpy_action(memcpy_ops)
+
+    mtl.st20p_tx_put_frame(st20p_tx, tx_frame)
+
+
+def get_video_resolution(video_path):
+    with av.open(video_path) as container:
+        video_stream = next(s for s in container.streams if s.type == "video")
+        width = video_stream.width
+        height = video_stream.height
+        return (width, height)
+
+
+def main():
+    # jellyfish-3-mbps-hd-hevc-10bit.mkv or output.mp4
+    input_file_path = "jellyfish-3-mbps-hd-hevc-10bit.mkv"
+    width, height = get_video_resolution(input_file_path)
+    print(f"input_file_path: {input_file_path}, width: {width}, height: {height}")
+    input_fmt = mtl.ST_FRAME_FMT_YUV422PLANAR10LE
+
+    # Init para
+    init_para = mtl.mtl_init_params()
+    mtl.mtl_para_port_set(init_para, mtl.MTL_PORT_P, "0000:af:01.1")
+    init_para.num_ports = 1
+    mtl.mtl_para_sip_set(init_para, mtl.MTL_PORT_P, "192.168.108.101")
+    init_para.flags = (
+        mtl.MTL_FLAG_BIND_NUMA
+        | mtl.MTL_FLAG_DEV_AUTO_START_STOP
+        | mtl.MTL_FLAG_PTP_ENABLE
+    )
+    mtl.mtl_para_tx_queues_cnt_set(init_para, mtl.MTL_PORT_P, 1)
+    mtl.mtl_para_rx_queues_cnt_set(init_para, mtl.MTL_PORT_P, 0)
+
+    # Create MTL instance
+    mtl_handle = mtl.mtl_init(init_para)
+    if not mtl_handle:
+        print("mtl_init fail")
+        sys.exit(1)
+
+    # Create st20p tx session
+    tx_para = mtl.st20p_tx_ops()
+    tx_para.name = "st20p_tx_python"
+    tx_para.width = width
+    tx_para.height = height
+    tx_para.fps = mtl.ST_FPS_P59_94
+    tx_para.framebuff_cnt = 3
+    tx_para.transport_fmt = mtl.ST20_FMT_YUV_422_10BIT
+    tx_para.input_fmt = input_fmt
+    # tx port
+    tx_port = mtl.st_tx_port()
+    mtl.st_txp_para_port_set(
+        tx_port,
+        mtl.MTL_SESSION_PORT_P,
+        mtl.mtl_para_port_get(init_para, mtl.MTL_SESSION_PORT_P),
+    )
+    tx_port.num_port = 1
+    mtl.st_txp_para_dip_set(tx_port, mtl.MTL_SESSION_PORT_P, "239.168.85.20")
+    mtl.st_txp_para_udp_port_set(tx_port, mtl.MTL_SESSION_PORT_P, 20000)
+    tx_port.payload_type = 112
+    tx_para.port = tx_port
+    # enable block get mode
+    tx_para.flags = mtl.ST20P_TX_FLAG_BLOCK_GET
+    # create st20p_tx session
+    st20p_tx = mtl.st20p_tx_create(mtl_handle, tx_para)
+    if not st20p_tx:
+        print("st20p_tx_create fail")
+        sys.exit(1)
+    frame_sz = mtl.st20p_tx_frame_size(st20p_tx)
+    print(f"frame_sz: {hex(frame_sz)}")
+
+    # loop until ctrl-c
+    container = av.open(input_file_path)
+    try:
+        stream = next(s for s in container.streams if s.type == "video")
+        while True:
+            for frame in container.decode(stream):
+                process_frame(mtl_handle, st20p_tx, frame)
+
+            # seek to the first frame
+            print("Finish play, seek to first frame")
+            container.seek(0, stream=stream)
+
+    except KeyboardInterrupt:
+        print("KeyboardInterrupt")
+
+    finally:
+        container.close()
+
+    print("Everything fine, bye")
+
+    # Free st20p_tx session
+    mtl.st20p_tx_free(st20p_tx)
+
+    # Free MTL instance
+    mtl.mtl_uninit(mtl_handle)
+
+    print("Everything fine, bye")
+
+
+if __name__ == "__main__":
+    main()