def pad_buffer_probe()

in awstreamer/gst_pipeline/ds_pipeline.py [0:0]

• 48 lines of code
• 10 McCabe index (conditional complexity)

    def pad_buffer_probe(self, pad, info, func):
        '''
        Parse DeepStream metadata and send it over to the callback function.
        Probe should be attached to either src or sink of plug-ins after nvinfer in the pipeline.
        '''
        # Get the gst buffer
        gst_buffer = info.get_buffer()
        if not gst_buffer:
            logger.error("Unable to get GstBuffer ")
            return Gst.PadProbeReturn.OK

        # Retrieve batch metadata from the gst_buffer
        # Note that pyds.gst_buffer_get_nvds_batch_meta() expects the
        # C address of gst_buffer as input, which is obtained with hash(gst_buffer)
        batch_meta = pyds.gst_buffer_get_nvds_batch_meta(hash(gst_buffer))
        frame_meta_list = batch_meta.frame_meta_list
        while frame_meta_list is not None:
            try:
                # Note that frame_meta_list.data needs a cast to pyds.NvDsFrameMeta
                # The casting is done by pyds.glist_get_nvds_frame_meta()
                # The casting also keeps ownership of the underlying memory
                # in the C code, so the Python garbage collector will leave
                # it alone.
                frame_meta = pyds.NvDsFrameMeta.cast(frame_meta_list.data)
            except StopIteration:
                break

            frame_number = frame_meta.frame_num
            num_rects = frame_meta.num_obj_meta
            obj_meta_list = frame_meta.obj_meta_list
            object_list = []
            while obj_meta_list is not None:
                try:
                    obj_meta = pyds.NvDsObjectMeta.cast(obj_meta_list.data)
                except StopIteration:
                    break

                # Add object info to the list
                bbox = obj_meta.rect_params
                object_info = {
                    "class_id": obj_meta.class_id,
                    "confidence": obj_meta.confidence,
                    "bbox": [bbox.left, bbox.top, bbox.width, bbox.height]
                }
                object_list.append(object_info)

                # Move to the next object
                try:
                    obj_meta_list = obj_meta_list.next
                except StopIteration:
                    break

            # Collect metadata and call callback function
            if len(object_list) > 0:
                metadata = {
                    "frame_number": frame_number,
                    "num_rects": num_rects,
                    "objects": object_list
                }
                func(metadata)

            # Move to the next frame
            try:
                frame_meta_list = frame_meta_list.next
            except StopIteration:
                break

        return Gst.PadProbeReturn.OK