Process some problems about buffer size in GStreamer

[Hùng Triệu Việt]

Hello guys, I'm new to GStreamer. I have a Python code about object
detection app using PyTorch, OpenCV to draw bounding boxes and GStreamer to
process multiple RTSP links. In this pipeline I'm using a video file to
demo:
filesrc ---> decodebin ---> videorate ---> videoconvert ---> queue --->
autovideosink
I have a problem, when I implement a simple code to save the output image,
the expected image must be an image with a frame cutted from input video
and have bounding boxes on that. Instead of that, I have an image with a
series of frames, like a flow of frames cut from input video with bounding
boxes on it (check this image: https://imgur.com/a/QKr4tN2).
Now I have no idea what to do, can someone give me advice on this?

This is my code:
import gi
gi.require_version('Gst', '1.0')
gi.require_version('GObject', '2.0')
from gi.repository import Gst
import numpy as np
import torch
from PIL import Image
import cv2
import time

Gst.init(None)

src = Gst.ElementFactory.make('filesrc', 'source')
src.set_property('location', 'video/test.mp4')
src.set_property('num-buffers', 9000)
decodebin = Gst.ElementFactory.make('decodebin', 'decoder')
videorate = Gst.ElementFactory.make('videorate', 'rate')
videorate.set_property('max-rate', 1)
videoconvert = Gst.ElementFactory.make('videoconvert', 'converter')
queue = Gst.ElementFactory.make('queue', 'queue')
autovideosink = Gst.ElementFactory.make('autovideosink', 's')

pipeline = Gst.Pipeline.new('new-pipeline')

pipeline.add(src)
pipeline.add(decodebin)
pipeline.add(videorate)
pipeline.add(videoconvert)
pipeline.add(queue)
pipeline.add(autovideosink)

def on_pad_added(element, pad):
    sinkpad = videoconvert.get_static_pad('sink')
    pad.link(sinkpad)

    decodebin.link(videorate)
    videorate.link(videoconvert)
    videoconvert.link(queue)
    queue.link(autovideosink)

    sinkpad = autovideosink.get_static_pad('sink')
    videoconvert.link(queue)

src.link(decodebin)
decodebin.connect('pad-added', on_pad_added)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
detector = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.pt
').eval().to(device)

fps = 0
prev_time = time.time()

def on_frame_probe(pad, info):
    global fps, prev_time
    buf = info.get_buffer()
    # print(buf.get_size())
    print(f'[{buf.pts / Gst.SECOND:6.2f}]')
    image_tensor = buffer_to_image_tensor(buf, pad.get_current_caps())
    with torch.no_grad():
        detections = detector(image_tensor)
        current_time = time.time()
        fps = 1 / (current_time - prev_time)
        print("FPS: {0:.2f}".format(fps))
        prev_time = current_time
        print(detections)
        objects = (detections.xyxy[0]).tolist()
        img_np = np.array(image_tensor)
        for i in range(len(objects)):
            x_min, y_min, x_max, y_max = int(objects[i][0]),
int(objects[i][1]), int(objects[i][2]), int(objects[i][3])
            cv2.rectangle(img_np, (x_min, y_min), (x_max, y_max),
(255,0,0), 2)
        # Save the output image
        Image.fromarray(img_np[:,:,:3]).save("output.jpg")
    return Gst.PadProbeReturn.OK


def buffer_to_image_tensor(buf, caps):
    global fps
    caps_structure = caps.get_structure(0)
    height, width = caps_structure.get_value('height'),
caps_structure.get_value('width')
    channels = 3
    is_mapped, map_info = buf.map(Gst.MapFlags.READ)
    if is_mapped:
        try:
            image_array = np.frombuffer(map_info.data,
dtype=np.uint8).reshape((640, 1637, channels)).copy()
            image_array.resize((height, width, channels))
            image_array = cv2.cvtColor(image_array, cv2.COLOR_BGR2RGB)
            image_array = cv2.resize(image_array, (1920, 1080))
            return Image.fromarray(image_array) # RGBA -> RGB
        finally:
            buf.unmap(map_info)

pipeline.get_by_name('s').get_static_pad('sink').add_probe(
    Gst.PadProbeType.BUFFER,
    on_frame_probe
)

pipeline.set_state(Gst.State.PLAYING)

while True:
    msg = pipeline.get_bus().timed_pop_filtered(
        Gst.SECOND,
        Gst.MessageType.EOS | Gst.MessageType.ERROR
    )
    if msg:
        text = msg.get_structure().to_string() if msg.get_structure() else
''
        msg_type = Gst.message_type_get_name(msg.type)
        print(f'{msg.src.name}: [{msg_type}] {text}')
        break

pipeline.set_state(Gst.State.NULL)

The input video i got from this youtube link:
https://www.youtube.com/watch?v=MNn9qKG2UFI&list=PLcQZGj9lFR7y5WikozDSrdk6UCtAnM9mB
I appreciate you so much.
-- 
Trân trọng,

Triệu Việt Hùng
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <https://lists.freedesktop.org/archives/gstreamer-devel/attachments/20230228/8f5e1ca5/attachment.htm>