How to add custom labels in python code to extract metadata

Hi, I Hope you are doing well.
i have a custom label for my yolov4 model
which are as under and the order is the same like we have in labels.txt.

helmet
balcony
invalid_balcaony
person
rail

so I have modified code like this but keyerror problem arise.

PGIE_CLASS_ID_helmet = 0
PGIE_CLASS_ID_balcony = 1
PGIE_CLASS_ID_invalid_balcony = 2
PGIE_CLASS_ID_person = 3
PGIE_CLASS_ID_rail =4 


MUXER_OUTPUT_WIDTH=1920
MUXER_OUTPUT_HEIGHT=1080
MUXER_BATCH_TIMEOUT_USEC=4000000
TILED_OUTPUT_WIDTH=1280
TILED_OUTPUT_HEIGHT=720
GST_CAPS_FEATURES_NVMM="memory:NVMM"
OSD_PROCESS_MODE= 0
OSD_DISPLAY_TEXT= 1
pgie_classes_str= ["helmet", "balcony","invalid_balcony", "Person","rail"] # this variable is not used.

# pgie_src_pad_buffer_probe  will extract metadata received on tiler sink pad
# and update params for drawing rectangle, object information etc.
def pgie_src_pad_buffer_probe(pad,info,u_data):
    frame_number=0
    num_rects=0
    got_fps = False
    gst_buffer = info.get_buffer()
    if not gst_buffer:
        print("Unable to get GstBuffer ")
        return
    # 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))
    l_frame = batch_meta.frame_meta_list
    while l_frame is not None:
        try:
            # Note that l_frame.data needs a cast to pyds.NvDsFrameMeta
            # The casting is done by pyds.NvDsFrameMeta.cast()
            # 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(l_frame.data)
        except StopIteration:
            break

        frame_number=frame_meta.frame_num
        l_obj=frame_meta.obj_meta_list
        num_rects = frame_meta.num_obj_meta
        obj_counter = {
        PGIE_CLASS_ID_helmet:0,
        PGIE_CLASS_ID_balcony:0,
        PGIE_CLASS_ID_invalid_balcony:0,
        PGIE_CLASS_ID_person:0,
        PGIE_CLASS_ID_rail:0
        }
        while l_obj is not None:
            try: 
                # Casting l_obj.data to pyds.NvDsObjectMeta
                obj_meta=pyds.NvDsObjectMeta.cast(l_obj.data)
            except StopIteration:
                break
            obj_counter[obj_meta.class_id] += 1
            try: 
                l_obj=l_obj.next
            except StopIteration:
                break
        if not silent:
            print("Frame Number=", frame_number, "Number of Objects=",num_rects,"Vehicle_count=",obj_counter[PGIE_CLASS_ID_helmet],"Person_count=",obj_counter[PGIE_CLASS_ID_person])

        # Update frame rate through this probe
        stream_index = "stream{0}".format(frame_meta.pad_index)
        global perf_data
        perf_data.update_fps(stream_index)

        try:
            l_frame=l_frame.next
        except StopIteration:
            break

    return Gst.PadProbeReturn.OK

I have used the above code is it right way or we should change something in code? any help would be appreciated thanks.

please refer to https://github.com/NVIDIA-AI-IOT/deepstream_python_apps/blob/master/apps/deepstream-test4/deepstream_test_4.py#L259, and https://github.com/NVIDIA-AI-IOT/deepstream_python_apps/blob/master/bindings/src/bindnvdsmeta.cpp#L287

1 Like

Hi thanks for you kind reply.
it worked for me
and one other question is what is meant by this section.

 stream_index = "stream{0}".format(frame_meta.pad_index)

output: it this the frame per sec or what?
**PERF: {'stream0': 29.98}

How I set different bounding box colors for each class?

sorry for the late reply, you might modify NvOSD_RectParams’s NvOSD_ColorParams, please refer to deepstream_python_apps/deepstream_imagedata-multistream_redaction.py at master · NVIDIA-AI-IOT/deepstream_python_apps · GitHub

1 Like

Thnk you @fanzh for yu kind reply.
we love this type of answer where we can see example and get that code for our own use.
Thank you

thanks for the update! If need further support, please open a new one. Thanks

1 Like

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.