├── .gitignore ├── Makefile ├── README.md ├── common └── includes │ ├── deepstream_common.h │ ├── deepstream_config.h │ ├── gst-nvdssr.h │ └── nvdsgstutils.h ├── config └── pgie_config.txt ├── deepstream-rtspsrc-yolo-app ├── src ├── deepstream_rtspsrc_yolo.cpp └── main.h ├── system_drawing.png └── yolo ├── README └── nvdsinfer_custom_impl_Yolo ├── Makefile ├── kernels.cu ├── nvdsinfer_yolo_engine.cpp ├── nvdsparsebbox_Yolo.cpp ├── trt_utils.cpp ├── trt_utils.h ├── yolo.cpp ├── yolo.h ├── yoloPlugins.cpp └── yoloPlugins.h /.gitignore: -------------------------------------------------------------------------------- 1 | *.o 2 | *.so 3 | *.onnx 4 | *.engine 5 | *.mp4 6 | *.sh -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | ################################################################################ 2 | # Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # Permission is hereby granted, free of charge, to any person obtaining a 5 | # copy of this software and associated documentation files (the "Software"), 6 | # to deal in the Software without restriction, including without limitation 7 | # the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | # and/or sell copies of the Software, and to permit persons to whom the 9 | # Software is furnished to do so, subject to the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be included in 12 | # all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | # DEALINGS IN THE SOFTWARE. 21 | ################################################################################ 22 | 23 | CUDA_VER?= 24 | ifeq ($(CUDA_VER),) 25 | $(error "CUDA_VER is not set") 26 | endif 27 | 28 | APP:= deepstream-rtspsrc-yolo-app 29 | CC = g++ 30 | TARGET_DEVICE = $(shell gcc -dumpmachine | cut -f1 -d -) 31 | 32 | NVDS_VERSION:=5.0 33 | 34 | LIB_INSTALL_DIR?=/opt/nvidia/deepstream/deepstream-$(NVDS_VERSION)/lib/ 35 | APP_INSTALL_DIR?=/opt/nvidia/deepstream/deepstream-$(NVDS_VERSION)/bin/ 36 | 37 | SRCS:= $(wildcard src/*.cpp) 38 | INCS:= $(wildcard *.h) 39 | INCS:= $(wildcard common/includes/*.h) 40 | 41 | ifeq ($(TARGET_DEVICE),aarch64) 42 | PKGS:= gstreamer-1.0 opencv4 gstreamer-video-1.0 x11 gstreamer-rtsp-1.0 43 | else 44 | PKGS:= gstreamer-1.0 opencv gstreamer-video-1.0 x11 gstreamer-rtsp-1.0 45 | endif 46 | 47 | CFLAGS:= -fPIC -std=c++11 \ 48 | -I /opt/nvidia/deepstream/deepstream-$(NVDS_VERSION)/sources/includes \ 49 | -I /usr/local/cuda-$(CUDA_VER)/include 50 | 51 | CFLAGS+= `pkg-config --cflags $(PKGS)` 52 | LIBS:= `pkg-config --libs $(PKGS)` 53 | 54 | LIBS+= -L$(LIB_INSTALL_DIR) -lnvdsgst_meta -lnvds_meta -lnvds_inferutils -lnvds_utils -lnvdsgst_helper -lm \ 55 | -lgstrtspserver-1.0 -ldl -L/usr/local/cuda-$(CUDA_VER)/lib64/ -lcudart \ 56 | -Wl,-rpath,$(LIB_INSTALL_DIR) 57 | 58 | ifeq ($(TARGET_DEVICE),aarch64) 59 | CFLAGS+= -DPLATFORM_TEGRA 60 | SRCS+=/opt/nvidia/deepstream/deepstream-$(NVDS_VERSION)/sources/libs/nvdsinfer_customparser/nvdsinfer_custombboxparser.cpp 61 | else 62 | LIBS+= -L$(LIB_INSTALL_DIR) -lnvds_infercustomparser 63 | endif 64 | 65 | OBJS:= $(SRCS:.cpp=.o) 66 | 67 | all: $(APP) 68 | 69 | %.o: %.cpp $(INCS) Makefile 70 | $(CC) -c -o $@ $(CFLAGS) $< 71 | 72 | $(APP): $(OBJS) Makefile 73 | $(CC) -o $(APP) $(OBJS) $(LIBS) 74 | 75 | install: $(APP) 76 | cp -rv $(APP) $(APP_INSTALL_DIR) 77 | 78 | clean: 79 | rm -rf $(OBJS) $(APP) 80 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Deepstream rtspsrc YOLO 2 | 3 | system_drawing 4 | 5 | This app detect objects from RTSP source and create RTSP output. 6 | 7 | This app is made from Deepstream sample-apps. 8 | 9 | YOLOv4 pre-trained model is trained using [https://github.com/AlexeyAB/darknet](https://github.com/AlexeyAB/darknet). 10 | 11 | ONNX model is created using [https://github.com/Tianxiaomo/pytorch-YOLOv4](https://github.com/Tianxiaomo/pytorch-YOLOv4). 12 | 13 | Pre-trained model detects 10 classes of shells. 14 | 15 | When RTSP source is disconnected, the APP wait for restarting RTSP source and try reset pipeline. 16 | 17 | ## Demo 18 | 19 | [![demo](http://img.youtube.com/vi/t9wc6LU26wg/0.jpg)](http://www.youtube.com/watch?v=t9wc6LU26wg "demo") 20 | 21 | ## Environment 22 | 23 | ### Jetson 24 | 25 | - higher model than Jetson Nano Developer Kit(2GB) 26 | - Jetpack 4.4.1 27 | - CUDA 10.2 28 | - TensorRT 7.1.3.0 29 | - 64GB microSD card 30 | 31 | ## Set up Jetson 32 | 33 | Doc for Jetson Nano: [https://developer.nvidia.com/embedded/learn/get-started-jetson-nano-devkit](https://developer.nvidia.com/embedded/learn/get-started-jetson-nano-devkit) 34 | 35 | ## YOLO 36 | 37 | ### Build custom YOLOv4 38 | 39 | ``` 40 | cd yolo/nvdsinfer_custom_impl_Yolo 41 | make 42 | ``` 43 | 44 | You will find `yolo/nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so` and use it as pgie's `custom-lib-path`. 45 | 46 | ### Download pre-trained ONNX model 47 | 48 | This YOLOv4 model detects 10 classes of shell. 49 | 50 | Google Drive: [https://drive.google.com/file/d/1MDd2u1Rn9xemOjd4fjJhm8SIHGqoieF8/view?usp=sharing](https://drive.google.com/file/d/1MDd2u1Rn9xemOjd4fjJhm8SIHGqoieF8/view?usp=sharing) 51 | 52 | ### Convert ONNX model to TensorRT model 53 | 54 | Set path to `trtexec`. 55 | 56 | ``` 57 | export PATH=/usr/local/cuda-10.2/bin:/usr/src/tensorrt/bin:$PATH 58 | ``` 59 | 60 | Convert `` to ``. 61 | 62 | ``` 63 | cd yolo 64 | trtexec \ 65 | --onnx= \ 66 | --explicitBatch \ 67 | --saveEngine= \ 68 | --minShapes=input:1x3x416x416 \ 69 | --optShapes=input:1x3x416x416 \ 70 | --maxShapes=input:1x3x416x416 \ 71 | --shapes=input:1x3x416x416 \ 72 | --fp16 73 | ``` 74 | 75 | ## Run prediction 76 | 77 | Compile app. 78 | 79 | ``` 80 | export CUDA_VER=10.2 81 | make 82 | ``` 83 | 84 | Run app with ``. 85 | 86 | ``` 87 | ./deepstream-rtspsrc-yolo-app \ 88 | 89 | ``` 90 | 91 | Rum app with multi streams. 92 | 93 | ``` 94 | ./deepstream-rtspsrc-yolo-app \ 95 | \ 96 | 97 | ``` 98 | 99 | ### Watch predicted streams. 100 | 101 | `rtsp://:8554/dt-test` 102 | 103 | ### Others 104 | 105 | #### PGIE's interval 106 | 107 | Update interval in `config/pgie_config.txt` to decrease Jetson's work load. 108 | 109 | This config means skip frames of RTSP source every `interval` number. 110 | 111 | ``` 112 | interval=5 113 | ``` 114 | 115 | #### Custom YOLO model 116 | 117 | If you want to use custom YOLOv4 or YOLOv4-tiny model, you need the following tasks. 118 | 119 | - train model and get weight file. 120 | - convert weight file to onnx file. 121 | - convert onnx file to tensorRT file. 122 | 123 | #### Custom classes number 124 | 125 | If you want to change pgie classes number form 10, you neet to change 126 | 127 | - `NUM_CLASSES_YOLO` in `yolo/nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.cpp` 128 | - `PGIE_DETECTED_CLASS_NUM` and `pgie_classes_str` in `src/deepstream_rtspsrc_yolo.cpp` 129 | - `num-detected-classes` in `config/pgie_config.txt` 130 | 131 | ***************************************************************************** 132 | * Copyright (c) 2019-2020 NVIDIA Corporation. All rights reserved. 133 | * 134 | * NVIDIA Corporation and its licensors retain all intellectual property 135 | * and proprietary rights in and to this software, related documentation 136 | * and any modifications thereto. Any use, reproduction, disclosure or 137 | * distribution of this software and related documentation without an express 138 | * license agreement from NVIDIA Corporation is strictly prohibited. 139 | ***************************************************************************** 140 | -------------------------------------------------------------------------------- /common/includes/deepstream_common.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018-2019, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | */ 22 | 23 | #ifndef __NVGSTDS_COMMON_H__ 24 | #define __NVGSTDS_COMMON_H__ 25 | 26 | #ifdef __cplusplus 27 | extern "C" 28 | { 29 | #endif 30 | 31 | #include 32 | 33 | #include "deepstream_config.h" 34 | 35 | #define NVGSTDS_ERR_MSG_V(msg, ...) \ 36 | g_print("** ERROR: <%s:%d>: " msg "\n", __func__, __LINE__, ##__VA_ARGS__) 37 | 38 | #define NVGSTDS_INFO_MSG_V(msg, ...) \ 39 | g_print("** INFO: <%s:%d>: " msg "\n", __func__, __LINE__, ##__VA_ARGS__) 40 | 41 | #define NVGSTDS_WARN_MSG_V(msg, ...) \ 42 | g_print("** WARN: <%s:%d>: " msg "\n", __func__, __LINE__, ##__VA_ARGS__) 43 | 44 | #define NVGSTDS_LINK_ELEMENT(elem1, elem2) \ 45 | do { \ 46 | if (!gst_element_link (elem1,elem2)) { \ 47 | GstCaps *src_caps, *sink_caps; \ 48 | src_caps = gst_pad_query_caps ((GstPad *) (elem1)->srcpads->data, NULL); \ 49 | sink_caps = gst_pad_query_caps ((GstPad *) (elem2)->sinkpads->data, NULL); \ 50 | NVGSTDS_ERR_MSG_V ("Failed to link '%s' (%s) and '%s' (%s)", \ 51 | GST_ELEMENT_NAME (elem1), \ 52 | gst_caps_to_string (src_caps), \ 53 | GST_ELEMENT_NAME (elem2), \ 54 | gst_caps_to_string (sink_caps)); \ 55 | goto done; \ 56 | } \ 57 | } while (0) 58 | 59 | #define NVGSTDS_LINK_ELEMENT_FULL(elem1, elem1_pad_name, elem2, elem2_pad_name) \ 60 | do { \ 61 | GstPad *elem1_pad = gst_element_get_static_pad(elem1, elem1_pad_name); \ 62 | GstPad *elem2_pad = gst_element_get_static_pad(elem2, elem2_pad_name); \ 63 | GstPadLinkReturn ret = gst_pad_link (elem1_pad,elem2_pad); \ 64 | if (ret != GST_PAD_LINK_OK) { \ 65 | gchar *n1 = gst_pad_get_name (elem1_pad); \ 66 | gchar *n2 = gst_pad_get_name (elem2_pad); \ 67 | NVGSTDS_ERR_MSG_V ("Failed to link '%s' and '%s': %d", \ 68 | n1, n2, ret); \ 69 | g_free (n1); \ 70 | g_free (n2); \ 71 | gst_object_unref (elem1_pad); \ 72 | gst_object_unref (elem2_pad); \ 73 | goto done; \ 74 | } \ 75 | gst_object_unref (elem1_pad); \ 76 | gst_object_unref (elem2_pad); \ 77 | } while (0) 78 | 79 | #define NVGSTDS_BIN_ADD_GHOST_PAD_NAMED(bin, elem, pad, ghost_pad_name) \ 80 | do { \ 81 | GstPad *gstpad = gst_element_get_static_pad (elem, pad); \ 82 | if (!gstpad) { \ 83 | NVGSTDS_ERR_MSG_V ("Could not find '%s' in '%s'", pad, \ 84 | GST_ELEMENT_NAME(elem)); \ 85 | goto done; \ 86 | } \ 87 | gst_element_add_pad (bin, gst_ghost_pad_new (ghost_pad_name, gstpad)); \ 88 | gst_object_unref (gstpad); \ 89 | } while (0) 90 | 91 | #define NVGSTDS_BIN_ADD_GHOST_PAD(bin, elem, pad) \ 92 | NVGSTDS_BIN_ADD_GHOST_PAD_NAMED (bin, elem, pad, pad) 93 | 94 | #define NVGSTDS_ELEM_ADD_PROBE(probe_id, elem, pad, probe_func, probe_type, probe_data) \ 95 | do { \ 96 | GstPad *gstpad = gst_element_get_static_pad (elem, pad); \ 97 | if (!gstpad) { \ 98 | NVGSTDS_ERR_MSG_V ("Could not find '%s' in '%s'", pad, \ 99 | GST_ELEMENT_NAME(elem)); \ 100 | goto done; \ 101 | } \ 102 | probe_id = gst_pad_add_probe(gstpad, (probe_type), probe_func, probe_data, NULL); \ 103 | gst_object_unref (gstpad); \ 104 | } while (0) 105 | 106 | #define NVGSTDS_ELEM_REMOVE_PROBE(probe_id, elem, pad) \ 107 | do { \ 108 | if (probe_id == 0 || !elem) { \ 109 | break; \ 110 | } \ 111 | GstPad *gstpad = gst_element_get_static_pad (elem, pad); \ 112 | if (!gstpad) { \ 113 | NVGSTDS_ERR_MSG_V ("Could not find '%s' in '%s'", pad, \ 114 | GST_ELEMENT_NAME(elem)); \ 115 | break; \ 116 | } \ 117 | gst_pad_remove_probe(gstpad, probe_id); \ 118 | gst_object_unref (gstpad); \ 119 | } while (0) 120 | 121 | #define GET_FILE_PATH(path) ((path) + (((path) && strstr ((path), "file://")) ? 7 : 0)) 122 | 123 | 124 | /** 125 | * Function to link sink pad of an element to source pad of tee. 126 | * 127 | * @param[in] tee Tee element. 128 | * @param[in] sinkelem downstream element. 129 | * 130 | * @return true if link successful. 131 | */ 132 | gboolean 133 | link_element_to_tee_src_pad (GstElement * tee, GstElement * sinkelem); 134 | 135 | /** 136 | * Function to link source pad of an element to sink pad of muxer element. 137 | * 138 | * @param[in] streammux muxer element. 139 | * @param[in] elem upstream element. 140 | * @param[in] index pad index of muxer element. 141 | * 142 | * @return true if link successful. 143 | */ 144 | gboolean 145 | link_element_to_streammux_sink_pad (GstElement *streammux, GstElement *elem, 146 | gint index); 147 | 148 | /** 149 | * Function to unlink source pad of an element from sink pad of muxer element. 150 | * 151 | * @param[in] streammux muxer element. 152 | * @param[in] elem upstream element. 153 | * 154 | * @return true if unlinking was successful. 155 | */ 156 | gboolean 157 | unlink_element_from_streammux_sink_pad (GstElement *streammux, GstElement *elem); 158 | 159 | /** 160 | * Function to link sink pad of an element to source pad of demux element. 161 | * 162 | * @param[in] demux demuxer element. 163 | * @param[in] elem downstream element. 164 | * @param[in] index pad index of demuxer element. 165 | * 166 | * @return true if link successful. 167 | */ 168 | gboolean 169 | link_element_to_demux_src_pad (GstElement *demux, GstElement *elem, 170 | guint index); 171 | 172 | /* 173 | * Function to replace string with another string. 174 | * Make sure @ref src is big enough to accommodate replacements. 175 | * 176 | * @param[in] str string to search in. 177 | * @param[in] replace string to replace. 178 | * @param[in] replace_with string to replace @ref replace with. 179 | */ 180 | void 181 | str_replace (gchar * str, const gchar * replace, const gchar * replace_with); 182 | 183 | #ifdef __cplusplus 184 | } 185 | #endif 186 | 187 | #endif 188 | -------------------------------------------------------------------------------- /common/includes/deepstream_config.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018-2020, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | */ 22 | 23 | #ifndef __NVGSTDS_CONFIG_H__ 24 | #define __NVGSTDS_CONFIG_H__ 25 | 26 | #ifdef __aarch64__ 27 | #define IS_TEGRA 28 | #endif 29 | 30 | #define MEMORY_FEATURES "memory:NVMM" 31 | 32 | #ifdef IS_TEGRA 33 | #define NVDS_ELEM_SRC_CAMERA_CSI "nvarguscamerasrc" 34 | #else 35 | #define NVDS_ELEM_SRC_CAMERA_CSI "videotestsrc" 36 | #endif 37 | #define NVDS_ELEM_SRC_CAMERA_V4L2 "v4l2src" 38 | #define NVDS_ELEM_SRC_URI "uridecodebin" 39 | #define NVDS_ELEM_SRC_MULTIFILE "multifilesrc" 40 | #define NVDS_ELEM_SRC_ALSA "alsasrc" 41 | 42 | #define NVDS_ELEM_DECODEBIN "decodebin" 43 | #define NVDS_ELEM_WAVPARSE "wavparse" 44 | 45 | #define NVDS_ELEM_QUEUE "queue" 46 | #define NVDS_ELEM_CAPS_FILTER "capsfilter" 47 | #define NVDS_ELEM_TEE "tee" 48 | 49 | #define NVDS_ELEM_PGIE "nvinfer" 50 | #define NVDS_ELEM_SGIE "nvinfer" 51 | #define NVDS_ELEM_INFER_SERVER "nvinferserver" 52 | #define NVDS_ELEM_TRACKER "nvtracker" 53 | 54 | #define NVDS_ELEM_VIDEO_CONV "nvvideoconvert" 55 | #define NVDS_ELEM_STREAM_MUX "nvstreammux" 56 | #define NVDS_ELEM_STREAM_DEMUX "nvstreamdemux" 57 | #define NVDS_ELEM_TILER "nvmultistreamtiler" 58 | #define NVDS_ELEM_OSD "nvdsosd" 59 | #define NVDS_ELEM_DSANALYTICS_ELEMENT "nvdsanalytics" 60 | #define NVDS_ELEM_DSEXAMPLE_ELEMENT "dsexample" 61 | 62 | #define NVDS_ELEM_DEWARPER "nvdewarper" 63 | #define NVDS_ELEM_SPOTANALYSIS "nvspot" 64 | #define NVDS_ELEM_NVAISLE "nvaisle" 65 | #define NVDS_ELEM_BBOXFILTER "nvbboxfilter" 66 | #define NVDS_ELEM_MSG_CONV "nvmsgconv" 67 | #define NVDS_ELEM_MSG_BROKER "nvmsgbroker" 68 | 69 | #define NVDS_ELEM_SINK_FAKESINK "fakesink" 70 | #define NVDS_ELEM_SINK_FILE "filesink" 71 | #define NVDS_ELEM_SINK_EGL "nveglglessink" 72 | 73 | #define NVDS_ELEM_SINK_OVERLAY "nvoverlaysink" 74 | #define NVDS_ELEM_EGLTRANSFORM "nvegltransform" 75 | 76 | #define NVDS_ELEM_MUX_MP4 "qtmux" 77 | #define NVDS_ELEM_MKV "matroskamux" 78 | 79 | #define NVDS_ELEM_ENC_H264_HW "nvv4l2h264enc" 80 | #define NVDS_ELEM_ENC_H265_HW "nvv4l2h265enc" 81 | #define NVDS_ELEM_ENC_MPEG4 "avenc_mpeg4" 82 | 83 | #define NVDS_ELEM_ENC_H264_SW "x264enc" 84 | #define NVDS_ELEM_ENC_H265_SW "x265enc" 85 | 86 | #define MAX_SOURCE_BINS 1024 87 | #define MAX_SINK_BINS (1024) 88 | #define MAX_SECONDARY_GIE_BINS (16) 89 | #define MAX_MESSAGE_CONSUMERS (16) 90 | 91 | #endif 92 | -------------------------------------------------------------------------------- /common/includes/gst-nvdssr.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019 NVIDIA Corporation. All rights reserved. 3 | * 4 | * NVIDIA Corporation and its licensors retain all intellectual property 5 | * and proprietary rights in and to this software, related documentation 6 | * and any modifications thereto. Any use, reproduction, disclosure or 7 | * distribution of this software and related documentation without an express 8 | * license agreement from NVIDIA Corporation is strictly prohibited. 9 | * 10 | */ 11 | /** 12 | * @file 13 | * NVIDIA DeepStream: Smart recording API 14 | */ 15 | /** 16 | * @defgroup custom_gstreamer Custom Gstreamer APIs 17 | * 18 | * This section defines custom Gstreamer APIs 19 | * 20 | */ 21 | 22 | #ifndef NVDSSR_H_ 23 | #define NVDSSR_H_ 24 | 25 | #include 26 | 27 | /** 28 | * 29 | * @defgroup gstreamer_nvdssr Smart Record 30 | * 31 | * Specifies APIs relating to smart recording. 32 | * 33 | * @ingroup custom_gstreamer 34 | * @{ 35 | */ 36 | #ifdef __cplusplus 37 | extern "C" 38 | { 39 | #endif 40 | 41 | typedef struct NvDsSRRecordingInfo NvDsSRRecordingInfo; 42 | 43 | typedef gpointer (*NvDsSRCallbackFunc) (NvDsSRRecordingInfo *info, gpointer userData); 44 | 45 | typedef guint32 NvDsSRSessionId; 46 | 47 | /** 48 | * Specifies container types. 49 | */ 50 | typedef enum { 51 | NVDSSR_CONTAINER_MP4, 52 | NVDSSR_CONTAINER_MKV 53 | } NvDsSRContainerType; 54 | 55 | /** 56 | * Specifies API return status. 57 | */ 58 | typedef enum { 59 | NVDSSR_STATUS_OK, 60 | NVDSSR_STATUS_INVALID_VAL, 61 | NVDSSR_STATUS_INVALID_OP, 62 | NVDSSR_STATUS_ERROR, 63 | 64 | NVDSSR_STATUS_CUSTOM1 = 100, 65 | NVDSSR_STATUS_CUSTOM2 = 101, 66 | NVDSSR_STATUS_CUSTOM3 = 102 67 | } NvDsSRStatus; 68 | 69 | /** 70 | * Holds initializtion paramerters required to create \ref NvDsSRContext. 71 | */ 72 | typedef struct NvDsSRInitParams 73 | { 74 | /** callback function gets called once recording is complete */ 75 | NvDsSRCallbackFunc callback; 76 | /** recording video container, MP4 / MKV */ 77 | NvDsSRContainerType containerType; 78 | /** optional, recording video width, 0 means no transcode */ 79 | guint width; 80 | /** optional, recording video height, 0 means no transcode */ 81 | guint height; 82 | /** recorded file name prefix */ 83 | gchar *fileNamePrefix; 84 | /** store recorded file under directory path */ 85 | gchar *dirpath; 86 | /** default recording duration in seconds */ 87 | guint defaultDuration; 88 | /** size of video cache in seconds. */ 89 | guint videoCacheSize; 90 | } NvDsSRInitParams; 91 | 92 | /** 93 | * Holds information about smart record instance. 94 | */ 95 | typedef struct NvDsSRContext 96 | { 97 | /** parent bin element. */ 98 | GstElement *recordbin; 99 | /** queue element to cache the content. */ 100 | GstElement *recordQue; 101 | /** child bin to save the content to file. */ 102 | GstElement *encodebin; 103 | /** filesink element */ 104 | GstElement *filesink; 105 | /** flag to check the key frame. */ 106 | gboolean gotKeyFrame; 107 | /** flag to check if recording is on */ 108 | gboolean recordOn; 109 | /** flag to check if encodebin is reset */ 110 | gboolean resetDone; 111 | /** flag to check if encodebin is in playing state. */ 112 | gboolean isPlaying; 113 | /** initialization parameters */ 114 | NvDsSRInitParams initParams; 115 | /** mutex to control the flow */ 116 | GMutex flowLock; 117 | /** thread to reset the encodebin */ 118 | GThread *resetThread; 119 | /** pointer to user provided data */ 120 | gpointer uData; 121 | /** pointer to private data */ 122 | gpointer privData; 123 | } NvDsSRContext; 124 | 125 | /** 126 | * Hold information about video recorded. 127 | */ 128 | typedef struct NvDsSRRecordingInfo 129 | { 130 | /** SR bin context */ 131 | NvDsSRContext *ctx; 132 | /** recording session-id */ 133 | NvDsSRSessionId sessionId; 134 | /** recorded file name */ 135 | gchar *filename; 136 | /** recorded file dir path */ 137 | gchar *dirpath; 138 | /** duration in milliseconds */ 139 | guint64 duration; 140 | /** recorded video container, MP4 / MKV */ 141 | NvDsSRContainerType containerType; 142 | /** recorded video width*/ 143 | guint width; 144 | /** recorded video height*/ 145 | guint height; 146 | } NvDsSRRecordingInfo; 147 | 148 | /** 149 | * \brief Creates the instance of smart record. 150 | * 151 | * This function creates the instance of smart record and returns the pointer 152 | * to an allocated \ref NvDsSRContext. The \a params structure must be filled 153 | * with initialization parameters required to create the instance. 154 | * 155 | * recordbin of \ref NvDsSRContext is smart record bin which must be added 156 | * to the pipeline. It expects encoded frames which will be muxed and saved to 157 | * the file. Add this bin after parser element in the pipeline. 158 | * 159 | * Call NvDsSRDestroy() to free resources allocated by this function. 160 | * 161 | * @param[out] ctx An indirect pointer to the smart record instance. 162 | * @param[in] params A pointer to a \ref NvDsSRInitParams structure. 163 | * 164 | * @return NVDSSR_STATUS_OK if successful, or corresponding error otherwise. 165 | */ 166 | NvDsSRStatus NvDsSRCreate (NvDsSRContext **ctx, NvDsSRInitParams *params); 167 | 168 | /** 169 | * \brief Starts the video recording. 170 | * 171 | * This function starts writing the cached video data to a file. It returns 172 | * the session id which later can be used in NvDsSRStop() to stop the 173 | * corresponding recording. 174 | * 175 | * Here startTime specifies the seconds before the current time and duration 176 | * specifies the seconds after the start of recording. 177 | * If current time is t1, content from t1 - startTime to t1 + duration will 178 | * be saved to file. Therefore a total of startTime + duration seconds of data 179 | * will be recorded. 180 | * 181 | * @param[in] ctx A pointer to a \ref NvDsSRContext. 182 | * @param[out] sessionId A pointer to a \ref NvDsSRSessionId. 183 | * @param[in] startTime Seconds before the current time. Should be less than video cache size. 184 | * @param[in] duration Duration value in seconds after the start of recording. 185 | * @param[in] userData A pointer to user specified data. 186 | * 187 | * @return NVDSSR_STATUS_OK if successful, or corresponding error otherwise. 188 | */ 189 | NvDsSRStatus NvDsSRStart (NvDsSRContext *ctx, NvDsSRSessionId *sessionId, 190 | guint startTime, guint duration, gpointer userData); 191 | 192 | /** 193 | * \brief Stops the previously started recording. 194 | * 195 | * @param[in] ctx A pointer to a \ref NvDsSRContext. 196 | * @param[in] sessionId Id of seesion to stop. 197 | * 198 | * @return NVDSSR_STATUS_OK if successful, or corresponding error otherwise. 199 | */ 200 | NvDsSRStatus NvDsSRStop (NvDsSRContext *ctx, NvDsSRSessionId sessionId); 201 | 202 | /** 203 | * \brief Destroys the instance of smart record. 204 | * 205 | * This function releases the resources previously allocated by NvDsSRCreate(). 206 | * 207 | * @param[in] ctx A pointer to a \ref NvDsSRContext to be freed. 208 | * 209 | * @return NVDSSR_STATUS_OK if successful, or corresponding error otherwise. 210 | */ 211 | NvDsSRStatus NvDsSRDestroy (NvDsSRContext *ctx); 212 | 213 | #ifdef __cplusplus 214 | } 215 | #endif 216 | #endif /* NVDSSR_H_ */ 217 | 218 | /** @} */ 219 | -------------------------------------------------------------------------------- /common/includes/nvdsgstutils.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * NVIDIA Corporation and its licensors retain all intellectual property 5 | * and proprietary rights in and to this software, related documentation 6 | * and any modifications thereto. Any use, reproduction, disclosure or 7 | * distribution of this software and related documentation without an express 8 | * license agreement from NVIDIA Corporation is strictly prohibited. 9 | */ 10 | 11 | /** 12 | * @file 13 | * Defines NVIDIA DeepStream GStreamer Utilities 14 | * 15 | * @b Description: This file specifies the NVIDIA DeepStream GStreamer utility 16 | * functions. 17 | * 18 | */ 19 | /** 20 | * @defgroup gstreamer_utils Utilities: Gstreamer utilities API 21 | * 22 | * Specifies GStreamer utilities functions, used to configure the source to generate NTP Sync values. 23 | * 24 | * @ingroup NvDsUtilsApi 25 | * @{ 26 | */ 27 | #ifndef __NVDS_GSTUTILS_H__ 28 | #define __NVDS_GSTUTILS_H__ 29 | 30 | #include 31 | 32 | #ifdef __cplusplus 33 | extern "C" { 34 | #endif 35 | #include 36 | 37 | /** 38 | * Configure the source to generate NTP sync values for RTSP sources. 39 | * 40 | * These values are used by the DeepStream GStreamer element NvStreamMux to 41 | * calculate the NTP time of the frames at the source. 42 | * 43 | * This functionality is dependent on the RTSP sending the RTCP Sender Reports. 44 | * source. 45 | * 46 | * This function only works for RTSP sources i.e. GStreamer elements "rtspsrc" 47 | * or "uridecodebin" with an RTSP uri. 48 | * 49 | * params[in] src_elem GStreamer source element to be configured. 50 | */ 51 | void configure_source_for_ntp_sync (GstElement *src_elem); 52 | 53 | #ifdef __cplusplus 54 | } 55 | #endif 56 | 57 | #endif 58 | 59 | /** @} */ -------------------------------------------------------------------------------- /config/pgie_config.txt: -------------------------------------------------------------------------------- 1 | ################################################################################ 2 | # Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # Permission is hereby granted, free of charge, to any person obtaining a 5 | # copy of this software and associated documentation files (the "Software"), 6 | # to deal in the Software without restriction, including without limitation 7 | # the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | # and/or sell copies of the Software, and to permit persons to whom the 9 | # Software is furnished to do so, subject to the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be included in 12 | # all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | # DEALINGS IN THE SOFTWARE. 21 | ################################################################################ 22 | 23 | # Following properties are mandatory when engine files are not specified: 24 | # int8-calib-file(Only in INT8) 25 | # Caffemodel mandatory properties: model-file, proto-file, output-blob-names 26 | # UFF: uff-file, input-dims, uff-input-blob-name, output-blob-names 27 | # ONNX: onnx-file 28 | # 29 | # Mandatory properties for detectors: 30 | # num-detected-classes 31 | # 32 | # Optional properties for detectors: 33 | # cluster-mode(Default=Group Rectangles), interval(Primary mode only, Default=0) 34 | # custom-lib-path, 35 | # parse-bbox-func-name 36 | # 37 | # Mandatory properties for classifiers: 38 | # classifier-threshold, network-type (Default=0 i.e. Detector) 39 | # 40 | # Optional properties for classifiers: 41 | # classifier-async-mode(Secondary mode only, Default=false) 42 | # 43 | # Optional properties in secondary mode: 44 | # operate-on-gie-id(Default=0), operate-on-class-ids(Defaults to all classes), 45 | # input-object-min-width, input-object-min-height, input-object-max-width, 46 | # input-object-max-height 47 | # 48 | # Following properties are always recommended: 49 | # batch-size(Default=1) 50 | # 51 | # Other optional properties: 52 | # net-scale-factor(Default=1), network-mode(Default=0 i.e FP32), 53 | # model-color-format(Default=0 i.e. RGB) model-engine-file, labelfile-path, 54 | # mean-file, gie-unique-id(Default=0), offsets, process-mode (Default=1 i.e. primary), 55 | # custom-lib-path, network-mode(Default=0 i.e FP32) 56 | # 57 | # The values in the config file are overridden by values set through GObject 58 | # properties. 59 | 60 | [property] 61 | gpu-id=0 62 | net-scale-factor=0.0039215697906911373 63 | #0=RGB, 1=BGR, 2=GRAY 64 | model-color-format=0 65 | model-engine-file=../yolo/yolov4_1_3_416_416_asari.engine 66 | ## 0=FP32, 1=INT8, 2=FP16 mode 67 | network-mode=2 68 | num-detected-classes=10 69 | gie-unique-id=1 70 | network-type=0 71 | is-classifier=0 72 | ## 0=Group Rectangles, 1=DBSCAN, 2=NMS, 3= DBSCAN+NMS Hybrid, 4 = None(No clustering) 73 | cluster-mode=2 74 | maintain-aspect-ratio=0 75 | parse-bbox-func-name=NvDsInferParseCustomYoloV4 76 | custom-lib-path=../yolo/nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so 77 | engine-create-func-name=NvDsInferYoloCudaEngineGet 78 | #scaling-filter=0 79 | #scaling-compute-hw=0 80 | batch-size=1 81 | interval=1 82 | 83 | [class-attrs-all] 84 | pre-cluster-threshold=0.2 85 | eps=0.2 86 | group-threshold=1 87 | -------------------------------------------------------------------------------- /deepstream-rtspsrc-yolo-app: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/belarbi2733/deepstream-rtspsrc-yolo/d90fc9c325f01478cb7ffc26879bd89f5c7b6735/deepstream-rtspsrc-yolo-app -------------------------------------------------------------------------------- /src/deepstream_rtspsrc_yolo.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include "cuda_runtime_api.h" 11 | #include "gstnvdsmeta.h" 12 | #include "gstnvdsinfer.h" 13 | #include "nvdsinfer_custom_impl.h" 14 | #include "nvds_version.h" 15 | #include "main.h" 16 | 17 | #define NVINFER_PLUGIN "nvinfer" 18 | #define INFER_PGIE_CONFIG_FILE "config/pgie_config.txt" 19 | 20 | #define MUXER_OUTPUT_WIDTH 1920 21 | #define MUXER_OUTPUT_HEIGHT 1080 22 | #define MUXER_BATCH_TIMEOUT_USEC 40000 23 | #define MEMORY_FEATURES "memory:NVMM" 24 | #define TILED_OUTPUT_WIDTH 1920 25 | #define TILED_OUTPUT_HEIGHT 1080 26 | #define PGIE_CLASS_ID_0 0 27 | #define PGIE_CLASS_ID_1 1 28 | #define PGIE_CLASS_ID_2 2 29 | #define PGIE_CLASS_ID_3 3 30 | #define PGIE_CLASS_ID_4 4 31 | #define PGIE_CLASS_ID_5 5 32 | #define PGIE_CLASS_ID_6 6 33 | #define PGIE_CLASS_ID_7 7 34 | #define PGIE_CLASS_ID_8 8 35 | #define PGIE_CLASS_ID_9 9 36 | #define PGIE_DETECTED_CLASS_NUM 10 37 | #define MAX_SINK_BINS (1024) 38 | #define MAX_INSTANCES 128 39 | #define MAX_DISPLAY_LEN 64 40 | 41 | AppCtx *appCtx[MAX_INSTANCES]; 42 | 43 | GMainLoop *loop = NULL; 44 | static gboolean install_mux_eosmonitor_probe = FALSE; 45 | gint frame_number = 0; 46 | const gchar pgie_classes_str[PGIE_DETECTED_CLASS_NUM][32] = 47 | { "ibokisago", 48 | "karasu", 49 | "hototogisu", 50 | "aramujiro", 51 | "mate", 52 | "tsumeta", 53 | "baka", 54 | "shiofuki", 55 | "kagami", 56 | "asari" }; 57 | 58 | unsigned int nvds_lib_major_version = NVDS_VERSION_MAJOR; 59 | unsigned int nvds_lib_minor_version = NVDS_VERSION_MINOR; 60 | static GstRTSPServer *server [MAX_SINK_BINS]; 61 | static guint server_count = 0; 62 | static GMutex server_cnt_lock; 63 | guint num_sources = 0; 64 | 65 | static GstPadProbeReturn 66 | osd_sink_pad_buffer_probe (GstPad * pad, GstPadProbeInfo * info, gpointer u_data) 67 | { 68 | GstBuffer *buf = (GstBuffer *) info->data; 69 | guint num_rects = 0; 70 | NvDsObjectMeta *obj_meta = NULL; 71 | guint count_0 = 0; 72 | guint count_1 = 0; 73 | guint count_2 = 0; 74 | guint count_3 = 0; 75 | guint count_4 = 0; 76 | guint count_5 = 0; 77 | guint count_6 = 0; 78 | guint count_7 = 0; 79 | guint count_8 = 0; 80 | guint count_9 = 0; 81 | guint truck_count = 0; 82 | NvDsMetaList *l_frame = NULL; 83 | NvDsMetaList *l_obj = NULL; 84 | NvDsMetaList *l_cls = NULL; 85 | NvDsDisplayMeta *display_meta3 = NULL; 86 | NvDsBatchMeta *batch_meta = gst_buffer_get_nvds_batch_meta (buf); 87 | 88 | for (l_frame = batch_meta->frame_meta_list; l_frame != NULL; 89 | l_frame = l_frame->next) { 90 | NvDsFrameMeta *frame_meta = (NvDsFrameMeta *) (l_frame->data); 91 | int offset = 0; 92 | for (l_obj = frame_meta->obj_meta_list; l_obj != NULL; l_obj = l_obj->next) { 93 | obj_meta = (NvDsObjectMeta *) (l_obj->data); 94 | display_meta3 = nvds_acquire_display_meta_from_pool (batch_meta); 95 | NvOSD_TextParams *txt_params3 = &display_meta3->text_params[0]; 96 | for (l_cls = obj_meta->classifier_meta_list; l_cls != NULL; l_cls = l_cls->next) { 97 | NvDsClassifierMeta *cls_meta = (NvDsClassifierMeta *) (l_cls->data); 98 | } 99 | switch (obj_meta->class_id) { 100 | case PGIE_CLASS_ID_0: 101 | count_0++; 102 | num_rects++; 103 | txt_params3->text_bg_clr.red = 0.0; 104 | txt_params3->text_bg_clr.green = 0.0; 105 | txt_params3->text_bg_clr.blue = 0.0; 106 | break; 107 | case PGIE_CLASS_ID_1: 108 | count_1++; 109 | num_rects++; 110 | txt_params3->text_bg_clr.red = 0.5; 111 | txt_params3->text_bg_clr.green = 0.0; 112 | txt_params3->text_bg_clr.blue = 0.0; 113 | break; 114 | case PGIE_CLASS_ID_2: 115 | count_2++; 116 | num_rects++; 117 | txt_params3->text_bg_clr.red = 0.0; 118 | txt_params3->text_bg_clr.green = 0.5; 119 | txt_params3->text_bg_clr.blue = 0.0; 120 | break; 121 | case PGIE_CLASS_ID_3: 122 | count_3++; 123 | num_rects++; 124 | txt_params3->text_bg_clr.red = 1.0; 125 | txt_params3->text_bg_clr.green = 1.0; 126 | txt_params3->text_bg_clr.blue = 1.0; 127 | break; 128 | case PGIE_CLASS_ID_4: 129 | count_4++; 130 | num_rects++; 131 | txt_params3->text_bg_clr.red = 0.0; 132 | txt_params3->text_bg_clr.green = 1.0; 133 | txt_params3->text_bg_clr.blue = 1.0; 134 | break; 135 | case PGIE_CLASS_ID_5: 136 | count_5++; 137 | num_rects++; 138 | txt_params3->text_bg_clr.red = 1.0; 139 | txt_params3->text_bg_clr.green = 0.0; 140 | txt_params3->text_bg_clr.blue = 1.0; 141 | break; 142 | case PGIE_CLASS_ID_6: 143 | count_6++; 144 | num_rects++; 145 | txt_params3->text_bg_clr.red = 1.0; 146 | txt_params3->text_bg_clr.green = 1.0; 147 | txt_params3->text_bg_clr.blue = 0.0; 148 | break; 149 | case PGIE_CLASS_ID_7: 150 | count_7++; 151 | num_rects++; 152 | txt_params3->text_bg_clr.red = 0.0; 153 | txt_params3->text_bg_clr.green = 0.0; 154 | txt_params3->text_bg_clr.blue = 1.0; 155 | break; 156 | case PGIE_CLASS_ID_8: 157 | count_8++; 158 | num_rects++; 159 | txt_params3->text_bg_clr.red = 0.0; 160 | txt_params3->text_bg_clr.green = 1.0; 161 | txt_params3->text_bg_clr.blue = 0.0; 162 | break; 163 | case PGIE_CLASS_ID_9: 164 | count_9++; 165 | num_rects++; 166 | txt_params3->text_bg_clr.red = 1.0; 167 | txt_params3->text_bg_clr.green = 0.0; 168 | txt_params3->text_bg_clr.blue = 0.0; 169 | break; 170 | } 171 | 172 | NvOSD_RectParams rect_params = obj_meta->rect_params; 173 | float left = rect_params.left; 174 | float top = rect_params.top; 175 | float width = rect_params.width; 176 | float height = rect_params.height; 177 | float confidence = obj_meta->confidence; 178 | 179 | std::string result; 180 | char result_0[32]; 181 | char result_1[32]; 182 | display_meta3->num_labels = 1; 183 | txt_params3->display_text = (gchar *) g_malloc0 (MAX_DISPLAY_LEN); 184 | 185 | snprintf(result_0, 32, "%s ", pgie_classes_str[obj_meta->class_id]); 186 | snprintf(result_1, 32, "%.0f％", confidence * 100); 187 | 188 | result.append(result_0); 189 | result.append(result_1); 190 | 191 | offset = snprintf ( 192 | txt_params3->display_text, MAX_DISPLAY_LEN, "%s", result.c_str() 193 | ); 194 | /* Now set the offsets where the string should appear */ 195 | txt_params3->x_offset = left; 196 | txt_params3->y_offset = top-40; 197 | 198 | /* Font , font-color and font-size */ 199 | txt_params3->font_params.font_name = (gchar *) "Serif"; 200 | txt_params3->font_params.font_size = 20; 201 | txt_params3->font_params.font_color.red = 0.0; 202 | txt_params3->font_params.font_color.green = 0.0; 203 | txt_params3->font_params.font_color.blue = 0.0; 204 | txt_params3->font_params.font_color.alpha = 1.0; 205 | 206 | /* Text background color */ 207 | txt_params3->set_bg_clr = 1; 208 | // txt_params3->text_bg_clr.red = 1.0; 209 | // txt_params3->text_bg_clr.green = 1.0; 210 | // txt_params3->text_bg_clr.blue = 1.0; 211 | txt_params3->text_bg_clr.alpha = 0.5; 212 | 213 | nvds_add_display_meta_to_frame (frame_meta, display_meta3); 214 | 215 | } 216 | if (num_rects > 0) { 217 | g_print ("Frame Number = %d Number of objects = %d " 218 | "ID_0: %d ID_1: %d ID_2: %d ID_3: %d ID_4: %d ID_5: %d ID_6: %d ID_7: %d ID_8: %d ID_9: %d \n", 219 | frame_meta->frame_num, num_rects, count_0, count_1, count_2, count_3, count_4, count_5, count_6, count_7, count_8, count_9); 220 | } 221 | } 222 | 223 | frame_number++; 224 | return GST_PAD_PROBE_OK; 225 | } 226 | 227 | gboolean 228 | link_element_to_streammux_sink_pad (GstElement *streammux, GstElement *elem, 229 | gint index) 230 | { 231 | gboolean ret = FALSE; 232 | GstPad *mux_sink_pad = NULL; 233 | GstPad *src_pad = NULL; 234 | gchar pad_name[16]; 235 | 236 | if (index >= 0) { 237 | g_snprintf (pad_name, 16, "sink_%u", index); 238 | pad_name[15] = '\0'; 239 | } else { 240 | strcpy (pad_name, "sink_%u"); 241 | } 242 | 243 | mux_sink_pad = gst_element_get_request_pad (streammux, pad_name); 244 | if (!mux_sink_pad) { 245 | NVGSTDS_ERR_MSG_V ("Failed to get sink pad from streammux"); 246 | goto done; 247 | } 248 | 249 | src_pad = gst_element_get_static_pad (elem, "src"); 250 | if (!src_pad) { 251 | NVGSTDS_ERR_MSG_V ("Failed to get src pad from '%s'", 252 | GST_ELEMENT_NAME (elem)); 253 | goto done; 254 | } 255 | 256 | if (gst_pad_link (src_pad, mux_sink_pad) != GST_PAD_LINK_OK) { 257 | NVGSTDS_ERR_MSG_V ("Failed to link '%s' and '%s'", GST_ELEMENT_NAME (streammux), 258 | GST_ELEMENT_NAME (elem)); 259 | goto done; 260 | } 261 | 262 | ret = TRUE; 263 | 264 | done: 265 | if (mux_sink_pad) { 266 | gst_object_unref (mux_sink_pad); 267 | } 268 | if (src_pad) { 269 | gst_object_unref (src_pad); 270 | } 271 | return ret; 272 | } 273 | 274 | gboolean 275 | link_element_to_tee_src_pad (GstElement *tee, GstElement *sinkelem) 276 | { 277 | gboolean ret = FALSE; 278 | GstPad *tee_src_pad = NULL; 279 | GstPad *sinkpad = NULL; 280 | GstPadTemplate *padtemplate = NULL; 281 | 282 | padtemplate = (GstPadTemplate *) 283 | gst_element_class_get_pad_template (GST_ELEMENT_GET_CLASS (tee), 284 | "src_%u"); 285 | tee_src_pad = gst_element_request_pad (tee, padtemplate, NULL, NULL); 286 | if (!tee_src_pad) { 287 | NVGSTDS_ERR_MSG_V ("Failed to get src pad from tee"); 288 | goto done; 289 | } 290 | 291 | sinkpad = gst_element_get_static_pad (sinkelem, "sink"); 292 | if (!sinkpad) { 293 | NVGSTDS_ERR_MSG_V ("Failed to get sink pad from '%s'", 294 | GST_ELEMENT_NAME (sinkelem)); 295 | goto done; 296 | } 297 | 298 | if (gst_pad_link (tee_src_pad, sinkpad) != GST_PAD_LINK_OK) { 299 | NVGSTDS_ERR_MSG_V ("Failed to link '%s' and '%s'", GST_ELEMENT_NAME (tee), 300 | GST_ELEMENT_NAME (sinkelem)); 301 | goto done; 302 | } 303 | 304 | ret = TRUE; 305 | 306 | done: 307 | if (tee_src_pad) { 308 | gst_object_unref (tee_src_pad); 309 | } 310 | if (sinkpad) { 311 | gst_object_unref (sinkpad); 312 | } 313 | return ret; 314 | } 315 | 316 | /** 317 | * Function called at regular interval when source bin is 318 | * changing state async. This function watches the state of 319 | * the source bin and sets it to PLAYING if the state of source 320 | * bin stops at PAUSED when changing state ASYNC. 321 | */ 322 | static gboolean 323 | watch_source_async_state_change (gpointer data) 324 | { 325 | NvDsSrcBin *src_bin = (NvDsSrcBin *) data; 326 | GstState state, pending; 327 | GstStateChangeReturn ret; 328 | 329 | ret = gst_element_get_state (src_bin->bin, &state, &pending, 0); 330 | 331 | // Bin is still changing state ASYNC. Wait for some more time. 332 | if (ret == GST_STATE_CHANGE_ASYNC) 333 | return TRUE; 334 | 335 | // Bin state change failed / failed to get state 336 | if (ret == GST_STATE_CHANGE_FAILURE) { 337 | src_bin->async_state_watch_running = FALSE; 338 | return FALSE; 339 | } 340 | 341 | // Bin successfully changed state to PLAYING. Stop watching state 342 | if (state == GST_STATE_PLAYING) { 343 | src_bin->reconfiguring = FALSE; 344 | src_bin->async_state_watch_running = FALSE; 345 | return FALSE; 346 | } 347 | 348 | // Bin has stopped ASYNC state change but has not gone into 349 | // PLAYING. Expliclity set state to PLAYING and keep watching 350 | // state 351 | gst_element_set_state (src_bin->bin, GST_STATE_PLAYING); 352 | 353 | return TRUE; 354 | } 355 | 356 | /** 357 | * Probe function to monitor data output from rtspsrc. 358 | */ 359 | static GstPadProbeReturn 360 | rtspsrc_monitor_probe_func (GstPad * pad, GstPadProbeInfo * info, 361 | gpointer u_data) 362 | { 363 | NvDsSrcBin *bin = (NvDsSrcBin *) u_data; 364 | if (info->type & GST_PAD_PROBE_TYPE_BUFFER) { 365 | g_mutex_lock(&bin->bin_lock); 366 | gettimeofday (&bin->last_buffer_time, NULL); 367 | g_mutex_unlock(&bin->bin_lock); 368 | } 369 | return GST_PAD_PROBE_OK; 370 | } 371 | 372 | gboolean 373 | reset_source_pipeline (gpointer data) 374 | { 375 | NvDsSrcBin *src_bin = (NvDsSrcBin *) data; 376 | GstState state, pending; 377 | GstStateChangeReturn ret; 378 | 379 | g_mutex_lock(&src_bin->bin_lock); 380 | gettimeofday (&src_bin->last_buffer_time, NULL); 381 | gettimeofday (&src_bin->last_reconnect_time, NULL); 382 | g_mutex_unlock(&src_bin->bin_lock); 383 | 384 | if (gst_element_set_state (src_bin->bin, 385 | GST_STATE_NULL) == GST_STATE_CHANGE_FAILURE) { 386 | GST_ERROR_OBJECT (src_bin->bin, "Can't set source bin to NULL"); 387 | return FALSE; 388 | } 389 | NVGSTDS_INFO_MSG_V ("Resetting source %d", src_bin->bin_id); 390 | 391 | // GST_CAT_INFO (NVDS_APP, "Reset source pipeline %s %p\n,", __func__, src_bin); 392 | if (!gst_element_sync_state_with_parent (src_bin->bin)) { 393 | GST_ERROR_OBJECT (src_bin->bin, "Couldn't sync state with parent"); 394 | } 395 | 396 | ret = gst_element_get_state (src_bin->bin, &state, &pending, 0); 397 | 398 | if (ret == GST_STATE_CHANGE_ASYNC || ret == GST_STATE_CHANGE_NO_PREROLL) { 399 | if (!src_bin->async_state_watch_running) 400 | g_timeout_add (20, watch_source_async_state_change, src_bin); 401 | src_bin->async_state_watch_running = TRUE; 402 | src_bin->reconfiguring = TRUE; 403 | } else if (ret == GST_STATE_CHANGE_SUCCESS && state == GST_STATE_PLAYING) { 404 | src_bin->reconfiguring = FALSE; 405 | } 406 | return FALSE; 407 | } 408 | 409 | /** 410 | * Function called at regular interval to check if NV_DS_SOURCE_RTSP type 411 | * source in the pipeline is down / disconnected. This function try to 412 | * reconnect the source by resetting that source pipeline. 413 | */ 414 | static gboolean 415 | watch_source_status (gpointer data) 416 | { 417 | NvDsSrcBin *src_bin = (NvDsSrcBin *) data; 418 | struct timeval current_time; 419 | gettimeofday (¤t_time, NULL); 420 | static struct timeval last_reset_time_global = {0, 0}; 421 | gdouble time_diff_msec_since_last_reset = 422 | 1000.0 * (current_time.tv_sec - last_reset_time_global.tv_sec) + 423 | (current_time.tv_usec - last_reset_time_global.tv_usec) / 1000.0; 424 | 425 | if (src_bin->reconfiguring) { 426 | guint time_since_last_reconnect_sec = 427 | current_time.tv_sec - src_bin->last_reconnect_time.tv_sec; 428 | if (time_since_last_reconnect_sec >= SOURCE_RESET_INTERVAL_SEC) { 429 | if (time_diff_msec_since_last_reset > 3000) { 430 | last_reset_time_global = current_time; 431 | // source is still not up, reconfigure it again. 432 | reset_source_pipeline (src_bin); 433 | } 434 | } 435 | } else { 436 | gint time_since_last_buf_sec = 0; 437 | g_mutex_lock (&src_bin->bin_lock); 438 | if (src_bin->last_buffer_time.tv_sec != 0) { 439 | time_since_last_buf_sec = 440 | current_time.tv_sec - src_bin->last_buffer_time.tv_sec; 441 | } 442 | g_mutex_unlock (&src_bin->bin_lock); 443 | 444 | // Reset source bin if no buffers are received in the last 445 | // `rtsp_reconnect_interval_sec` seconds. 446 | if (src_bin->rtsp_reconnect_interval_sec > 0 && 447 | time_since_last_buf_sec >= src_bin->rtsp_reconnect_interval_sec) { 448 | if (time_diff_msec_since_last_reset > 3000) { 449 | last_reset_time_global = current_time; 450 | 451 | NVGSTDS_WARN_MSG_V ("No data from source %d since last %u sec. Trying reconnection", 452 | src_bin->bin_id, time_since_last_buf_sec); 453 | reset_source_pipeline (src_bin); 454 | } 455 | } 456 | 457 | } 458 | return TRUE; 459 | } 460 | 461 | /** 462 | * callback function to receive messages from components 463 | * in the pipeline. 464 | */ 465 | static gboolean 466 | bus_callback (GstBus * bus, GstMessage * message, gpointer data) 467 | { 468 | AppCtx *appCtx = (AppCtx *) data; 469 | switch (GST_MESSAGE_TYPE (message)) { 470 | case GST_MESSAGE_INFO:{ 471 | GError *error = NULL; 472 | gchar *debuginfo = NULL; 473 | gst_message_parse_info (message, &error, &debuginfo); 474 | g_printerr ("INFO from %s: %s\n", 475 | GST_OBJECT_NAME (message->src), error->message); 476 | if (debuginfo) { 477 | g_printerr ("Debug info: %s\n", debuginfo); 478 | } 479 | g_error_free (error); 480 | g_free (debuginfo); 481 | break; 482 | } 483 | case GST_MESSAGE_WARNING:{ 484 | GError *error = NULL; 485 | gchar *debuginfo = NULL; 486 | gst_message_parse_warning (message, &error, &debuginfo); 487 | g_printerr ("WARNING from %s: %s\n", 488 | GST_OBJECT_NAME (message->src), error->message); 489 | if (debuginfo) { 490 | g_printerr ("Debug info: %s\n", debuginfo); 491 | } 492 | g_error_free (error); 493 | g_free (debuginfo); 494 | break; 495 | } 496 | case GST_MESSAGE_ERROR:{ 497 | GError *error = NULL; 498 | gchar *debuginfo = NULL; 499 | guint i = 0; 500 | gst_message_parse_error (message, &error, &debuginfo); 501 | g_printerr ("ERROR from %s: %s\n", 502 | GST_OBJECT_NAME (message->src), error->message); 503 | if (debuginfo) { 504 | g_printerr ("Debug info: %s\n", debuginfo); 505 | } 506 | NvDsSrcParentBin *bin = &appCtx->pipeline.multi_src_bin; 507 | GstElement *msg_src_elem = (GstElement *) GST_MESSAGE_SRC (message); 508 | gboolean bin_found = FALSE; 509 | /* Find the source bin which generated the error. */ 510 | while (msg_src_elem && !bin_found) { 511 | for (i = 0; i < bin->num_bins && !bin_found; i++) { 512 | printf("msg_src_elem: %s\n", gst_element_get_name(msg_src_elem)); 513 | 514 | if (!bin->sub_bins[i].src_elem) { 515 | printf("bin->sub_bins[i].src_elem is empty\n"); 516 | goto done; 517 | } 518 | if (bin->sub_bins[i].src_elem == msg_src_elem || 519 | bin->sub_bins[i].bin == msg_src_elem) { 520 | bin_found = TRUE; 521 | printf("bin to reset is found\n"); 522 | break; 523 | } 524 | } 525 | msg_src_elem = GST_ELEMENT_PARENT (msg_src_elem); 526 | } 527 | 528 | if (i != bin->num_bins) { 529 | // Error from one of RTSP source. 530 | NvDsSrcBin *subBin = &bin->sub_bins[i]; 531 | if (!subBin->reconfiguring || 532 | g_strrstr(debuginfo, "500 (Internal Server Error)")) { 533 | subBin->reconfiguring = TRUE; 534 | printf("trying to call reset_source_pipeline\n"); 535 | g_timeout_add (0, reset_source_pipeline, subBin); 536 | } 537 | g_error_free (error); 538 | g_free (debuginfo); 539 | return TRUE; 540 | } 541 | 542 | done: 543 | g_error_free (error); 544 | g_free (debuginfo); 545 | appCtx->return_value = -1; 546 | appCtx->quit = TRUE; 547 | break; 548 | } 549 | case GST_MESSAGE_STATE_CHANGED:{ 550 | GstState oldstate, newstate; 551 | gst_message_parse_state_changed (message, &oldstate, &newstate, NULL); 552 | if (GST_ELEMENT (GST_MESSAGE_SRC (message)) == appCtx->pipeline.pipeline) { 553 | switch (newstate) { 554 | case GST_STATE_PLAYING: 555 | NVGSTDS_INFO_MSG_V ("Pipeline running\n"); 556 | GST_DEBUG_BIN_TO_DOT_FILE_WITH_TS (GST_BIN (appCtx-> 557 | pipeline.pipeline), GST_DEBUG_GRAPH_SHOW_ALL, 558 | "ds-app-playing"); 559 | break; 560 | case GST_STATE_PAUSED: 561 | if (oldstate == GST_STATE_PLAYING) { 562 | NVGSTDS_INFO_MSG_V ("Pipeline paused\n"); 563 | } 564 | break; 565 | case GST_STATE_READY: 566 | GST_DEBUG_BIN_TO_DOT_FILE_WITH_TS (GST_BIN (appCtx->pipeline. 567 | pipeline), GST_DEBUG_GRAPH_SHOW_ALL, "ds-app-ready"); 568 | if (oldstate == GST_STATE_NULL) { 569 | NVGSTDS_INFO_MSG_V ("Pipeline ready\n"); 570 | } else { 571 | NVGSTDS_INFO_MSG_V ("Pipeline stopped\n"); 572 | } 573 | break; 574 | case GST_STATE_NULL: 575 | g_mutex_lock (&appCtx->app_lock); 576 | g_cond_broadcast (&appCtx->app_cond); 577 | g_mutex_unlock (&appCtx->app_lock); 578 | break; 579 | default: 580 | break; 581 | } 582 | } 583 | break; 584 | } 585 | case GST_MESSAGE_EOS:{ 586 | /* 587 | * In normal scenario, this would use g_main_loop_quit() to exit the 588 | * loop and release the resources. Since this application might be 589 | * running multiple pipelines through configuration files, it should wait 590 | * till all pipelines are done. 591 | */ 592 | NVGSTDS_INFO_MSG_V ("Received EOS. Exiting ...\n"); 593 | appCtx->quit = TRUE; 594 | return FALSE; 595 | break; 596 | } 597 | default: 598 | break; 599 | } 600 | return TRUE; 601 | } 602 | 603 | 604 | /* Returning FALSE from this callback will make rtspsrc ignore the stream. 605 | * Ignore audio and add the proper depay element based on codec. */ 606 | static gboolean 607 | cb_rtspsrc_select_stream (GstElement *rtspsrc, guint num, GstCaps *caps, 608 | gpointer user_data) 609 | { 610 | GstStructure *str = gst_caps_get_structure (caps, 0); 611 | const gchar *media = gst_structure_get_string (str, "media"); 612 | const gchar *encoding_name = gst_structure_get_string (str, "encoding-name"); 613 | gchar elem_name[50]; 614 | NvDsSrcBin *bin = (NvDsSrcBin *) user_data; 615 | gboolean ret = FALSE; 616 | 617 | gboolean is_video = (!g_strcmp0 (media, "video")); 618 | 619 | if (!is_video) 620 | return FALSE; 621 | 622 | /* Create and add depay element only if it is not created yet. */ 623 | if (!bin->depay) { 624 | g_snprintf (elem_name, sizeof (elem_name), "depay_elem%d", bin->bin_id); 625 | 626 | /* Add the proper depay element based on codec. */ 627 | if (!g_strcmp0 (encoding_name, "H264")) { 628 | bin->depay = gst_element_factory_make ("rtph264depay", elem_name); 629 | g_snprintf (elem_name, sizeof (elem_name), "h264parse_elem%d", bin->bin_id); 630 | bin->parser = gst_element_factory_make ("h264parse", elem_name); 631 | } else if (!g_strcmp0 (encoding_name, "H265")) { 632 | printf("passes rtph265depay\n"); 633 | bin->depay = gst_element_factory_make ("rtph265depay", elem_name); 634 | g_snprintf (elem_name, sizeof (elem_name), "h265parse_elem%d", bin->bin_id); 635 | bin->parser = gst_element_factory_make ("h265parse", elem_name); 636 | } else { 637 | NVGSTDS_WARN_MSG_V ("%s not supported", encoding_name); 638 | return FALSE; 639 | } 640 | 641 | if (!bin->depay) { 642 | NVGSTDS_ERR_MSG_V ("Failed to create '%s'", elem_name); 643 | return FALSE; 644 | } 645 | gst_bin_add_many (GST_BIN (bin->bin), bin->depay, bin->parser, NULL); 646 | 647 | NVGSTDS_LINK_ELEMENT (bin->depay, bin->parser); 648 | NVGSTDS_LINK_ELEMENT (bin->parser, bin->tee_rtsp_pre_decode); 649 | 650 | if (!gst_element_sync_state_with_parent (bin->depay)) { 651 | NVGSTDS_ERR_MSG_V ("'%s' failed to sync state with parent", elem_name); 652 | return FALSE; 653 | } 654 | gst_element_sync_state_with_parent (bin->parser); 655 | } 656 | 657 | ret = TRUE; 658 | done: 659 | return ret; 660 | } 661 | 662 | static void 663 | cb_newpad2 (GstElement * decodebin, GstPad * pad, gpointer data) 664 | { 665 | GstCaps *caps = gst_pad_query_caps (pad, NULL); 666 | const GstStructure *str = gst_caps_get_structure (caps, 0); 667 | const gchar *name = gst_structure_get_name (str); 668 | gint input_width = 1920; 669 | gint input_height = 1080; 670 | gint input_fps_n = 1; 671 | gint input_fps_d = 1; 672 | if (!strncmp (name, "video", 5)) { 673 | NvDsSrcBin *bin = (NvDsSrcBin *) data; 674 | GstPad *sinkpad = gst_element_get_static_pad (bin->cap_filter, "sink"); 675 | if (gst_pad_link (pad, sinkpad) != GST_PAD_LINK_OK) { 676 | 677 | NVGSTDS_ERR_MSG_V ("Failed to link decodebin to pipeline"); 678 | } else { 679 | gst_structure_get_int (str, "width", &input_width); 680 | gst_structure_get_int (str, "height", &input_height); 681 | gst_structure_get_fraction (str, "framerate", &input_fps_n, &input_fps_d); 682 | } 683 | gst_object_unref (sinkpad); 684 | } 685 | gst_caps_unref (caps); 686 | } 687 | 688 | static void 689 | cb_newpad3 (GstElement * decodebin, GstPad * pad, gpointer data) 690 | { 691 | GstCaps *caps = gst_pad_query_caps (pad, NULL); 692 | const GstStructure *str = gst_caps_get_structure (caps, 0); 693 | const gchar *name = gst_structure_get_name (str); 694 | if (g_strrstr (name, "x-rtp")) { 695 | NvDsSrcBin *bin = (NvDsSrcBin *) data; 696 | GstPad *sinkpad = gst_element_get_static_pad (bin->depay, "sink"); 697 | if (gst_pad_link (pad, sinkpad) != GST_PAD_LINK_OK) { 698 | 699 | NVGSTDS_ERR_MSG_V ("Failed to link depay loader to rtsp src"); 700 | } 701 | gst_object_unref (sinkpad); 702 | } 703 | gst_caps_unref (caps); 704 | } 705 | 706 | static void 707 | cb_newpad (GstElement * decodebin, GstPad * decoder_src_pad, gpointer data) 708 | { 709 | g_print ("In cb_newpad\n"); 710 | GstCaps *caps = gst_pad_get_current_caps (decoder_src_pad); 711 | const GstStructure *str = gst_caps_get_structure (caps, 0); 712 | const gchar *name = gst_structure_get_name (str); 713 | GstElement *source_bin = (GstElement *) data; 714 | GstCapsFeatures *features = gst_caps_get_features (caps, 0); 715 | 716 | if (!strncmp (name, "video", 5)) { 717 | if (gst_caps_features_contains (features, MEMORY_FEATURES)) { 718 | GstPad *bin_ghost_pad = gst_element_get_static_pad (source_bin, "src"); 719 | if (!gst_ghost_pad_set_target (GST_GHOST_PAD (bin_ghost_pad), 720 | decoder_src_pad)) { 721 | g_printerr ("Failed to link decoder src pad to source bin ghost pad\n"); 722 | } 723 | gst_object_unref (bin_ghost_pad); 724 | } else { 725 | g_printerr ("Error: Decodebin did not pick nvidia decoder plugin.\n"); 726 | } 727 | } 728 | } 729 | 730 | static void 731 | decodebin_child_added (GstChildProxy * child_proxy, GObject * object, 732 | gchar * name, gpointer user_data) 733 | { 734 | NvDsSrcBin *bin = (NvDsSrcBin *) user_data; 735 | 736 | if (g_strrstr (name, "decodebin") == name) { 737 | g_signal_connect (G_OBJECT (object), "child-added", 738 | G_CALLBACK (decodebin_child_added), user_data); 739 | } 740 | if ((g_strrstr (name, "h264parse") == name) || 741 | (g_strrstr (name, "h265parse") == name)) { 742 | g_object_set (object, "config-interval", -1, NULL); 743 | } 744 | if (g_strrstr (name, "fakesink") == name) { 745 | g_object_set (object, "enable-last-sample", FALSE, NULL); 746 | } 747 | if (g_strrstr (name, "nvcuvid") == name) { 748 | g_object_set (object, "gpu-id", 0, NULL); 749 | } 750 | if (g_strstr_len (name, -1, "omx") == name) { 751 | g_object_set (object, "skip-frames", 2, NULL); 752 | g_object_set (object, "disable-dvfs", TRUE, NULL); 753 | } 754 | if (g_strstr_len (name, -1, "nvjpegdec") == name) { 755 | g_object_set (object, "DeepStream", TRUE, NULL); 756 | } 757 | if (g_strstr_len (name, -1, "nvv4l2decoder") == name) { 758 | #ifdef __aarch64__ 759 | g_object_set (object, "enable-max-performance", TRUE, NULL); 760 | #endif 761 | } 762 | done: 763 | return; 764 | } 765 | 766 | static gboolean 767 | create_rtsp_src_bin (guint index, gchar * location, NvDsSrcBin * bin) 768 | { 769 | NvDsSRContext *ctx = NULL; 770 | gboolean ret = FALSE; 771 | gchar elem_name[50]; 772 | GstCaps *caps = NULL; 773 | GstCapsFeatures *feature = NULL; 774 | 775 | bin->rtsp_reconnect_interval_sec = 10; 776 | 777 | g_snprintf (elem_name, sizeof (elem_name), "src_elem%d", bin->bin_id); 778 | bin->src_elem = gst_element_factory_make ("rtspsrc", elem_name); 779 | if (!bin->src_elem) { 780 | NVGSTDS_ERR_MSG_V ("Failed to create '%s'", elem_name); 781 | goto done; 782 | } 783 | 784 | g_signal_connect (G_OBJECT(bin->src_elem), "select-stream", 785 | G_CALLBACK(cb_rtspsrc_select_stream), 786 | bin); 787 | 788 | g_object_set (G_OBJECT (bin->src_elem), "location", location, NULL); 789 | g_object_set (G_OBJECT (bin->src_elem), "latency", bin->latency, NULL); 790 | g_object_set (G_OBJECT (bin->src_elem), "drop-on-latency", TRUE, NULL); 791 | configure_source_for_ntp_sync (bin->src_elem); 792 | g_signal_connect (G_OBJECT (bin->src_elem), "pad-added", 793 | G_CALLBACK (cb_newpad3), bin); 794 | 795 | g_snprintf (elem_name, sizeof (elem_name), "tee_rtsp_elem%d", bin->bin_id); 796 | bin->tee_rtsp_pre_decode = gst_element_factory_make ("tee", elem_name); 797 | if (!bin->tee_rtsp_pre_decode) { 798 | NVGSTDS_ERR_MSG_V ("Failed to create '%s'", elem_name); 799 | goto done; 800 | } 801 | 802 | g_snprintf (elem_name, sizeof (elem_name), "dec_que%d", bin->bin_id); 803 | bin->dec_que = gst_element_factory_make ("queue", elem_name); 804 | if (!bin->dec_que) { 805 | NVGSTDS_ERR_MSG_V ("Failed to create '%s'", elem_name); 806 | goto done; 807 | } 808 | 809 | if (bin->rtsp_reconnect_interval_sec > 0) { 810 | printf("rtsp_reconnect_interval_sec: %u\n", bin->rtsp_reconnect_interval_sec); 811 | NVGSTDS_ELEM_ADD_PROBE (bin->rtspsrc_monitor_probe, bin->dec_que, 812 | "sink", rtspsrc_monitor_probe_func, 813 | GST_PAD_PROBE_TYPE_BUFFER, 814 | bin); 815 | install_mux_eosmonitor_probe = TRUE; 816 | } 817 | 818 | g_snprintf (elem_name, sizeof (elem_name), "decodebin_elem%d", bin->bin_id); 819 | bin->decodebin = gst_element_factory_make ("decodebin", elem_name); 820 | if (!bin->decodebin) { 821 | NVGSTDS_ERR_MSG_V ("Failed to create '%s'", elem_name); 822 | goto done; 823 | } 824 | 825 | g_signal_connect (G_OBJECT (bin->decodebin), "pad-added", 826 | G_CALLBACK (cb_newpad2), bin); 827 | g_signal_connect (G_OBJECT (bin->decodebin), "child-added", 828 | G_CALLBACK (decodebin_child_added), bin); 829 | 830 | 831 | g_snprintf (elem_name, sizeof (elem_name), "src_que%d", bin->bin_id); 832 | bin->cap_filter = gst_element_factory_make ("queue", elem_name); 833 | if (!bin->cap_filter) { 834 | NVGSTDS_ERR_MSG_V ("Failed to create '%s'", elem_name); 835 | goto done; 836 | } 837 | 838 | g_mutex_init (&bin->bin_lock); 839 | 840 | g_snprintf(elem_name, sizeof(elem_name), "nvvidconv_elem%d", bin->bin_id); 841 | bin->nvvidconv = gst_element_factory_make("nvvideoconvert", elem_name); 842 | if (!bin->nvvidconv) 843 | { 844 | NVGSTDS_ERR_MSG_V("Could not create element 'nvvidconv_elem'"); 845 | goto done; 846 | } 847 | caps = gst_caps_new_empty_simple("video/x-raw"); 848 | feature = gst_caps_features_new("memory:NVMM", NULL); 849 | gst_caps_set_features(caps, 0, feature); 850 | 851 | bin->cap_filter1 = 852 | gst_element_factory_make("capsfilter", "src_cap_filter_nvvidconv"); 853 | if (!bin->cap_filter1) 854 | { 855 | NVGSTDS_ERR_MSG_V("Could not create 'queue'"); 856 | goto done; 857 | } 858 | 859 | g_object_set(G_OBJECT(bin->cap_filter1), "caps", caps, NULL); 860 | gst_caps_unref(caps); 861 | 862 | gst_bin_add_many (GST_BIN(bin->bin), 863 | bin->src_elem, 864 | bin->tee_rtsp_pre_decode, 865 | bin->dec_que, 866 | bin->decodebin, 867 | bin->cap_filter, 868 | bin->nvvidconv, 869 | bin->cap_filter1, 870 | NULL); 871 | 872 | link_element_to_tee_src_pad(bin->tee_rtsp_pre_decode, bin->dec_que); 873 | NVGSTDS_LINK_ELEMENT (bin->dec_que, bin->decodebin); 874 | 875 | if (ctx) 876 | link_element_to_tee_src_pad(bin->tee_rtsp_pre_decode, ctx->recordbin); 877 | 878 | NVGSTDS_LINK_ELEMENT (bin->cap_filter, bin->nvvidconv); 879 | NVGSTDS_LINK_ELEMENT (bin->nvvidconv, bin->cap_filter1); 880 | NVGSTDS_BIN_ADD_GHOST_PAD (bin->bin, bin->cap_filter1, "src"); 881 | 882 | ret = TRUE; 883 | 884 | g_timeout_add (1000, watch_source_status, bin); 885 | 886 | done: 887 | if (!ret) { 888 | NVGSTDS_ERR_MSG_V ("%s failed", __func__); 889 | } 890 | return ret; 891 | 892 | } 893 | 894 | /** 895 | * Probe function to drop EOS events from nvstreammux when RTSP sources 896 | * are being used so that app does not quit from EOS in case of RTSP 897 | * connection errors and tries to reconnect. 898 | */ 899 | static GstPadProbeReturn 900 | nvstreammux_eosmonitor_probe_func (GstPad * pad, GstPadProbeInfo * info, 901 | gpointer u_data) 902 | { 903 | if (info->type & GST_PAD_PROBE_TYPE_EVENT_DOWNSTREAM) { 904 | GstEvent *event = (GstEvent *) info->data; 905 | if (GST_EVENT_TYPE (event) == GST_EVENT_EOS) 906 | return GST_PAD_PROBE_DROP; 907 | } 908 | return GST_PAD_PROBE_OK; 909 | } 910 | 911 | static gboolean 912 | start_rtsp_streaming (guint rtsp_port_num, guint updsink_port_num, guint64 udp_buffer_size) 913 | { 914 | GstRTSPMountPoints *mounts; 915 | GstRTSPMediaFactory *factory; 916 | char udpsrc_pipeline[512]; 917 | char port_num_Str[64] = { 0 }; 918 | 919 | if (udp_buffer_size == 0) 920 | udp_buffer_size = 512 * 1024; 921 | 922 | sprintf (udpsrc_pipeline, 923 | "( udpsrc name=pay0 port=%d buffer-size=%lu caps=\"application/x-rtp, media=video, " 924 | "clock-rate=90000, encoding-name=%s, payload=96 \" )", 925 | updsink_port_num, udp_buffer_size, "H264"); 926 | 927 | sprintf (port_num_Str, "%d", rtsp_port_num); 928 | 929 | g_mutex_lock (&server_cnt_lock); 930 | 931 | server [server_count] = gst_rtsp_server_new (); 932 | g_object_set (server [server_count], "service", port_num_Str, NULL); 933 | 934 | mounts = gst_rtsp_server_get_mount_points (server [server_count]); 935 | 936 | factory = gst_rtsp_media_factory_new (); 937 | gst_rtsp_media_factory_set_launch (factory, udpsrc_pipeline); 938 | 939 | gst_rtsp_mount_points_add_factory (mounts, "/ds-test", factory); 940 | 941 | g_object_unref (mounts); 942 | 943 | gst_rtsp_server_attach (server [server_count], NULL); 944 | 945 | server_count++; 946 | 947 | g_mutex_unlock (&server_cnt_lock); 948 | 949 | g_print 950 | ("\n *** DeepStream: Launched RTSP Streaming at rtsp://localhost:%d/ds-test ***\n\n", 951 | rtsp_port_num); 952 | 953 | return TRUE; 954 | } 955 | 956 | 957 | static void 958 | usage(const char *bin) 959 | { 960 | g_printerr 961 | ("Usage: %s [-t infer-type] ... \n", 962 | bin); 963 | g_printerr 964 | (" -t infer-type: select form [infer, inferserver], infer by default\n"); 965 | } 966 | 967 | int 968 | main (int argc, char *argv[]) 969 | { 970 | printf("app start..\n"); 971 | appCtx[0] = (AppCtx *) g_malloc0 (sizeof (AppCtx)); 972 | NvDsPipeline *pipeline = &appCtx[0]->pipeline; 973 | GstElement *streammux = NULL, 974 | *sink = NULL, 975 | *pgie = NULL, 976 | *nvosd = NULL, 977 | *queue2 = NULL, 978 | *nvvidconv2 = NULL, 979 | *caps_filter = NULL, 980 | *encoder = NULL, 981 | *codecparser = NULL, 982 | *rtppay = NULL, 983 | *tiler = NULL; 984 | GstCaps *caps = NULL; 985 | GstBus *bus = NULL; 986 | guint i=0; 987 | guint tiler_rows, tiler_columns; 988 | GstPad *osd_sink_pad = NULL; 989 | guint pgie_batch_size; 990 | const char *infer_plugin = NVINFER_PLUGIN; 991 | gboolean ret = FALSE; 992 | 993 | /* define numsorces from argc */ 994 | if (argc < 2) { 995 | g_printerr ("Usage: %s \n", argv[0]); 996 | return -1; 997 | } 998 | num_sources = argc - 1; 999 | /* define numsorces from argc */ 1000 | 1001 | nvds_version(&nvds_lib_major_version, &nvds_lib_minor_version); 1002 | 1003 | gst_init (&argc, &argv); 1004 | 1005 | pipeline->pipeline = gst_pipeline_new ("pipeline"); 1006 | if (!pipeline->pipeline) { 1007 | NVGSTDS_ERR_MSG_V ("Failed to create pipeline"); 1008 | } 1009 | 1010 | bus = gst_pipeline_get_bus (GST_PIPELINE (pipeline->pipeline)); 1011 | pipeline->bus_id = gst_bus_add_watch (bus, bus_callback, appCtx[0]); 1012 | gst_object_unref (bus); 1013 | 1014 | pipeline->multi_src_bin.reset_thread = NULL; 1015 | 1016 | pipeline->multi_src_bin.bin = gst_bin_new ("multi_src_bin"); 1017 | if (!pipeline->multi_src_bin.bin) { 1018 | NVGSTDS_ERR_MSG_V ("Failed to create element 'multi_src_bin'"); 1019 | goto done; 1020 | } 1021 | g_object_set (pipeline->multi_src_bin.bin, "message-forward", TRUE, NULL); 1022 | 1023 | pipeline->multi_src_bin.streammux = gst_element_factory_make ("nvstreammux", "src_bin_muxer"); 1024 | if (!pipeline->multi_src_bin.streammux) { 1025 | NVGSTDS_ERR_MSG_V ("Failed to create element 'src_bin_muxer'"); 1026 | goto done; 1027 | } 1028 | gst_bin_add (GST_BIN (pipeline->multi_src_bin.bin), pipeline->multi_src_bin.streammux); 1029 | 1030 | /* create rtsp src bin */ 1031 | for (i = 0; i < num_sources; i++) { 1032 | printf("%s: %u\n", "start createing source bins", i); 1033 | GstPad *sinkpad, *srcpad; 1034 | gchar elem_name[50]; 1035 | g_snprintf (elem_name, sizeof (elem_name), "src_sub_bin%d", i); 1036 | pipeline->multi_src_bin.sub_bins[i].bin = gst_bin_new (elem_name); 1037 | if (!pipeline->multi_src_bin.sub_bins[i].bin) { 1038 | NVGSTDS_ERR_MSG_V ("Failed to create '%s'", elem_name); 1039 | goto done; 1040 | } 1041 | 1042 | pipeline->multi_src_bin.sub_bins[i].bin_id = pipeline->multi_src_bin.sub_bins[i].source_id = i; 1043 | pipeline->multi_src_bin.live_source = TRUE; 1044 | pipeline->multi_src_bin.sub_bins[i].eos_done = TRUE; 1045 | pipeline->multi_src_bin.sub_bins[i].reset_done = TRUE; 1046 | 1047 | printf("argv[2]: %s\n", argv[i+1]); 1048 | if(!create_rtsp_src_bin (i, argv[i + 1], &pipeline->multi_src_bin.sub_bins[i])){ 1049 | g_printerr ("Failed to create source bin. Exiting.\n"); 1050 | } 1051 | gst_bin_add (GST_BIN (pipeline->multi_src_bin.bin), pipeline->multi_src_bin.sub_bins[i].bin); 1052 | 1053 | if (!link_element_to_streammux_sink_pad (pipeline->multi_src_bin.streammux, 1054 | pipeline->multi_src_bin.sub_bins[i].bin, i)) { 1055 | NVGSTDS_ERR_MSG_V ("source %d cannot be linked to mux's sink pad %p\n", i, pipeline->multi_src_bin.streammux); 1056 | goto done; 1057 | } 1058 | pipeline->multi_src_bin.num_bins++; 1059 | printf("pipeline->multi_src_bin.num_bins: %d\n", pipeline->multi_src_bin.num_bins); 1060 | } 1061 | /* create rtsp src bin */ 1062 | 1063 | NVGSTDS_BIN_ADD_GHOST_PAD (pipeline->multi_src_bin.bin, pipeline->multi_src_bin.streammux, "src"); 1064 | 1065 | if (install_mux_eosmonitor_probe) { 1066 | NVGSTDS_ELEM_ADD_PROBE (pipeline->multi_src_bin.nvstreammux_eosmonitor_probe, pipeline->multi_src_bin.streammux, 1067 | "src", nvstreammux_eosmonitor_probe_func, 1068 | GST_PAD_PROBE_TYPE_EVENT_DOWNSTREAM, 1069 | &pipeline->multi_src_bin); 1070 | } 1071 | 1072 | done: 1073 | 1074 | loop = g_main_loop_new (NULL, FALSE); 1075 | 1076 | pgie = gst_element_factory_make (infer_plugin, "primary-nvinference-engine"); 1077 | tiler = gst_element_factory_make ("nvmultistreamtiler", "nvtiler"); 1078 | nvosd = gst_element_factory_make ("nvdsosd", "nv-onscreendisplay"); 1079 | g_object_set (G_OBJECT (nvosd), "display-text", 1, NULL); 1080 | g_object_set (G_OBJECT (nvosd), "display-clock", 1, NULL); 1081 | g_object_set (G_OBJECT (nvosd), "display-bbox", 1, NULL); 1082 | g_object_set (G_OBJECT (nvosd), "display-mask", 1, NULL); 1083 | g_object_set (G_OBJECT (nvosd), "process-mode", 2, NULL); 1084 | queue2 = gst_element_factory_make ("queue", "queue2"); 1085 | nvvidconv2 = gst_element_factory_make ("nvvideoconvert", "convertor2"); 1086 | guint gpu_id = 0; 1087 | caps_filter = gst_element_factory_make ("capsfilter", "capsfilter"); 1088 | caps = gst_caps_from_string ("video/x-raw(memory:NVMM), format=I420"); 1089 | encoder = gst_element_factory_make ("nvv4l2h264enc", "encoder"); 1090 | guint profile = 0; 1091 | guint bitrate = 1000000; 1092 | guint iframeinterval = 60; 1093 | codecparser = gst_element_factory_make ("h264parse", "h264-parser2"); 1094 | rtppay = gst_element_factory_make ("rtph264pay", "rtppay"); 1095 | sink = gst_element_factory_make ("udpsink", "udpsink"); 1096 | 1097 | g_object_set (G_OBJECT (pipeline->multi_src_bin.streammux), "width", MUXER_OUTPUT_WIDTH, "height", 1098 | MUXER_OUTPUT_HEIGHT, "batch-size", num_sources, 1099 | "batched-push-timeout", MUXER_BATCH_TIMEOUT_USEC, NULL); 1100 | g_object_get (G_OBJECT (pgie), "batch-size", &pgie_batch_size, NULL); 1101 | if (pgie_batch_size != num_sources) { 1102 | g_printerr 1103 | ("WARNING: Overriding infer-config batch-size (%d) with number of sources (%d)\n", 1104 | pgie_batch_size, num_sources); 1105 | g_object_set (G_OBJECT (pgie), "batch-size", num_sources, NULL); 1106 | } 1107 | g_object_set (G_OBJECT (pgie), "batch-size", 1, NULL); 1108 | g_object_set (G_OBJECT (pgie), "config-file-path", INFER_PGIE_CONFIG_FILE, NULL); 1109 | g_object_set (G_OBJECT (nvvidconv2), "gpu-id", gpu_id, NULL); 1110 | g_object_set (G_OBJECT (caps_filter), "caps", caps, NULL); 1111 | g_object_set (G_OBJECT (encoder), "profile", profile, NULL); 1112 | g_object_set (G_OBJECT (encoder), "iframeinterval", iframeinterval, NULL); 1113 | g_object_set (G_OBJECT (encoder), "bitrate", bitrate, NULL); 1114 | g_object_set (G_OBJECT (encoder), "preset-level", 1, NULL); 1115 | g_object_set (G_OBJECT (encoder), "insert-sps-pps", 1, NULL); 1116 | g_object_set (G_OBJECT (encoder), "bufapi-version", 1, NULL); 1117 | g_object_set (G_OBJECT (sink), "host", "224.224.255.255", "port", 1118 | 5400, "async", FALSE, "sync", 0, NULL); 1119 | 1120 | if (!pgie) { 1121 | g_printerr ("pgie could not be created. Exiting.\n"); 1122 | return -1; 1123 | } 1124 | if (!tiler) { 1125 | g_printerr ("tiler could not be created. Exiting.\n"); 1126 | return -1; 1127 | } 1128 | if (!nvosd) { 1129 | g_printerr ("nvosd could not be created. Exiting.\n"); 1130 | return -1; 1131 | } 1132 | if (!sink) { 1133 | g_printerr ("sink could not be created. Exiting.\n"); 1134 | return -1; 1135 | } 1136 | if (!queue2 || !nvvidconv2 || !caps_filter || !encoder || !codecparser ) { 1137 | g_printerr ("One element could not be created. Exiting.\n"); 1138 | return -1; 1139 | } 1140 | if (!rtppay) { 1141 | g_printerr ("rtppay could not be created. Exiting.\n"); 1142 | return -1; 1143 | } 1144 | 1145 | /* a tee after the tiler which shall be connected to sink(s) */ 1146 | pipeline->tiler_tee = gst_element_factory_make ("tee", "tiler_tee"); 1147 | if (!pipeline->tiler_tee) { 1148 | NVGSTDS_ERR_MSG_V ("Failed to create element 'tiler_tee'"); 1149 | } 1150 | /* Create demuxer only if tiled display is disabled. */ 1151 | pipeline->demuxer = 1152 | gst_element_factory_make ("nvstreamdemux", "demuxer"); 1153 | if (!pipeline->demuxer) { 1154 | NVGSTDS_ERR_MSG_V ("Failed to create element 'demuxer'"); 1155 | } 1156 | 1157 | /* setting tiler */ 1158 | tiler_rows = (guint) sqrt (num_sources); 1159 | tiler_columns = (guint) ceil (1.0 * num_sources / tiler_rows); 1160 | g_object_set (G_OBJECT (tiler), "rows", tiler_rows, "columns", tiler_columns, 1161 | "width", TILED_OUTPUT_WIDTH, "height", TILED_OUTPUT_HEIGHT, NULL); 1162 | 1163 | gst_bin_add_many (GST_BIN (pipeline->pipeline), 1164 | pipeline->multi_src_bin.bin, 1165 | pipeline->multi_src_bin.streammux, 1166 | pgie, 1167 | tiler, 1168 | nvosd, 1169 | queue2, 1170 | nvvidconv2, 1171 | caps_filter, 1172 | encoder, 1173 | codecparser, 1174 | rtppay, 1175 | sink, 1176 | NULL); 1177 | 1178 | if (!gst_element_link_many (pipeline->multi_src_bin.bin, 1179 | pgie, 1180 | tiler, 1181 | nvosd, 1182 | queue2, 1183 | nvvidconv2, 1184 | caps_filter, 1185 | encoder, 1186 | codecparser, 1187 | rtppay, 1188 | sink, 1189 | NULL)) { 1190 | g_printerr ("Elements could not be linked. Exiting.\n"); 1191 | return -1; 1192 | } 1193 | 1194 | osd_sink_pad = gst_element_get_static_pad (pgie, "src"); 1195 | if (!osd_sink_pad) 1196 | g_print ("Unable to get sink pad\n"); 1197 | else 1198 | gst_pad_add_probe (osd_sink_pad, GST_PAD_PROBE_TYPE_BUFFER, 1199 | osd_sink_pad_buffer_probe, NULL, NULL); 1200 | 1201 | ret = TRUE; 1202 | ret = start_rtsp_streaming (8554, 5400, 0); 1203 | if (ret != TRUE) { 1204 | g_print ("%s: start_rtsp_straming function failed\n", __func__); 1205 | } 1206 | 1207 | g_print ("Now playing...\n"); 1208 | gst_element_set_state (pipeline->pipeline, GST_STATE_PLAYING); 1209 | g_print ("Running...\n"); 1210 | g_main_loop_run (loop); 1211 | g_print ("Returned, stopping playback\n"); 1212 | gst_element_set_state (pipeline->pipeline, GST_STATE_NULL); 1213 | g_print ("Deleting pipeline\n"); 1214 | gst_object_unref (GST_OBJECT (pipeline->pipeline)); 1215 | g_source_remove (pipeline->bus_id); 1216 | g_main_loop_unref (loop); 1217 | return 0; 1218 | } 1219 | -------------------------------------------------------------------------------- /src/main.h: -------------------------------------------------------------------------------- 1 | #ifndef __NVGSTDS_APP_H__ 2 | #define __NVGSTDS_APP_H__ 3 | 4 | #ifdef __cplusplus 5 | extern "C" 6 | { 7 | #endif 8 | 9 | #include 10 | #include 11 | #include 12 | #include "../common/includes/gst-nvdssr.h" 13 | #include "../common/includes/deepstream_common.h" 14 | #include "../common/includes/nvdsgstutils.h" 15 | 16 | #define NVGSTDS_ERR_MSG_V(msg, ...) \ 17 | g_print("** ERROR: <%s:%d>: " msg "\n", __func__, __LINE__, ##__VA_ARGS__) 18 | 19 | #define NVGSTDS_INFO_MSG_V(msg, ...) \ 20 | g_print("** INFO: <%s:%d>: " msg "\n", __func__, __LINE__, ##__VA_ARGS__) 21 | 22 | #define NVGSTDS_WARN_MSG_V(msg, ...) \ 23 | g_print("** WARN: <%s:%d>: " msg "\n", __func__, __LINE__, ##__VA_ARGS__) 24 | 25 | #define MAX_SOURCE_BINS 1024 26 | #define SOURCE_RESET_INTERVAL_SEC 30 27 | 28 | typedef struct _AppCtx AppCtx; 29 | 30 | typedef enum 31 | { 32 | NV_DS_SOURCE_CAMERA_V4L2 = 1, 33 | NV_DS_SOURCE_URI, 34 | NV_DS_SOURCE_URI_MULTIPLE, 35 | NV_DS_SOURCE_RTSP, 36 | NV_DS_SOURCE_CAMERA_CSI, 37 | } NvDsSourceType; 38 | 39 | 40 | typedef struct 41 | { 42 | GstElement *bin; 43 | GstElement *src_elem; 44 | GstElement *cap_filter; 45 | GstElement *cap_filter1; 46 | GstElement *depay; 47 | GstElement *parser; 48 | GstElement *enc_que; 49 | GstElement *dec_que; 50 | GstElement *decodebin; 51 | GstElement *enc_filter; 52 | GstElement *encbin_que; 53 | GstElement *tee; 54 | GstElement *tee_rtsp_pre_decode; 55 | GstElement *fakesink_queue; 56 | GstElement *fakesink; 57 | GstElement *nvvidconv; 58 | 59 | gboolean do_record; 60 | guint64 pre_event_rec; 61 | GMutex bin_lock; 62 | guint bin_id; 63 | gint rtsp_reconnect_interval_sec; 64 | struct timeval last_buffer_time; 65 | struct timeval last_reconnect_time; 66 | gulong src_buffer_probe; 67 | gulong rtspsrc_monitor_probe; 68 | gpointer bbox_meta; 69 | GstBuffer *inbuf; 70 | gchar *location; 71 | gchar *file; 72 | gchar *direction; 73 | gint latency; 74 | gboolean got_key_frame; 75 | gboolean eos_done; 76 | gboolean reset_done; 77 | gboolean live_source; 78 | gboolean reconfiguring; 79 | gboolean async_state_watch_running; 80 | // NvDsDewarperBin dewarper_bin; 81 | gulong probe_id; 82 | guint64 accumulated_base; 83 | guint64 prev_accumulated_base; 84 | guint source_id; 85 | // NvDsSourceConfig *config; 86 | gpointer recordCtx; 87 | } NvDsSrcBin; 88 | 89 | typedef struct 90 | { 91 | GstElement *bin; 92 | GstElement *streammux; 93 | GThread *reset_thread; 94 | NvDsSrcBin sub_bins[MAX_SOURCE_BINS]; 95 | guint num_bins; 96 | guint num_fr_on; 97 | gboolean live_source; 98 | gulong nvstreammux_eosmonitor_probe; 99 | } NvDsSrcParentBin; 100 | 101 | typedef struct 102 | { 103 | // gulong primary_bbox_buffer_probe_id; 104 | guint bus_id; 105 | GstElement *pipeline; 106 | NvDsSrcParentBin multi_src_bin; 107 | // NvDsInstanceBin instance_bins[MAX_SOURCE_BINS]; 108 | // NvDsInstanceBin demux_instance_bins[MAX_SOURCE_BINS]; 109 | // NvDsInstanceBin common_elements; 110 | GstElement *tiler_tee; 111 | // NvDsTiledDisplayBin tiled_display_bin; 112 | GstElement *demuxer; 113 | // NvDsDsExampleBin dsexample_bin; 114 | AppCtx *appCtx; 115 | } NvDsPipeline; 116 | 117 | struct _AppCtx 118 | { 119 | // gboolean version; 120 | // gboolean cintr; 121 | // gboolean show_bbox_text; 122 | // gboolean seeking; 123 | gboolean quit; 124 | // gint person_class_id; 125 | // gint car_class_id; 126 | gint return_value; 127 | // guint index; 128 | // gint active_source_index; 129 | 130 | GMutex app_lock; 131 | GCond app_cond; 132 | 133 | NvDsPipeline pipeline; 134 | // NvDsConfig config; 135 | // NvDsConfig override_config; 136 | // NvDsInstanceData instance_data[MAX_SOURCE_BINS]; 137 | // NvDsC2DContext *c2d_ctx[MAX_MESSAGE_CONSUMERS]; 138 | // NvDsAppPerfStructInt perf_struct; 139 | // bbox_generated_callback bbox_generated_post_analytics_cb; 140 | // bbox_generated_callback all_bbox_generated_cb; 141 | // overlay_graphics_callback overlay_graphics_cb; 142 | // NvDsFrameLatencyInfo *latency_info; 143 | // GMutex latency_lock; 144 | // GThread *ota_handler_thread; 145 | // guint ota_inotify_fd; 146 | // guint ota_watch_desc; 147 | }; 148 | 149 | #ifdef __cplusplus 150 | } 151 | #endif 152 | 153 | #endif -------------------------------------------------------------------------------- /system_drawing.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/belarbi2733/deepstream-rtspsrc-yolo/d90fc9c325f01478cb7ffc26879bd89f5c7b6735/system_drawing.png -------------------------------------------------------------------------------- /yolo/README: -------------------------------------------------------------------------------- 1 | ################################################################################ 2 | # Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # NVIDIA Corporation and its licensors retain all intellectual property 5 | # and proprietary rights in and to this software, related documentation 6 | # and any modifications thereto. Any use, reproduction, disclosure or 7 | # distribution of this software and related documentation without an express 8 | # license agreement from NVIDIA Corporation is strictly prohibited. 9 | # 10 | ################################################################################ 11 | 12 | This sample shows how to integrate a detector Model like Yolo having layers 13 | supported through IPlugin interface by TensorRT and with custom output layer 14 | parsing for detected objects with DeepStreamSDK. The sample uses the 15 | IPluginFactory interface to create specific layers in Yolo and configure yolo 16 | plugin provided by TensorRT. This sample demonstrates how to reconstruct 17 | user-defined TensorRT layers from engine file in runtime. 18 | 19 | -------------------------------------------------------------------------------- 20 | Sample contents: 21 | - deepstream_app_config_yolo[V3,V3_tiny,V2,V2_tiny,tlt].txt - DeepStream reference 22 | app configuration file for using YoloV2/yoloV2-tiny/yolo/yolo-tiny/tlt model 23 | as the primary detector. 24 | - config_infer_primary_yolo[V3,V3_tiny,V2,V2_tiny,tlt].txt - Configuration file for the GStreamer 25 | nvinfer plugin for the Yolo detector model. 26 | - yolov3-calibration.table.trt7.0 - yoloV3 INT8 calibration binary on TensorRT 7.0+ 27 | - nvdsinfer_custom_impl_Yolo/nvdsinfer_yolo_engine.cpp - 28 | Implementation of 'NvDsInferCreateModelParser'/IModelParser for nvdsinfer to 29 | parse custom models. Alternatively, also contains implementation of 30 | 'NvDsInferYoloCudaEngineGet' for nvdsinfer to directly create cuda engine. 31 | To use the 'NvDsInferYoloCudaEngineGet' interface, enable the macro 32 | USE_CUDA_ENGINE_GET_API in nvdsinfer_yolo_engine.cpp 33 | - nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.cpp - Output layer 34 | parsing function for detected objects for the Yolo model. 35 | - nvdsinfer_custom_impl_Yolo/yoloPlugins.h - 36 | Declaration of YoloLayerV3 and YoloLayerV3PluginCreator. 37 | - nvdsinfer_custom_impl_Yolo/yoloPlugins.cpp - 38 | Implementation of YoloLayerV3 and YoloLayerV3PluginCreator. 39 | - nvdsinfer_custom_impl_Yolo/kernels.cu - Implementation of cuda kernels for 40 | - nvdsinfer_custom_impl_Yolo/trt_utils.h - Utilities to setup tensorRT networks 41 | - nvdsinfer_custom_impl_Yolo/trt_utils.cpp - Implementation of Utilities to setup 42 | tensorRT networks 43 | - nvdsinfer_custom_impl_Yolo/yolo.h - Interface to create Yolo Cuda-Engine 44 | - nvdsinfer_custom_impl_Yolo/yolo.cpp - Implementation to create Yolo Cuda-Engine 45 | 46 | -------------------------------------------------------------------------------- 47 | Pre-requisites: 48 | - Download yolo config and weights files 49 | $ ./prebuild.sh 50 | - Set correct yolo config/weights file in config_infer_primary_yolo[...].txt. 51 | custom-network-config # path to yolo config 52 | model-file # path to yolo weights 53 | - Enable INT8 precision detection if there is a calibration cache file, update 54 | config_infer_primary_yolo[...].txt. 55 | int8-calib-file=yolo[...]-calibration.table.trt5.1 56 | - Other INT8 precision calibration table need to be calibrated by user. 57 | - The yolo-tlt sample makes use of a few TensorRT OSS plugins. Download the 58 | TensorRT OSS repo (https://github.com/NVIDIA/TensorRT/) and checkout the 59 | 'release/7.0' branch. Follow the instructions in the README to build the 60 | plugin library 'libnvinfer_plugin.so.7.0.0' corresponding to your dGPU/Jetson 61 | platform. This library needs to be used with LD_PRELOAD to let nvinfer plugin 62 | access the TRT OSS plugin layers. 63 | - Set the path of the etlt model (tlt-encoded-model) and the tlt model key 64 | (tlt-model-key) in config_infer_primary_yolo_tlt.txt before running the sample. 65 | 66 | Note: 67 | - There is a quality regression of masks in yolov3-tiny.cfg, please see comments 68 | https://github.com/pjreddie/darknet/commit/f86901f6177dfc6116360a13cc06ab680e0c86b0#r30200016 69 | Before official fix in yolov3-tiny.weights, end users can revert the masks to 70 | [1, 2, 3]. There is also a hardcode of masks in nvdsparsebbox_Yolo.cpp. 71 | 72 | -------------------------------------------------------------------------------- 73 | Compile the custom library: 74 | # Based on the API to use 'NvDsInferCreateModelParser' or 'NvDsInferCudaEngineGet' 75 | # set the macro USE_CUDA_ENGINE_GET_API to 0 or 1 in 76 | # nvdsinfer_custom_impl_Yolo/nvdsinfer_yolo_engine.cpp 77 | 78 | # Export correct CUDA version (e.g. 10.2, 10.1) 79 | $ export CUDA_VER=10.2 80 | $ make -C nvdsinfer_custom_impl_Yolo 81 | 82 | -------------------------------------------------------------------------------- 83 | Run the sample: 84 | The "nvinfer" config file config_infer_primary_yolo.txt specifies the path to 85 | the custom library and the custom output parsing function through the properties 86 | "custom-lib-path" and "parse-bbox-func-name" respectively. 87 | The first-time a "model_b1_int8.engine" would be generated as the engine-file 88 | 89 | - With deepstream-app 90 | $ deepstream-app -c deepstream_app_config_yoloV3.txt 91 | $ deepstream-app -c deepstream_app_config_yoloV3_tiny.txt 92 | $ deepstream-app -c deepstream_app_config_yoloV2.txt 93 | $ deepstream-app -c deepstream_app_config_yoloV2_tiny.txt 94 | $ LD_PRELOAD= deepstream-app -c deepstream_app_config_yolo_tlt.txt 95 | -------------------------------------------------------------------------------- /yolo/nvdsinfer_custom_impl_Yolo/Makefile: -------------------------------------------------------------------------------- 1 | ################################################################################ 2 | # Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # Permission is hereby granted, free of charge, to any person obtaining a 5 | # copy of this software and associated documentation files (the "Software"), 6 | # to deal in the Software without restriction, including without limitation 7 | # the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | # and/or sell copies of the Software, and to permit persons to whom the 9 | # Software is furnished to do so, subject to the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be included in 12 | # all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | # DEALINGS IN THE SOFTWARE. 21 | ################################################################################ 22 | 23 | CUDA_VER:=10.2 24 | NVDS_VERSION:=5.0 25 | 26 | CC:= g++ 27 | NVCC:=/usr/local/cuda-$(CUDA_VER)/bin/nvcc 28 | 29 | CFLAGS:= -Wall -std=c++11 -shared -fPIC -Wno-error=deprecated-declarations 30 | CFLAGS+= -I/opt/nvidia/deepstream/deepstream-$(NVDS_VERSION)/sources/includes -I/usr/local/cuda-$(CUDA_VER)/include 31 | 32 | LIBS:= -lnvinfer_plugin -lnvinfer -lnvparsers -L/usr/local/cuda-$(CUDA_VER)/lib64 -lcudart -lcublas -lstdc++fs 33 | LFLAGS:= -shared -Wl,--start-group $(LIBS) -Wl,--end-group 34 | 35 | INCS:= $(wildcard *.h) 36 | SRCFILES:= nvdsinfer_yolo_engine.cpp \ 37 | nvdsparsebbox_Yolo.cpp \ 38 | yoloPlugins.cpp \ 39 | trt_utils.cpp \ 40 | yolo.cpp \ 41 | kernels.cu 42 | TARGET_LIB:= libnvdsinfer_custom_impl_Yolo.so 43 | 44 | TARGET_OBJS:= $(SRCFILES:.cpp=.o) 45 | TARGET_OBJS:= $(TARGET_OBJS:.cu=.o) 46 | 47 | all: $(TARGET_LIB) 48 | 49 | %.o: %.cpp $(INCS) Makefile 50 | $(CC) -c -o $@ $(CFLAGS) $< 51 | 52 | %.o: %.cu $(INCS) Makefile 53 | $(NVCC) -c -o $@ --compiler-options '-fPIC' $< 54 | 55 | $(TARGET_LIB) : $(TARGET_OBJS) 56 | $(CC) -o $@ $(TARGET_OBJS) $(LFLAGS) 57 | 58 | clean: 59 | rm -rf $(TARGET_LIB) 60 | -------------------------------------------------------------------------------- /yolo/nvdsinfer_custom_impl_Yolo/kernels.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018-2019 NVIDIA Corporation. All rights reserved. 3 | * 4 | * NVIDIA Corporation and its licensors retain all intellectual property 5 | * and proprietary rights in and to this software, related documentation 6 | * and any modifications thereto. Any use, reproduction, disclosure or 7 | * distribution of this software and related documentation without an express 8 | * license agreement from NVIDIA Corporation is strictly prohibited. 9 | * 10 | */ 11 | 12 | #include 13 | #include 14 | #include 15 | #include 16 | #include 17 | 18 | inline __device__ float sigmoidGPU(const float& x) { return 1.0f / (1.0f + __expf(-x)); } 19 | 20 | __global__ void gpuYoloLayerV3(const float* input, float* output, const uint gridSize, const uint numOutputClasses, 21 | const uint numBBoxes) 22 | { 23 | uint x_id = blockIdx.x * blockDim.x + threadIdx.x; 24 | uint y_id = blockIdx.y * blockDim.y + threadIdx.y; 25 | uint z_id = blockIdx.z * blockDim.z + threadIdx.z; 26 | 27 | if ((x_id >= gridSize) || (y_id >= gridSize) || (z_id >= numBBoxes)) 28 | { 29 | return; 30 | } 31 | 32 | const int numGridCells = gridSize * gridSize; 33 | const int bbindex = y_id * gridSize + x_id; 34 | 35 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)] 36 | = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)]); 37 | 38 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)] 39 | = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)]); 40 | 41 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)] 42 | = __expf(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)]); 43 | 44 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)] 45 | = __expf(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)]); 46 | 47 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)] 48 | = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)]); 49 | 50 | for (uint i = 0; i < numOutputClasses; ++i) 51 | { 52 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))] 53 | = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))]); 54 | } 55 | } 56 | 57 | cudaError_t cudaYoloLayerV3(const void* input, void* output, const uint& batchSize, const uint& gridSize, 58 | const uint& numOutputClasses, const uint& numBBoxes, 59 | uint64_t outputSize, cudaStream_t stream); 60 | 61 | cudaError_t cudaYoloLayerV3(const void* input, void* output, const uint& batchSize, const uint& gridSize, 62 | const uint& numOutputClasses, const uint& numBBoxes, 63 | uint64_t outputSize, cudaStream_t stream) 64 | { 65 | dim3 threads_per_block(16, 16, 4); 66 | dim3 number_of_blocks((gridSize / threads_per_block.x) + 1, 67 | (gridSize / threads_per_block.y) + 1, 68 | (numBBoxes / threads_per_block.z) + 1); 69 | for (unsigned int batch = 0; batch < batchSize; ++batch) 70 | { 71 | gpuYoloLayerV3<<>>( 72 | reinterpret_cast(input) + (batch * outputSize), 73 | reinterpret_cast(output) + (batch * outputSize), gridSize, numOutputClasses, 74 | numBBoxes); 75 | } 76 | return cudaGetLastError(); 77 | } 78 | -------------------------------------------------------------------------------- /yolo/nvdsinfer_custom_impl_Yolo/nvdsinfer_yolo_engine.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | */ 22 | 23 | #include "nvdsinfer_custom_impl.h" 24 | #include "nvdsinfer_context.h" 25 | #include "yoloPlugins.h" 26 | #include "yolo.h" 27 | 28 | #include 29 | 30 | #define USE_CUDA_ENGINE_GET_API 1 31 | 32 | static bool getYoloNetworkInfo (NetworkInfo &networkInfo, const NvDsInferContextInitParams* initParams) 33 | { 34 | std::string yoloCfg = initParams->customNetworkConfigFilePath; 35 | std::string yoloType; 36 | 37 | std::transform (yoloCfg.begin(), yoloCfg.end(), yoloCfg.begin(), [] (uint8_t c) { 38 | return std::tolower (c);}); 39 | 40 | if (yoloCfg.find("yolov2") != std::string::npos) { 41 | if (yoloCfg.find("yolov2-tiny") != std::string::npos) 42 | yoloType = "yolov2-tiny"; 43 | else 44 | yoloType = "yolov2"; 45 | } else if (yoloCfg.find("yolov3") != std::string::npos) { 46 | if (yoloCfg.find("yolov3-tiny") != std::string::npos) 47 | yoloType = "yolov3-tiny"; 48 | else 49 | yoloType = "yolov3"; 50 | } else { 51 | std::cerr << "Yolo type is not defined from config file name:" 52 | << yoloCfg << std::endl; 53 | return false; 54 | } 55 | 56 | networkInfo.networkType = yoloType; 57 | networkInfo.configFilePath = initParams->customNetworkConfigFilePath; 58 | networkInfo.wtsFilePath = initParams->modelFilePath; 59 | networkInfo.deviceType = (initParams->useDLA ? "kDLA" : "kGPU"); 60 | networkInfo.inputBlobName = "data"; 61 | 62 | if (networkInfo.configFilePath.empty() || 63 | networkInfo.wtsFilePath.empty()) { 64 | std::cerr << "Yolo config file or weights file is NOT specified." 65 | << std::endl; 66 | return false; 67 | } 68 | 69 | if (!fileExists(networkInfo.configFilePath) || 70 | !fileExists(networkInfo.wtsFilePath)) { 71 | std::cerr << "Yolo config file or weights file is NOT exist." 72 | << std::endl; 73 | return false; 74 | } 75 | 76 | return true; 77 | } 78 | 79 | #if !USE_CUDA_ENGINE_GET_API 80 | IModelParser* NvDsInferCreateModelParser( 81 | const NvDsInferContextInitParams* initParams) { 82 | NetworkInfo networkInfo; 83 | if (!getYoloNetworkInfo(networkInfo, initParams)) { 84 | return nullptr; 85 | } 86 | 87 | return new Yolo(networkInfo); 88 | } 89 | #else 90 | extern "C" 91 | bool NvDsInferYoloCudaEngineGet(nvinfer1::IBuilder * const builder, 92 | const NvDsInferContextInitParams * const initParams, 93 | nvinfer1::DataType dataType, 94 | nvinfer1::ICudaEngine *& cudaEngine); 95 | 96 | extern "C" 97 | bool NvDsInferYoloCudaEngineGet(nvinfer1::IBuilder * const builder, 98 | const NvDsInferContextInitParams * const initParams, 99 | nvinfer1::DataType dataType, 100 | nvinfer1::ICudaEngine *& cudaEngine) 101 | { 102 | NetworkInfo networkInfo; 103 | if (!getYoloNetworkInfo(networkInfo, initParams)) { 104 | return false; 105 | } 106 | 107 | Yolo yolo(networkInfo); 108 | cudaEngine = yolo.createEngine (builder); 109 | if (cudaEngine == nullptr) 110 | { 111 | std::cerr << "Failed to build cuda engine on " 112 | << networkInfo.configFilePath << std::endl; 113 | return false; 114 | } 115 | 116 | return true; 117 | } 118 | #endif 119 | -------------------------------------------------------------------------------- /yolo/nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | */ 22 | 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | #include "nvdsinfer_custom_impl.h" 31 | #include "trt_utils.h" 32 | 33 | static const int NUM_CLASSES_YOLO = 10; 34 | 35 | extern "C" bool NvDsInferParseCustomYoloV3( 36 | std::vector const& outputLayersInfo, 37 | NvDsInferNetworkInfo const& networkInfo, 38 | NvDsInferParseDetectionParams const& detectionParams, 39 | std::vector& objectList); 40 | 41 | extern "C" bool NvDsInferParseCustomYoloV3Tiny( 42 | std::vector const& outputLayersInfo, 43 | NvDsInferNetworkInfo const& networkInfo, 44 | NvDsInferParseDetectionParams const& detectionParams, 45 | std::vector& objectList); 46 | 47 | extern "C" bool NvDsInferParseCustomYoloV2( 48 | std::vector const& outputLayersInfo, 49 | NvDsInferNetworkInfo const& networkInfo, 50 | NvDsInferParseDetectionParams const& detectionParams, 51 | std::vector& objectList); 52 | 53 | extern "C" bool NvDsInferParseCustomYoloV2Tiny( 54 | std::vector const& outputLayersInfo, 55 | NvDsInferNetworkInfo const& networkInfo, 56 | NvDsInferParseDetectionParams const& detectionParams, 57 | std::vector& objectList); 58 | 59 | extern "C" bool NvDsInferParseCustomYoloTLT( 60 | std::vector const& outputLayersInfo, 61 | NvDsInferNetworkInfo const& networkInfo, 62 | NvDsInferParseDetectionParams const& detectionParams, 63 | std::vector& objectList); 64 | 65 | extern "C" bool NvDsInferParseCustomYoloV4( 66 | std::vector const& outputLayersInfo, 67 | NvDsInferNetworkInfo const& networkInfo, 68 | NvDsInferParseDetectionParams const& detectionParams, 69 | std::vector& objectList); 70 | 71 | /* This is a sample bounding box parsing function for the sample YoloV3 detector model */ 72 | static NvDsInferParseObjectInfo convertBBox(const float& bx, const float& by, const float& bw, 73 | const float& bh, const int& stride, const uint& netW, 74 | const uint& netH) 75 | { 76 | NvDsInferParseObjectInfo b; 77 | // Restore coordinates to network input resolution 78 | float xCenter = bx * stride; 79 | float yCenter = by * stride; 80 | float x0 = xCenter - bw / 2; 81 | float y0 = yCenter - bh / 2; 82 | float x1 = x0 + bw; 83 | float y1 = y0 + bh; 84 | 85 | x0 = clamp(x0, 0, netW); 86 | y0 = clamp(y0, 0, netH); 87 | x1 = clamp(x1, 0, netW); 88 | y1 = clamp(y1, 0, netH); 89 | 90 | b.left = x0; 91 | b.width = clamp(x1 - x0, 0, netW); 92 | b.top = y0; 93 | b.height = clamp(y1 - y0, 0, netH); 94 | 95 | return b; 96 | } 97 | 98 | static void addBBoxProposal(const float bx, const float by, const float bw, const float bh, 99 | const uint stride, const uint& netW, const uint& netH, const int maxIndex, 100 | const float maxProb, std::vector& binfo) 101 | { 102 | NvDsInferParseObjectInfo bbi = convertBBox(bx, by, bw, bh, stride, netW, netH); 103 | if (bbi.width < 1 || bbi.height < 1) return; 104 | 105 | bbi.detectionConfidence = maxProb; 106 | bbi.classId = maxIndex; 107 | binfo.push_back(bbi); 108 | } 109 | 110 | static std::vector 111 | decodeYoloV2Tensor( 112 | const float* detections, const std::vector &anchors, 113 | const uint gridSizeW, const uint gridSizeH, const uint stride, const uint numBBoxes, 114 | const uint numOutputClasses, const uint& netW, 115 | const uint& netH) 116 | { 117 | std::vector binfo; 118 | for (uint y = 0; y < gridSizeH; ++y) { 119 | for (uint x = 0; x < gridSizeW; ++x) { 120 | for (uint b = 0; b < numBBoxes; ++b) 121 | { 122 | const float pw = anchors[b * 2]; 123 | const float ph = anchors[b * 2 + 1]; 124 | 125 | const int numGridCells = gridSizeH * gridSizeW; 126 | const int bbindex = y * gridSizeW + x; 127 | const float bx 128 | = x + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 0)]; 129 | const float by 130 | = y + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 1)]; 131 | const float bw 132 | = pw * exp (detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 2)]); 133 | const float bh 134 | = ph * exp (detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 3)]); 135 | 136 | const float objectness 137 | = detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 4)]; 138 | 139 | float maxProb = 0.0f; 140 | int maxIndex = -1; 141 | 142 | for (uint i = 0; i < numOutputClasses; ++i) 143 | { 144 | float prob 145 | = (detections[bbindex 146 | + numGridCells * (b * (5 + numOutputClasses) + (5 + i))]); 147 | 148 | if (prob > maxProb) 149 | { 150 | maxProb = prob; 151 | maxIndex = i; 152 | } 153 | } 154 | maxProb = objectness * maxProb; 155 | 156 | addBBoxProposal(bx, by, bw, bh, stride, netW, netH, maxIndex, maxProb, binfo); 157 | } 158 | } 159 | } 160 | return binfo; 161 | } 162 | 163 | static std::vector 164 | decodeYoloV3Tensor( 165 | const float* detections, const std::vector &mask, const std::vector &anchors, 166 | const uint gridSizeW, const uint gridSizeH, const uint stride, const uint numBBoxes, 167 | const uint numOutputClasses, const uint& netW, 168 | const uint& netH) 169 | { 170 | std::vector binfo; 171 | for (uint y = 0; y < gridSizeH; ++y) { 172 | for (uint x = 0; x < gridSizeW; ++x) { 173 | for (uint b = 0; b < numBBoxes; ++b) 174 | { 175 | const float pw = anchors[mask[b] * 2]; 176 | const float ph = anchors[mask[b] * 2 + 1]; 177 | 178 | const int numGridCells = gridSizeH * gridSizeW; 179 | const int bbindex = y * gridSizeW + x; 180 | const float bx 181 | = x + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 0)]; 182 | const float by 183 | = y + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 1)]; 184 | const float bw 185 | = pw * detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 2)]; 186 | const float bh 187 | = ph * detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 3)]; 188 | 189 | const float objectness 190 | = detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 4)]; 191 | 192 | float maxProb = 0.0f; 193 | int maxIndex = -1; 194 | 195 | for (uint i = 0; i < numOutputClasses; ++i) 196 | { 197 | float prob 198 | = (detections[bbindex 199 | + numGridCells * (b * (5 + numOutputClasses) + (5 + i))]); 200 | 201 | if (prob > maxProb) 202 | { 203 | maxProb = prob; 204 | maxIndex = i; 205 | } 206 | } 207 | maxProb = objectness * maxProb; 208 | 209 | addBBoxProposal(bx, by, bw, bh, stride, netW, netH, maxIndex, maxProb, binfo); 210 | } 211 | } 212 | } 213 | return binfo; 214 | } 215 | 216 | static inline std::vector 217 | SortLayers(const std::vector & outputLayersInfo) 218 | { 219 | std::vector outLayers; 220 | for (auto const &layer : outputLayersInfo) { 221 | outLayers.push_back (&layer); 222 | } 223 | std::sort(outLayers.begin(), outLayers.end(), 224 | [](const NvDsInferLayerInfo* a, const NvDsInferLayerInfo* b) { 225 | return a->inferDims.d[1] < b->inferDims.d[1]; 226 | }); 227 | return outLayers; 228 | } 229 | 230 | static bool NvDsInferParseYoloV3( 231 | std::vector const& outputLayersInfo, 232 | NvDsInferNetworkInfo const& networkInfo, 233 | NvDsInferParseDetectionParams const& detectionParams, 234 | std::vector& objectList, 235 | const std::vector &anchors, 236 | const std::vector> &masks) 237 | { 238 | const uint kNUM_BBOXES = 3; 239 | 240 | const std::vector sortedLayers = 241 | SortLayers (outputLayersInfo); 242 | 243 | if (sortedLayers.size() != masks.size()) { 244 | std::cerr << "ERROR: yoloV3 output layer.size: " << sortedLayers.size() 245 | << " does not match mask.size: " << masks.size() << std::endl; 246 | return false; 247 | } 248 | 249 | if (NUM_CLASSES_YOLO != detectionParams.numClassesConfigured) 250 | { 251 | std::cerr << "WARNING: Num classes mismatch. Configured:" 252 | << detectionParams.numClassesConfigured 253 | << ", detected by network: " << NUM_CLASSES_YOLO << std::endl; 254 | } 255 | 256 | std::vector objects; 257 | 258 | for (uint idx = 0; idx < masks.size(); ++idx) { 259 | const NvDsInferLayerInfo &layer = *sortedLayers[idx]; // 255 x Grid x Grid 260 | 261 | assert(layer.inferDims.numDims == 3); 262 | const uint gridSizeH = layer.inferDims.d[1]; 263 | const uint gridSizeW = layer.inferDims.d[2]; 264 | const uint stride = DIVUP(networkInfo.width, gridSizeW); 265 | assert(stride == DIVUP(networkInfo.height, gridSizeH)); 266 | 267 | std::vector outObjs = 268 | decodeYoloV3Tensor((const float*)(layer.buffer), masks[idx], anchors, gridSizeW, gridSizeH, stride, kNUM_BBOXES, 269 | NUM_CLASSES_YOLO, networkInfo.width, networkInfo.height); 270 | objects.insert(objects.end(), outObjs.begin(), outObjs.end()); 271 | } 272 | 273 | 274 | objectList = objects; 275 | 276 | return true; 277 | } 278 | 279 | static NvDsInferParseObjectInfo convertBBoxYoloV4(const float& bx1, const float& by1, const float& bx2, 280 | const float& by2, const uint& netW, const uint& netH) 281 | { 282 | NvDsInferParseObjectInfo b; 283 | // Restore coordinates to network input resolution 284 | 285 | float x1 = bx1 * netW; 286 | float y1 = by1 * netH; 287 | float x2 = bx2 * netW; 288 | float y2 = by2 * netH; 289 | 290 | x1 = clamp(x1, 0, netW); 291 | y1 = clamp(y1, 0, netH); 292 | x2 = clamp(x2, 0, netW); 293 | y2 = clamp(y2, 0, netH); 294 | 295 | b.left = x1; 296 | b.width = clamp(x2 - x1, 0, netW); 297 | b.top = y1; 298 | b.height = clamp(y2 - y1, 0, netH); 299 | 300 | return b; 301 | } 302 | 303 | static void addBBoxProposalYoloV4(const float bx, const float by, const float bw, const float bh, 304 | const uint& netW, const uint& netH, const int maxIndex, 305 | const float maxProb, std::vector& binfo) 306 | { 307 | NvDsInferParseObjectInfo bbi = convertBBoxYoloV4(bx, by, bw, bh, netW, netH); 308 | if (bbi.width < 1 || bbi.height < 1) return; 309 | 310 | bbi.detectionConfidence = maxProb; 311 | bbi.classId = maxIndex; 312 | binfo.push_back(bbi); 313 | } 314 | 315 | static std::vector 316 | decodeYoloV4Tensor( 317 | const float* boxes, const float* scores, 318 | const uint num_bboxes, NvDsInferParseDetectionParams const& detectionParams, 319 | const uint& netW, const uint& netH) 320 | { 321 | std::vector binfo; 322 | 323 | uint bbox_location = 0; 324 | uint score_location = 0; 325 | for (uint b = 0; b < num_bboxes; ++b) 326 | { 327 | float bx1 = boxes[bbox_location]; 328 | float by1 = boxes[bbox_location + 1]; 329 | float bx2 = boxes[bbox_location + 2]; 330 | float by2 = boxes[bbox_location + 3]; 331 | 332 | float maxProb = 0.0f; 333 | int maxIndex = -1; 334 | 335 | for (uint c = 0; c < detectionParams.numClassesConfigured; ++c) 336 | { 337 | float prob = scores[score_location + c]; 338 | if (prob > maxProb) 339 | { 340 | maxProb = prob; 341 | maxIndex = c; 342 | } 343 | } 344 | 345 | if (maxProb > detectionParams.perClassPreclusterThreshold[maxIndex]) 346 | { 347 | addBBoxProposalYoloV4(bx1, by1, bx2, by2, netW, netH, maxIndex, maxProb, binfo); 348 | } 349 | 350 | bbox_location += 4; 351 | score_location += detectionParams.numClassesConfigured; 352 | } 353 | 354 | return binfo; 355 | } 356 | 357 | 358 | /* C-linkage to prevent name-mangling */ 359 | 360 | static bool NvDsInferParseYoloV4( 361 | std::vector const& outputLayersInfo, 362 | NvDsInferNetworkInfo const& networkInfo, 363 | NvDsInferParseDetectionParams const& detectionParams, 364 | std::vector& objectList) 365 | { 366 | if (NUM_CLASSES_YOLO != detectionParams.numClassesConfigured) 367 | { 368 | std::cerr << "WARNING: Num classes mismatch. Configured:" 369 | << detectionParams.numClassesConfigured 370 | << ", detected by network: " << NUM_CLASSES_YOLO << std::endl; 371 | } 372 | 373 | std::vector objects; 374 | 375 | const NvDsInferLayerInfo &boxes = outputLayersInfo[0]; // num_boxes x 4 376 | const NvDsInferLayerInfo &scores = outputLayersInfo[1]; // num_boxes x num_classes 377 | const NvDsInferLayerInfo &subbox = outputLayersInfo[2]; 378 | //* printf("%d\n", subbox.inferDims.numDims); 379 | // 3 dimensional: [num_boxes, 1, 4] 380 | assert(boxes.inferDims.numDims == 3); 381 | // 2 dimensional: [num_boxes, num_classes] 382 | assert(scores.inferDims.numDims == 2); 383 | 384 | // The second dimension should be num_classes 385 | assert(detectionParams.numClassesConfigured == scores.inferDims.d[1]); 386 | 387 | uint num_bboxes = boxes.inferDims.d[0]; 388 | 389 | // std::cout << "Network Info: " << networkInfo.height << " " << networkInfo.width << std::endl; 390 | 391 | std::vector outObjs = 392 | decodeYoloV4Tensor( 393 | (const float*)(boxes.buffer), (const float*)(scores.buffer), num_bboxes, detectionParams, 394 | networkInfo.width, networkInfo.height); 395 | 396 | objects.insert(objects.end(), outObjs.begin(), outObjs.end()); 397 | 398 | objectList = objects; 399 | 400 | return true; 401 | } 402 | 403 | extern "C" bool NvDsInferParseCustomYoloV4( 404 | std::vector const& outputLayersInfo, 405 | NvDsInferNetworkInfo const& networkInfo, 406 | NvDsInferParseDetectionParams const& detectionParams, 407 | std::vector& objectList) 408 | { 409 | return NvDsInferParseYoloV4 ( 410 | outputLayersInfo, networkInfo, detectionParams, objectList); 411 | } 412 | 413 | extern "C" bool NvDsInferParseCustomYoloV3( 414 | std::vector const& outputLayersInfo, 415 | NvDsInferNetworkInfo const& networkInfo, 416 | NvDsInferParseDetectionParams const& detectionParams, 417 | std::vector& objectList) 418 | { 419 | static const std::vector kANCHORS = { 420 | 10.0, 13.0, 16.0, 30.0, 33.0, 23.0, 30.0, 61.0, 62.0, 421 | 45.0, 59.0, 119.0, 116.0, 90.0, 156.0, 198.0, 373.0, 326.0}; 422 | static const std::vector> kMASKS = { 423 | {6, 7, 8}, 424 | {3, 4, 5}, 425 | {0, 1, 2}}; 426 | return NvDsInferParseYoloV3 ( 427 | outputLayersInfo, networkInfo, detectionParams, objectList, 428 | kANCHORS, kMASKS); 429 | } 430 | 431 | extern "C" bool NvDsInferParseCustomYoloV3Tiny( 432 | std::vector const& outputLayersInfo, 433 | NvDsInferNetworkInfo const& networkInfo, 434 | NvDsInferParseDetectionParams const& detectionParams, 435 | std::vector& objectList) 436 | { 437 | static const std::vector kANCHORS = { 438 | 10, 14, 23, 27, 37, 58, 81, 82, 135, 169, 344, 319}; 439 | static const std::vector> kMASKS = { 440 | {3, 4, 5}, 441 | //{0, 1, 2}}; // as per output result, select {1,2,3} 442 | {1, 2, 3}}; 443 | 444 | return NvDsInferParseYoloV3 ( 445 | outputLayersInfo, networkInfo, detectionParams, objectList, 446 | kANCHORS, kMASKS); 447 | } 448 | 449 | static bool NvDsInferParseYoloV2( 450 | std::vector const& outputLayersInfo, 451 | NvDsInferNetworkInfo const& networkInfo, 452 | NvDsInferParseDetectionParams const& detectionParams, 453 | std::vector& objectList) 454 | { 455 | // copy anchor data from yolov2.cfg file 456 | std::vector anchors = {0.57273, 0.677385, 1.87446, 2.06253, 3.33843, 457 | 5.47434, 7.88282, 3.52778, 9.77052, 9.16828}; 458 | const uint kNUM_BBOXES = 5; 459 | 460 | if (outputLayersInfo.empty()) { 461 | std::cerr << "Could not find output layer in bbox parsing" << std::endl;; 462 | return false; 463 | } 464 | const NvDsInferLayerInfo &layer = outputLayersInfo[0]; 465 | 466 | if (NUM_CLASSES_YOLO != detectionParams.numClassesConfigured) 467 | { 468 | std::cerr << "WARNING: Num classes mismatch. Configured:" 469 | << detectionParams.numClassesConfigured 470 | << ", detected by network: " << NUM_CLASSES_YOLO << std::endl; 471 | } 472 | 473 | assert(layer.inferDims.numDims == 3); 474 | const uint gridSizeH = layer.inferDims.d[1]; 475 | const uint gridSizeW = layer.inferDims.d[2]; 476 | const uint stride = DIVUP(networkInfo.width, gridSizeW); 477 | assert(stride == DIVUP(networkInfo.height, gridSizeH)); 478 | for (auto& anchor : anchors) { 479 | anchor *= stride; 480 | } 481 | std::vector objects = 482 | decodeYoloV2Tensor((const float*)(layer.buffer), anchors, gridSizeW, gridSizeH, stride, kNUM_BBOXES, 483 | NUM_CLASSES_YOLO, networkInfo.width, networkInfo.height); 484 | 485 | objectList = objects; 486 | 487 | return true; 488 | } 489 | 490 | extern "C" bool NvDsInferParseCustomYoloV2( 491 | std::vector const& outputLayersInfo, 492 | NvDsInferNetworkInfo const& networkInfo, 493 | NvDsInferParseDetectionParams const& detectionParams, 494 | std::vector& objectList) 495 | { 496 | return NvDsInferParseYoloV2 ( 497 | outputLayersInfo, networkInfo, detectionParams, objectList); 498 | } 499 | 500 | extern "C" bool NvDsInferParseCustomYoloV2Tiny( 501 | std::vector const& outputLayersInfo, 502 | NvDsInferNetworkInfo const& networkInfo, 503 | NvDsInferParseDetectionParams const& detectionParams, 504 | std::vector& objectList) 505 | { 506 | return NvDsInferParseYoloV2 ( 507 | outputLayersInfo, networkInfo, detectionParams, objectList); 508 | } 509 | 510 | extern "C" bool NvDsInferParseCustomYoloTLT( 511 | std::vector const& outputLayersInfo, 512 | NvDsInferNetworkInfo const& networkInfo, 513 | NvDsInferParseDetectionParams const& detectionParams, 514 | std::vector& objectList) 515 | { 516 | 517 | if(outputLayersInfo.size() != 4) 518 | { 519 | std::cerr << "Mismatch in the number of output buffers." 520 | << "Expected 4 output buffers, detected in the network :" 521 | << outputLayersInfo.size() << std::endl; 522 | return false; 523 | } 524 | 525 | const int topK = 200; 526 | const int* keepCount = static_cast (outputLayersInfo.at(0).buffer); 527 | const float* boxes = static_cast (outputLayersInfo.at(1).buffer); 528 | const float* scores = static_cast (outputLayersInfo.at(2).buffer); 529 | const float* cls = static_cast (outputLayersInfo.at(3).buffer); 530 | 531 | for (int i = 0; (i < keepCount[0]) && (objectList.size() <= topK); ++i) 532 | { 533 | const float* loc = &boxes[0] + (i * 4); 534 | const float* conf = &scores[0] + i; 535 | const float* cls_id = &cls[0] + i; 536 | 537 | if(conf[0] > 1.001) 538 | continue; 539 | 540 | if((loc[0] < 0) || (loc[1] < 0) || (loc[2] < 0) || (loc[3] < 0)) 541 | continue; 542 | 543 | if((loc[0] > networkInfo.width) || (loc[2] > networkInfo.width) || (loc[1] > networkInfo.height) || (loc[3] > networkInfo.width)) 544 | continue; 545 | 546 | if((loc[2] < loc[0]) || (loc[3] < loc[1])) 547 | continue; 548 | 549 | if(((loc[3] - loc[1]) > networkInfo.height) || ((loc[2]-loc[0]) > networkInfo.width)) 550 | continue; 551 | 552 | NvDsInferParseObjectInfo curObj{static_cast(cls_id[0]), 553 | loc[0],loc[1],(loc[2]-loc[0]), 554 | (loc[3]-loc[1]), conf[0]}; 555 | objectList.push_back(curObj); 556 | 557 | } 558 | 559 | return true; 560 | } 561 | 562 | /* Check that the custom function has been defined correctly */ 563 | CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV4); 564 | CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV3); 565 | CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV3Tiny); 566 | CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV2); 567 | CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV2Tiny); 568 | CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloTLT); 569 | -------------------------------------------------------------------------------- /yolo/nvdsinfer_custom_impl_Yolo/trt_utils.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | */ 22 | 23 | #include "trt_utils.h" 24 | 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | #include 31 | 32 | #include "NvInferPlugin.h" 33 | 34 | static void leftTrim(std::string& s) 35 | { 36 | s.erase(s.begin(), std::find_if(s.begin(), s.end(), [](int ch) { return !isspace(ch); })); 37 | } 38 | 39 | static void rightTrim(std::string& s) 40 | { 41 | s.erase(std::find_if(s.rbegin(), s.rend(), [](int ch) { return !isspace(ch); }).base(), s.end()); 42 | } 43 | 44 | std::string trim(std::string s) 45 | { 46 | leftTrim(s); 47 | rightTrim(s); 48 | return s; 49 | } 50 | 51 | float clamp(const float val, const float minVal, const float maxVal) 52 | { 53 | assert(minVal <= maxVal); 54 | return std::min(maxVal, std::max(minVal, val)); 55 | } 56 | 57 | bool fileExists(const std::string fileName, bool verbose) 58 | { 59 | if (!std::experimental::filesystem::exists(std::experimental::filesystem::path(fileName))) 60 | { 61 | if (verbose) std::cout << "File does not exist : " << fileName << std::endl; 62 | return false; 63 | } 64 | return true; 65 | } 66 | 67 | std::vector loadWeights(const std::string weightsFilePath, const std::string& networkType) 68 | { 69 | assert(fileExists(weightsFilePath)); 70 | std::cout << "Loading pre-trained weights..." << std::endl; 71 | std::ifstream file(weightsFilePath, std::ios_base::binary); 72 | assert(file.good()); 73 | std::string line; 74 | 75 | if (networkType == "yolov2") 76 | { 77 | // Remove 4 int32 bytes of data from the stream belonging to the header 78 | file.ignore(4 * 4); 79 | } 80 | else if ((networkType == "yolov3") || (networkType == "yolov3-tiny") 81 | || (networkType == "yolov2-tiny")) 82 | { 83 | // Remove 5 int32 bytes of data from the stream belonging to the header 84 | file.ignore(4 * 5); 85 | } 86 | else 87 | { 88 | std::cout << "Invalid network type" << std::endl; 89 | assert(0); 90 | } 91 | 92 | std::vector weights; 93 | char floatWeight[4]; 94 | while (!file.eof()) 95 | { 96 | file.read(floatWeight, 4); 97 | assert(file.gcount() == 4); 98 | weights.push_back(*reinterpret_cast(floatWeight)); 99 | if (file.peek() == std::istream::traits_type::eof()) break; 100 | } 101 | std::cout << "Loading weights of " << networkType << " complete!" 102 | << std::endl; 103 | std::cout << "Total Number of weights read : " << weights.size() << std::endl; 104 | return weights; 105 | } 106 | 107 | std::string dimsToString(const nvinfer1::Dims d) 108 | { 109 | std::stringstream s; 110 | assert(d.nbDims >= 1); 111 | for (int i = 0; i < d.nbDims - 1; ++i) 112 | { 113 | s << std::setw(4) << d.d[i] << " x"; 114 | } 115 | s << std::setw(4) << d.d[d.nbDims - 1]; 116 | 117 | return s.str(); 118 | } 119 | 120 | void displayDimType(const nvinfer1::Dims d) 121 | { 122 | std::cout << "(" << d.nbDims << ") "; 123 | for (int i = 0; i < d.nbDims; ++i) 124 | { 125 | switch (d.type[i]) 126 | { 127 | case nvinfer1::DimensionType::kSPATIAL: std::cout << "kSPATIAL "; break; 128 | case nvinfer1::DimensionType::kCHANNEL: std::cout << "kCHANNEL "; break; 129 | case nvinfer1::DimensionType::kINDEX: std::cout << "kINDEX "; break; 130 | case nvinfer1::DimensionType::kSEQUENCE: std::cout << "kSEQUENCE "; break; 131 | } 132 | } 133 | std::cout << std::endl; 134 | } 135 | 136 | int getNumChannels(nvinfer1::ITensor* t) 137 | { 138 | nvinfer1::Dims d = t->getDimensions(); 139 | assert(d.nbDims == 3); 140 | 141 | return d.d[0]; 142 | } 143 | 144 | uint64_t get3DTensorVolume(nvinfer1::Dims inputDims) 145 | { 146 | assert(inputDims.nbDims == 3); 147 | return inputDims.d[0] * inputDims.d[1] * inputDims.d[2]; 148 | } 149 | 150 | nvinfer1::ILayer* netAddMaxpool(int layerIdx, std::map& block, 151 | nvinfer1::ITensor* input, nvinfer1::INetworkDefinition* network) 152 | { 153 | assert(block.at("type") == "maxpool"); 154 | assert(block.find("size") != block.end()); 155 | assert(block.find("stride") != block.end()); 156 | 157 | int size = std::stoi(block.at("size")); 158 | int stride = std::stoi(block.at("stride")); 159 | 160 | nvinfer1::IPoolingLayer* pool 161 | = network->addPooling(*input, nvinfer1::PoolingType::kMAX, nvinfer1::DimsHW{size, size}); 162 | assert(pool); 163 | std::string maxpoolLayerName = "maxpool_" + std::to_string(layerIdx); 164 | pool->setStride(nvinfer1::DimsHW{stride, stride}); 165 | pool->setPaddingMode(nvinfer1::PaddingMode::kSAME_UPPER); 166 | pool->setName(maxpoolLayerName.c_str()); 167 | 168 | return pool; 169 | } 170 | 171 | nvinfer1::ILayer* netAddConvLinear(int layerIdx, std::map& block, 172 | std::vector& weights, 173 | std::vector& trtWeights, int& weightPtr, 174 | int& inputChannels, nvinfer1::ITensor* input, 175 | nvinfer1::INetworkDefinition* network) 176 | { 177 | assert(block.at("type") == "convolutional"); 178 | assert(block.find("batch_normalize") == block.end()); 179 | assert(block.at("activation") == "linear"); 180 | assert(block.find("filters") != block.end()); 181 | assert(block.find("pad") != block.end()); 182 | assert(block.find("size") != block.end()); 183 | assert(block.find("stride") != block.end()); 184 | 185 | int filters = std::stoi(block.at("filters")); 186 | int padding = std::stoi(block.at("pad")); 187 | int kernelSize = std::stoi(block.at("size")); 188 | int stride = std::stoi(block.at("stride")); 189 | int pad; 190 | if (padding) 191 | pad = (kernelSize - 1) / 2; 192 | else 193 | pad = 0; 194 | // load the convolution layer bias 195 | nvinfer1::Weights convBias{nvinfer1::DataType::kFLOAT, nullptr, filters}; 196 | float* val = new float[filters]; 197 | for (int i = 0; i < filters; ++i) 198 | { 199 | val[i] = weights[weightPtr]; 200 | weightPtr++; 201 | } 202 | convBias.values = val; 203 | trtWeights.push_back(convBias); 204 | // load the convolutional layer weights 205 | int size = filters * inputChannels * kernelSize * kernelSize; 206 | nvinfer1::Weights convWt{nvinfer1::DataType::kFLOAT, nullptr, size}; 207 | val = new float[size]; 208 | for (int i = 0; i < size; ++i) 209 | { 210 | val[i] = weights[weightPtr]; 211 | weightPtr++; 212 | } 213 | convWt.values = val; 214 | trtWeights.push_back(convWt); 215 | nvinfer1::IConvolutionLayer* conv = network->addConvolution( 216 | *input, filters, nvinfer1::DimsHW{kernelSize, kernelSize}, convWt, convBias); 217 | assert(conv != nullptr); 218 | std::string convLayerName = "conv_" + std::to_string(layerIdx); 219 | conv->setName(convLayerName.c_str()); 220 | conv->setStride(nvinfer1::DimsHW{stride, stride}); 221 | conv->setPadding(nvinfer1::DimsHW{pad, pad}); 222 | 223 | return conv; 224 | } 225 | 226 | nvinfer1::ILayer* netAddConvBNLeaky(int layerIdx, std::map& block, 227 | std::vector& weights, 228 | std::vector& trtWeights, int& weightPtr, 229 | int& inputChannels, nvinfer1::ITensor* input, 230 | nvinfer1::INetworkDefinition* network) 231 | { 232 | assert(block.at("type") == "convolutional"); 233 | assert(block.find("batch_normalize") != block.end()); 234 | assert(block.at("batch_normalize") == "1"); 235 | assert(block.at("activation") == "leaky"); 236 | assert(block.find("filters") != block.end()); 237 | assert(block.find("pad") != block.end()); 238 | assert(block.find("size") != block.end()); 239 | assert(block.find("stride") != block.end()); 240 | 241 | bool batchNormalize, bias; 242 | if (block.find("batch_normalize") != block.end()) 243 | { 244 | batchNormalize = (block.at("batch_normalize") == "1"); 245 | bias = false; 246 | } 247 | else 248 | { 249 | batchNormalize = false; 250 | bias = true; 251 | } 252 | // all conv_bn_leaky layers assume bias is false 253 | assert(batchNormalize == true && bias == false); 254 | UNUSED(batchNormalize); 255 | UNUSED(bias); 256 | 257 | int filters = std::stoi(block.at("filters")); 258 | int padding = std::stoi(block.at("pad")); 259 | int kernelSize = std::stoi(block.at("size")); 260 | int stride = std::stoi(block.at("stride")); 261 | int pad; 262 | if (padding) 263 | pad = (kernelSize - 1) / 2; 264 | else 265 | pad = 0; 266 | 267 | /***** CONVOLUTION LAYER *****/ 268 | /*****************************/ 269 | // batch norm weights are before the conv layer 270 | // load BN biases (bn_biases) 271 | std::vector bnBiases; 272 | for (int i = 0; i < filters; ++i) 273 | { 274 | bnBiases.push_back(weights[weightPtr]); 275 | weightPtr++; 276 | } 277 | // load BN weights 278 | std::vector bnWeights; 279 | for (int i = 0; i < filters; ++i) 280 | { 281 | bnWeights.push_back(weights[weightPtr]); 282 | weightPtr++; 283 | } 284 | // load BN running_mean 285 | std::vector bnRunningMean; 286 | for (int i = 0; i < filters; ++i) 287 | { 288 | bnRunningMean.push_back(weights[weightPtr]); 289 | weightPtr++; 290 | } 291 | // load BN running_var 292 | std::vector bnRunningVar; 293 | for (int i = 0; i < filters; ++i) 294 | { 295 | // 1e-05 for numerical stability 296 | bnRunningVar.push_back(sqrt(weights[weightPtr] + 1.0e-5)); 297 | weightPtr++; 298 | } 299 | // load Conv layer weights (GKCRS) 300 | int size = filters * inputChannels * kernelSize * kernelSize; 301 | nvinfer1::Weights convWt{nvinfer1::DataType::kFLOAT, nullptr, size}; 302 | float* val = new float[size]; 303 | for (int i = 0; i < size; ++i) 304 | { 305 | val[i] = weights[weightPtr]; 306 | weightPtr++; 307 | } 308 | convWt.values = val; 309 | trtWeights.push_back(convWt); 310 | nvinfer1::Weights convBias{nvinfer1::DataType::kFLOAT, nullptr, 0}; 311 | trtWeights.push_back(convBias); 312 | nvinfer1::IConvolutionLayer* conv = network->addConvolution( 313 | *input, filters, nvinfer1::DimsHW{kernelSize, kernelSize}, convWt, convBias); 314 | assert(conv != nullptr); 315 | std::string convLayerName = "conv_" + std::to_string(layerIdx); 316 | conv->setName(convLayerName.c_str()); 317 | conv->setStride(nvinfer1::DimsHW{stride, stride}); 318 | conv->setPadding(nvinfer1::DimsHW{pad, pad}); 319 | 320 | /***** BATCHNORM LAYER *****/ 321 | /***************************/ 322 | size = filters; 323 | // create the weights 324 | nvinfer1::Weights shift{nvinfer1::DataType::kFLOAT, nullptr, size}; 325 | nvinfer1::Weights scale{nvinfer1::DataType::kFLOAT, nullptr, size}; 326 | nvinfer1::Weights power{nvinfer1::DataType::kFLOAT, nullptr, size}; 327 | float* shiftWt = new float[size]; 328 | for (int i = 0; i < size; ++i) 329 | { 330 | shiftWt[i] 331 | = bnBiases.at(i) - ((bnRunningMean.at(i) * bnWeights.at(i)) / bnRunningVar.at(i)); 332 | } 333 | shift.values = shiftWt; 334 | float* scaleWt = new float[size]; 335 | for (int i = 0; i < size; ++i) 336 | { 337 | scaleWt[i] = bnWeights.at(i) / bnRunningVar[i]; 338 | } 339 | scale.values = scaleWt; 340 | float* powerWt = new float[size]; 341 | for (int i = 0; i < size; ++i) 342 | { 343 | powerWt[i] = 1.0; 344 | } 345 | power.values = powerWt; 346 | trtWeights.push_back(shift); 347 | trtWeights.push_back(scale); 348 | trtWeights.push_back(power); 349 | // Add the batch norm layers 350 | nvinfer1::IScaleLayer* bn = network->addScale( 351 | *conv->getOutput(0), nvinfer1::ScaleMode::kCHANNEL, shift, scale, power); 352 | assert(bn != nullptr); 353 | std::string bnLayerName = "batch_norm_" + std::to_string(layerIdx); 354 | bn->setName(bnLayerName.c_str()); 355 | /***** ACTIVATION LAYER *****/ 356 | /****************************/ 357 | nvinfer1::ITensor* bnOutput = bn->getOutput(0); 358 | nvinfer1::IActivationLayer* leaky = network->addActivation( 359 | *bnOutput, nvinfer1::ActivationType::kLEAKY_RELU); 360 | leaky->setAlpha(0.1); 361 | assert(leaky != nullptr); 362 | std::string leakyLayerName = "leaky_" + std::to_string(layerIdx); 363 | leaky->setName(leakyLayerName.c_str()); 364 | 365 | return leaky; 366 | } 367 | 368 | nvinfer1::ILayer* netAddUpsample(int layerIdx, std::map& block, 369 | std::vector& weights, 370 | std::vector& trtWeights, int& inputChannels, 371 | nvinfer1::ITensor* input, nvinfer1::INetworkDefinition* network) 372 | { 373 | assert(block.at("type") == "upsample"); 374 | nvinfer1::Dims inpDims = input->getDimensions(); 375 | assert(inpDims.nbDims == 3); 376 | assert(inpDims.d[1] == inpDims.d[2]); 377 | int h = inpDims.d[1]; 378 | int w = inpDims.d[2]; 379 | int stride = std::stoi(block.at("stride")); 380 | // add pre multiply matrix as a constant 381 | nvinfer1::Dims preDims{3, 382 | {1, stride * h, w}, 383 | {nvinfer1::DimensionType::kCHANNEL, nvinfer1::DimensionType::kSPATIAL, 384 | nvinfer1::DimensionType::kSPATIAL}}; 385 | int size = stride * h * w; 386 | nvinfer1::Weights preMul{nvinfer1::DataType::kFLOAT, nullptr, size}; 387 | float* preWt = new float[size]; 388 | /* (2*h * w) 389 | [ [1, 0, ..., 0], 390 | [1, 0, ..., 0], 391 | [0, 1, ..., 0], 392 | [0, 1, ..., 0], 393 | ..., 394 | ..., 395 | [0, 0, ..., 1], 396 | [0, 0, ..., 1] ] 397 | */ 398 | for (int i = 0, idx = 0; i < h; ++i) 399 | { 400 | for (int s = 0; s < stride; ++s) 401 | { 402 | for (int j = 0; j < w; ++j, ++idx) 403 | { 404 | preWt[idx] = (i == j) ? 1.0 : 0.0; 405 | } 406 | } 407 | } 408 | preMul.values = preWt; 409 | trtWeights.push_back(preMul); 410 | nvinfer1::IConstantLayer* preM = network->addConstant(preDims, preMul); 411 | assert(preM != nullptr); 412 | std::string preLayerName = "preMul_" + std::to_string(layerIdx); 413 | preM->setName(preLayerName.c_str()); 414 | // add post multiply matrix as a constant 415 | nvinfer1::Dims postDims{3, 416 | {1, h, stride * w}, 417 | {nvinfer1::DimensionType::kCHANNEL, nvinfer1::DimensionType::kSPATIAL, 418 | nvinfer1::DimensionType::kSPATIAL}}; 419 | size = stride * h * w; 420 | nvinfer1::Weights postMul{nvinfer1::DataType::kFLOAT, nullptr, size}; 421 | float* postWt = new float[size]; 422 | /* (h * 2*w) 423 | [ [1, 1, 0, 0, ..., 0, 0], 424 | [0, 0, 1, 1, ..., 0, 0], 425 | ..., 426 | ..., 427 | [0, 0, 0, 0, ..., 1, 1] ] 428 | */ 429 | for (int i = 0, idx = 0; i < h; ++i) 430 | { 431 | for (int j = 0; j < stride * w; ++j, ++idx) 432 | { 433 | postWt[idx] = (j / stride == i) ? 1.0 : 0.0; 434 | } 435 | } 436 | postMul.values = postWt; 437 | trtWeights.push_back(postMul); 438 | nvinfer1::IConstantLayer* post_m = network->addConstant(postDims, postMul); 439 | assert(post_m != nullptr); 440 | std::string postLayerName = "postMul_" + std::to_string(layerIdx); 441 | post_m->setName(postLayerName.c_str()); 442 | // add matrix multiply layers for upsampling 443 | nvinfer1::IMatrixMultiplyLayer* mm1 444 | = network->addMatrixMultiply(*preM->getOutput(0), nvinfer1::MatrixOperation::kNONE, *input, 445 | nvinfer1::MatrixOperation::kNONE); 446 | assert(mm1 != nullptr); 447 | std::string mm1LayerName = "mm1_" + std::to_string(layerIdx); 448 | mm1->setName(mm1LayerName.c_str()); 449 | nvinfer1::IMatrixMultiplyLayer* mm2 450 | = network->addMatrixMultiply(*mm1->getOutput(0), nvinfer1::MatrixOperation::kNONE, 451 | *post_m->getOutput(0), nvinfer1::MatrixOperation::kNONE); 452 | assert(mm2 != nullptr); 453 | std::string mm2LayerName = "mm2_" + std::to_string(layerIdx); 454 | mm2->setName(mm2LayerName.c_str()); 455 | return mm2; 456 | } 457 | 458 | void printLayerInfo(std::string layerIndex, std::string layerName, std::string layerInput, 459 | std::string layerOutput, std::string weightPtr) 460 | { 461 | std::cout << std::setw(6) << std::left << layerIndex << std::setw(15) << std::left << layerName; 462 | std::cout << std::setw(20) << std::left << layerInput << std::setw(20) << std::left 463 | << layerOutput; 464 | std::cout << std::setw(6) << std::left << weightPtr << std::endl; 465 | } 466 | -------------------------------------------------------------------------------- /yolo/nvdsinfer_custom_impl_Yolo/trt_utils.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | */ 22 | 23 | 24 | #ifndef __TRT_UTILS_H__ 25 | #define __TRT_UTILS_H__ 26 | 27 | #include 28 | #include 29 | #include 30 | #include 31 | #include 32 | #include 33 | #include 34 | 35 | #include "NvInfer.h" 36 | 37 | #define UNUSED(expr) (void)(expr) 38 | #define DIVUP(n, d) ((n) + (d)-1) / (d) 39 | 40 | std::string trim(std::string s); 41 | float clamp(const float val, const float minVal, const float maxVal); 42 | bool fileExists(const std::string fileName, bool verbose = true); 43 | std::vector loadWeights(const std::string weightsFilePath, const std::string& networkType); 44 | std::string dimsToString(const nvinfer1::Dims d); 45 | void displayDimType(const nvinfer1::Dims d); 46 | int getNumChannels(nvinfer1::ITensor* t); 47 | uint64_t get3DTensorVolume(nvinfer1::Dims inputDims); 48 | 49 | // Helper functions to create yolo engine 50 | nvinfer1::ILayer* netAddMaxpool(int layerIdx, std::map& block, 51 | nvinfer1::ITensor* input, nvinfer1::INetworkDefinition* network); 52 | nvinfer1::ILayer* netAddConvLinear(int layerIdx, std::map& block, 53 | std::vector& weights, 54 | std::vector& trtWeights, int& weightPtr, 55 | int& inputChannels, nvinfer1::ITensor* input, 56 | nvinfer1::INetworkDefinition* network); 57 | nvinfer1::ILayer* netAddConvBNLeaky(int layerIdx, std::map& block, 58 | std::vector& weights, 59 | std::vector& trtWeights, int& weightPtr, 60 | int& inputChannels, nvinfer1::ITensor* input, 61 | nvinfer1::INetworkDefinition* network); 62 | nvinfer1::ILayer* netAddUpsample(int layerIdx, std::map& block, 63 | std::vector& weights, 64 | std::vector& trtWeights, int& inputChannels, 65 | nvinfer1::ITensor* input, nvinfer1::INetworkDefinition* network); 66 | void printLayerInfo(std::string layerIndex, std::string layerName, std::string layerInput, 67 | std::string layerOutput, std::string weightPtr); 68 | 69 | #endif 70 | -------------------------------------------------------------------------------- /yolo/nvdsinfer_custom_impl_Yolo/yolo.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | */ 22 | 23 | #include "yolo.h" 24 | #include "yoloPlugins.h" 25 | 26 | #include 27 | #include 28 | #include 29 | 30 | Yolo::Yolo(const NetworkInfo& networkInfo) 31 | : m_NetworkType(networkInfo.networkType), // yolov3 32 | m_ConfigFilePath(networkInfo.configFilePath), // yolov3.cfg 33 | m_WtsFilePath(networkInfo.wtsFilePath), // yolov3.weights 34 | m_DeviceType(networkInfo.deviceType), // kDLA, kGPU 35 | m_InputBlobName(networkInfo.inputBlobName), // data 36 | m_InputH(0), 37 | m_InputW(0), 38 | m_InputC(0), 39 | m_InputSize(0) 40 | {} 41 | 42 | Yolo::~Yolo() 43 | { 44 | destroyNetworkUtils(); 45 | } 46 | 47 | nvinfer1::ICudaEngine *Yolo::createEngine (nvinfer1::IBuilder* builder) 48 | { 49 | assert (builder); 50 | 51 | std::vector weights = loadWeights(m_WtsFilePath, m_NetworkType); 52 | std::vector trtWeights; 53 | 54 | nvinfer1::INetworkDefinition *network = builder->createNetwork(); 55 | if (parseModel(*network) != NVDSINFER_SUCCESS) { 56 | network->destroy(); 57 | return nullptr; 58 | } 59 | 60 | // Build the engine 61 | std::cout << "Building the TensorRT Engine..." << std::endl; 62 | nvinfer1::ICudaEngine * engine = builder->buildCudaEngine(*network); 63 | if (engine) { 64 | std::cout << "Building complete!" << std::endl; 65 | } else { 66 | std::cerr << "Building engine failed!" << std::endl; 67 | } 68 | 69 | // destroy 70 | network->destroy(); 71 | return engine; 72 | } 73 | 74 | NvDsInferStatus Yolo::parseModel(nvinfer1::INetworkDefinition& network) { 75 | destroyNetworkUtils(); 76 | 77 | m_ConfigBlocks = parseConfigFile(m_ConfigFilePath); 78 | parseConfigBlocks(); 79 | 80 | std::vector weights = loadWeights(m_WtsFilePath, m_NetworkType); 81 | // build yolo network 82 | std::cout << "Building Yolo network..." << std::endl; 83 | NvDsInferStatus status = buildYoloNetwork(weights, network); 84 | 85 | if (status == NVDSINFER_SUCCESS) { 86 | std::cout << "Building yolo network complete!" << std::endl; 87 | } else { 88 | std::cerr << "Building yolo network failed!" << std::endl; 89 | } 90 | 91 | return status; 92 | } 93 | 94 | NvDsInferStatus Yolo::buildYoloNetwork( 95 | std::vector& weights, nvinfer1::INetworkDefinition& network) { 96 | int weightPtr = 0; 97 | int channels = m_InputC; 98 | 99 | nvinfer1::ITensor* data = 100 | network.addInput(m_InputBlobName.c_str(), nvinfer1::DataType::kFLOAT, 101 | nvinfer1::DimsCHW{static_cast(m_InputC), 102 | static_cast(m_InputH), static_cast(m_InputW)}); 103 | assert(data != nullptr && data->getDimensions().nbDims > 0); 104 | 105 | nvinfer1::ITensor* previous = data; 106 | std::vector tensorOutputs; 107 | uint outputTensorCount = 0; 108 | 109 | // build the network using the network API 110 | for (uint i = 0; i < m_ConfigBlocks.size(); ++i) { 111 | // check if num. of channels is correct 112 | assert(getNumChannels(previous) == channels); 113 | std::string layerIndex = "(" + std::to_string(tensorOutputs.size()) + ")"; 114 | 115 | if (m_ConfigBlocks.at(i).at("type") == "net") { 116 | printLayerInfo("", "layer", " inp_size", " out_size", "weightPtr"); 117 | } else if (m_ConfigBlocks.at(i).at("type") == "convolutional") { 118 | std::string inputVol = dimsToString(previous->getDimensions()); 119 | nvinfer1::ILayer* out; 120 | std::string layerType; 121 | // check if batch_norm enabled 122 | if (m_ConfigBlocks.at(i).find("batch_normalize") != 123 | m_ConfigBlocks.at(i).end()) { 124 | out = netAddConvBNLeaky(i, m_ConfigBlocks.at(i), weights, 125 | m_TrtWeights, weightPtr, channels, previous, &network); 126 | layerType = "conv-bn-leaky"; 127 | } 128 | else 129 | { 130 | out = netAddConvLinear(i, m_ConfigBlocks.at(i), weights, 131 | m_TrtWeights, weightPtr, channels, previous, &network); 132 | layerType = "conv-linear"; 133 | } 134 | previous = out->getOutput(0); 135 | assert(previous != nullptr); 136 | channels = getNumChannels(previous); 137 | std::string outputVol = dimsToString(previous->getDimensions()); 138 | tensorOutputs.push_back(out->getOutput(0)); 139 | printLayerInfo(layerIndex, layerType, inputVol, outputVol, std::to_string(weightPtr)); 140 | } else if (m_ConfigBlocks.at(i).at("type") == "shortcut") { 141 | assert(m_ConfigBlocks.at(i).at("activation") == "linear"); 142 | assert(m_ConfigBlocks.at(i).find("from") != 143 | m_ConfigBlocks.at(i).end()); 144 | int from = stoi(m_ConfigBlocks.at(i).at("from")); 145 | 146 | std::string inputVol = dimsToString(previous->getDimensions()); 147 | // check if indexes are correct 148 | assert((i - 2 >= 0) && (i - 2 < tensorOutputs.size())); 149 | assert((i + from - 1 >= 0) && (i + from - 1 < tensorOutputs.size())); 150 | assert(i + from - 1 < i - 2); 151 | nvinfer1::IElementWiseLayer* ew = network.addElementWise( 152 | *tensorOutputs[i - 2], *tensorOutputs[i + from - 1], 153 | nvinfer1::ElementWiseOperation::kSUM); 154 | assert(ew != nullptr); 155 | std::string ewLayerName = "shortcut_" + std::to_string(i); 156 | ew->setName(ewLayerName.c_str()); 157 | previous = ew->getOutput(0); 158 | assert(previous != nullptr); 159 | std::string outputVol = dimsToString(previous->getDimensions()); 160 | tensorOutputs.push_back(ew->getOutput(0)); 161 | printLayerInfo(layerIndex, "skip", inputVol, outputVol, " -"); 162 | } else if (m_ConfigBlocks.at(i).at("type") == "yolo") { 163 | nvinfer1::Dims prevTensorDims = previous->getDimensions(); 164 | assert(prevTensorDims.d[1] == prevTensorDims.d[2]); 165 | TensorInfo& curYoloTensor = m_OutputTensors.at(outputTensorCount); 166 | curYoloTensor.gridSize = prevTensorDims.d[1]; 167 | curYoloTensor.stride = m_InputW / curYoloTensor.gridSize; 168 | m_OutputTensors.at(outputTensorCount).volume = curYoloTensor.gridSize 169 | * curYoloTensor.gridSize 170 | * (curYoloTensor.numBBoxes * (5 + curYoloTensor.numClasses)); 171 | std::string layerName = "yolo_" + std::to_string(i); 172 | curYoloTensor.blobName = layerName; 173 | nvinfer1::IPluginV2* yoloPlugin 174 | = new YoloLayerV3(m_OutputTensors.at(outputTensorCount).numBBoxes, 175 | m_OutputTensors.at(outputTensorCount).numClasses, 176 | m_OutputTensors.at(outputTensorCount).gridSize); 177 | assert(yoloPlugin != nullptr); 178 | nvinfer1::IPluginV2Layer* yolo = 179 | network.addPluginV2(&previous, 1, *yoloPlugin); 180 | assert(yolo != nullptr); 181 | yolo->setName(layerName.c_str()); 182 | std::string inputVol = dimsToString(previous->getDimensions()); 183 | previous = yolo->getOutput(0); 184 | assert(previous != nullptr); 185 | previous->setName(layerName.c_str()); 186 | std::string outputVol = dimsToString(previous->getDimensions()); 187 | network.markOutput(*previous); 188 | channels = getNumChannels(previous); 189 | tensorOutputs.push_back(yolo->getOutput(0)); 190 | printLayerInfo(layerIndex, "yolo", inputVol, outputVol, std::to_string(weightPtr)); 191 | ++outputTensorCount; 192 | } else if (m_ConfigBlocks.at(i).at("type") == "region") { 193 | nvinfer1::Dims prevTensorDims = previous->getDimensions(); 194 | assert(prevTensorDims.d[1] == prevTensorDims.d[2]); 195 | TensorInfo& curRegionTensor = m_OutputTensors.at(outputTensorCount); 196 | curRegionTensor.gridSize = prevTensorDims.d[1]; 197 | curRegionTensor.stride = m_InputW / curRegionTensor.gridSize; 198 | m_OutputTensors.at(outputTensorCount).volume = curRegionTensor.gridSize 199 | * curRegionTensor.gridSize 200 | * (curRegionTensor.numBBoxes * (5 + curRegionTensor.numClasses)); 201 | std::string layerName = "region_" + std::to_string(i); 202 | curRegionTensor.blobName = layerName; 203 | nvinfer1::plugin::RegionParameters RegionParameters{ 204 | static_cast(curRegionTensor.numBBoxes), 4, 205 | static_cast(curRegionTensor.numClasses), nullptr}; 206 | std::string inputVol = dimsToString(previous->getDimensions()); 207 | nvinfer1::IPluginV2* regionPlugin 208 | = createRegionPlugin(RegionParameters); 209 | assert(regionPlugin != nullptr); 210 | nvinfer1::IPluginV2Layer* region = 211 | network.addPluginV2(&previous, 1, *regionPlugin); 212 | assert(region != nullptr); 213 | region->setName(layerName.c_str()); 214 | previous = region->getOutput(0); 215 | assert(previous != nullptr); 216 | previous->setName(layerName.c_str()); 217 | std::string outputVol = dimsToString(previous->getDimensions()); 218 | network.markOutput(*previous); 219 | channels = getNumChannels(previous); 220 | tensorOutputs.push_back(region->getOutput(0)); 221 | printLayerInfo(layerIndex, "region", inputVol, outputVol, std::to_string(weightPtr)); 222 | std::cout << "Anchors are being converted to network input resolution i.e. Anchors x " 223 | << curRegionTensor.stride << " (stride)" << std::endl; 224 | for (auto& anchor : curRegionTensor.anchors) anchor *= curRegionTensor.stride; 225 | ++outputTensorCount; 226 | } else if (m_ConfigBlocks.at(i).at("type") == "reorg") { 227 | std::string inputVol = dimsToString(previous->getDimensions()); 228 | nvinfer1::IPluginV2* reorgPlugin = createReorgPlugin(2); 229 | assert(reorgPlugin != nullptr); 230 | nvinfer1::IPluginV2Layer* reorg = 231 | network.addPluginV2(&previous, 1, *reorgPlugin); 232 | assert(reorg != nullptr); 233 | 234 | std::string layerName = "reorg_" + std::to_string(i); 235 | reorg->setName(layerName.c_str()); 236 | previous = reorg->getOutput(0); 237 | assert(previous != nullptr); 238 | std::string outputVol = dimsToString(previous->getDimensions()); 239 | channels = getNumChannels(previous); 240 | tensorOutputs.push_back(reorg->getOutput(0)); 241 | printLayerInfo(layerIndex, "reorg", inputVol, outputVol, std::to_string(weightPtr)); 242 | } 243 | // route layers (single or concat) 244 | else if (m_ConfigBlocks.at(i).at("type") == "route") { 245 | std::string strLayers = m_ConfigBlocks.at(i).at("layers"); 246 | std::vector idxLayers; 247 | size_t lastPos = 0, pos = 0; 248 | while ((pos = strLayers.find(',', lastPos)) != std::string::npos) { 249 | int vL = std::stoi(trim(strLayers.substr(lastPos, pos - lastPos))); 250 | idxLayers.push_back (vL); 251 | lastPos = pos + 1; 252 | } 253 | if (lastPos < strLayers.length()) { 254 | std::string lastV = trim(strLayers.substr(lastPos)); 255 | if (!lastV.empty()) { 256 | idxLayers.push_back (std::stoi(lastV)); 257 | } 258 | } 259 | assert (!idxLayers.empty()); 260 | std::vector concatInputs; 261 | for (int idxLayer : idxLayers) { 262 | if (idxLayer < 0) { 263 | idxLayer = tensorOutputs.size() + idxLayer; 264 | } 265 | assert (idxLayer >= 0 && idxLayer < (int)tensorOutputs.size()); 266 | concatInputs.push_back (tensorOutputs[idxLayer]); 267 | } 268 | nvinfer1::IConcatenationLayer* concat = 269 | network.addConcatenation(concatInputs.data(), concatInputs.size()); 270 | assert(concat != nullptr); 271 | std::string concatLayerName = "route_" + std::to_string(i - 1); 272 | concat->setName(concatLayerName.c_str()); 273 | // concatenate along the channel dimension 274 | concat->setAxis(0); 275 | previous = concat->getOutput(0); 276 | assert(previous != nullptr); 277 | std::string outputVol = dimsToString(previous->getDimensions()); 278 | // set the output volume depth 279 | channels 280 | = getNumChannels(previous); 281 | tensorOutputs.push_back(concat->getOutput(0)); 282 | printLayerInfo(layerIndex, "route", " -", outputVol, 283 | std::to_string(weightPtr)); 284 | } else if (m_ConfigBlocks.at(i).at("type") == "upsample") { 285 | std::string inputVol = dimsToString(previous->getDimensions()); 286 | nvinfer1::ILayer* out = netAddUpsample(i - 1, m_ConfigBlocks[i], 287 | weights, m_TrtWeights, channels, previous, &network); 288 | previous = out->getOutput(0); 289 | std::string outputVol = dimsToString(previous->getDimensions()); 290 | tensorOutputs.push_back(out->getOutput(0)); 291 | printLayerInfo(layerIndex, "upsample", inputVol, outputVol, " -"); 292 | } else if (m_ConfigBlocks.at(i).at("type") == "maxpool") { 293 | std::string inputVol = dimsToString(previous->getDimensions()); 294 | nvinfer1::ILayer* out = 295 | netAddMaxpool(i, m_ConfigBlocks.at(i), previous, &network); 296 | previous = out->getOutput(0); 297 | assert(previous != nullptr); 298 | std::string outputVol = dimsToString(previous->getDimensions()); 299 | tensorOutputs.push_back(out->getOutput(0)); 300 | printLayerInfo(layerIndex, "maxpool", inputVol, outputVol, std::to_string(weightPtr)); 301 | } 302 | else 303 | { 304 | std::cout << "Unsupported layer type --> \"" 305 | << m_ConfigBlocks.at(i).at("type") << "\"" << std::endl; 306 | assert(0); 307 | } 308 | } 309 | 310 | if ((int)weights.size() != weightPtr) 311 | { 312 | std::cout << "Number of unused weights left : " << weights.size() - weightPtr << std::endl; 313 | assert(0); 314 | } 315 | 316 | std::cout << "Output yolo blob names :" << std::endl; 317 | for (auto& tensor : m_OutputTensors) { 318 | std::cout << tensor.blobName << std::endl; 319 | } 320 | 321 | int nbLayers = network.getNbLayers(); 322 | std::cout << "Total number of yolo layers: " << nbLayers << std::endl; 323 | 324 | return NVDSINFER_SUCCESS; 325 | } 326 | 327 | std::vector> 328 | Yolo::parseConfigFile (const std::string cfgFilePath) 329 | { 330 | assert(fileExists(cfgFilePath)); 331 | std::ifstream file(cfgFilePath); 332 | assert(file.good()); 333 | std::string line; 334 | std::vector> blocks; 335 | std::map block; 336 | 337 | while (getline(file, line)) 338 | { 339 | if (line.size() == 0) continue; 340 | if (line.front() == '#') continue; 341 | line = trim(line); 342 | if (line.front() == '[') 343 | { 344 | if (block.size() > 0) 345 | { 346 | blocks.push_back(block); 347 | block.clear(); 348 | } 349 | std::string key = "type"; 350 | std::string value = trim(line.substr(1, line.size() - 2)); 351 | block.insert(std::pair(key, value)); 352 | } 353 | else 354 | { 355 | int cpos = line.find('='); 356 | std::string key = trim(line.substr(0, cpos)); 357 | std::string value = trim(line.substr(cpos + 1)); 358 | block.insert(std::pair(key, value)); 359 | } 360 | } 361 | blocks.push_back(block); 362 | return blocks; 363 | } 364 | 365 | void Yolo::parseConfigBlocks() 366 | { 367 | for (auto block : m_ConfigBlocks) { 368 | if (block.at("type") == "net") 369 | { 370 | assert((block.find("height") != block.end()) 371 | && "Missing 'height' param in network cfg"); 372 | assert((block.find("width") != block.end()) && "Missing 'width' param in network cfg"); 373 | assert((block.find("channels") != block.end()) 374 | && "Missing 'channels' param in network cfg"); 375 | 376 | m_InputH = std::stoul(block.at("height")); 377 | m_InputW = std::stoul(block.at("width")); 378 | m_InputC = std::stoul(block.at("channels")); 379 | assert(m_InputW == m_InputH); 380 | m_InputSize = m_InputC * m_InputH * m_InputW; 381 | } 382 | else if ((block.at("type") == "region") || (block.at("type") == "yolo")) 383 | { 384 | assert((block.find("num") != block.end()) 385 | && std::string("Missing 'num' param in " + block.at("type") + " layer").c_str()); 386 | assert((block.find("classes") != block.end()) 387 | && std::string("Missing 'classes' param in " + block.at("type") + " layer") 388 | .c_str()); 389 | assert((block.find("anchors") != block.end()) 390 | && std::string("Missing 'anchors' param in " + block.at("type") + " layer") 391 | .c_str()); 392 | 393 | TensorInfo outputTensor; 394 | std::string anchorString = block.at("anchors"); 395 | while (!anchorString.empty()) 396 | { 397 | int npos = anchorString.find_first_of(','); 398 | if (npos != -1) 399 | { 400 | float anchor = std::stof(trim(anchorString.substr(0, npos))); 401 | outputTensor.anchors.push_back(anchor); 402 | anchorString.erase(0, npos + 1); 403 | } 404 | else 405 | { 406 | float anchor = std::stof(trim(anchorString)); 407 | outputTensor.anchors.push_back(anchor); 408 | break; 409 | } 410 | } 411 | 412 | if ((m_NetworkType == "yolov3") || (m_NetworkType == "yolov3-tiny")) 413 | { 414 | assert((block.find("mask") != block.end()) 415 | && std::string("Missing 'mask' param in " + block.at("type") + " layer") 416 | .c_str()); 417 | 418 | std::string maskString = block.at("mask"); 419 | while (!maskString.empty()) 420 | { 421 | int npos = maskString.find_first_of(','); 422 | if (npos != -1) 423 | { 424 | uint mask = std::stoul(trim(maskString.substr(0, npos))); 425 | outputTensor.masks.push_back(mask); 426 | maskString.erase(0, npos + 1); 427 | } 428 | else 429 | { 430 | uint mask = std::stoul(trim(maskString)); 431 | outputTensor.masks.push_back(mask); 432 | break; 433 | } 434 | } 435 | } 436 | 437 | outputTensor.numBBoxes = outputTensor.masks.size() > 0 438 | ? outputTensor.masks.size() 439 | : std::stoul(trim(block.at("num"))); 440 | outputTensor.numClasses = std::stoul(block.at("classes")); 441 | m_OutputTensors.push_back(outputTensor); 442 | } 443 | } 444 | } 445 | 446 | void Yolo::destroyNetworkUtils() { 447 | // deallocate the weights 448 | for (uint i = 0; i < m_TrtWeights.size(); ++i) { 449 | if (m_TrtWeights[i].count > 0) 450 | free(const_cast(m_TrtWeights[i].values)); 451 | } 452 | m_TrtWeights.clear(); 453 | } 454 | 455 | -------------------------------------------------------------------------------- /yolo/nvdsinfer_custom_impl_Yolo/yolo.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | */ 22 | 23 | #ifndef _YOLO_H_ 24 | #define _YOLO_H_ 25 | 26 | #include 27 | #include 28 | #include 29 | #include 30 | 31 | #include "NvInfer.h" 32 | #include "trt_utils.h" 33 | 34 | #include "nvdsinfer_custom_impl.h" 35 | 36 | /** 37 | * Holds all the file paths required to build a network. 38 | */ 39 | struct NetworkInfo 40 | { 41 | std::string networkType; 42 | std::string configFilePath; 43 | std::string wtsFilePath; 44 | std::string deviceType; 45 | std::string inputBlobName; 46 | }; 47 | 48 | /** 49 | * Holds information about an output tensor of the yolo network. 50 | */ 51 | struct TensorInfo 52 | { 53 | std::string blobName; 54 | uint stride{0}; 55 | uint gridSize{0}; 56 | uint numClasses{0}; 57 | uint numBBoxes{0}; 58 | uint64_t volume{0}; 59 | std::vector masks; 60 | std::vector anchors; 61 | int bindingIndex{-1}; 62 | float* hostBuffer{nullptr}; 63 | }; 64 | 65 | class Yolo : public IModelParser { 66 | public: 67 | Yolo(const NetworkInfo& networkInfo); 68 | ~Yolo() override; 69 | bool hasFullDimsSupported() const override { return false; } 70 | const char* getModelName() const override { 71 | return m_ConfigFilePath.empty() ? m_NetworkType.c_str() 72 | : m_ConfigFilePath.c_str(); 73 | } 74 | NvDsInferStatus parseModel(nvinfer1::INetworkDefinition& network) override; 75 | 76 | nvinfer1::ICudaEngine *createEngine (nvinfer1::IBuilder* builder); 77 | 78 | protected: 79 | const std::string m_NetworkType; 80 | const std::string m_ConfigFilePath; 81 | const std::string m_WtsFilePath; 82 | const std::string m_DeviceType; 83 | const std::string m_InputBlobName; 84 | std::vector m_OutputTensors; 85 | std::vector> m_ConfigBlocks; 86 | uint m_InputH; 87 | uint m_InputW; 88 | uint m_InputC; 89 | uint64_t m_InputSize; 90 | 91 | // TRT specific members 92 | std::vector m_TrtWeights; 93 | 94 | private: 95 | NvDsInferStatus buildYoloNetwork( 96 | std::vector& weights, nvinfer1::INetworkDefinition& network); 97 | std::vector> parseConfigFile( 98 | const std::string cfgFilePath); 99 | void parseConfigBlocks(); 100 | void destroyNetworkUtils(); 101 | }; 102 | 103 | #endif // _YOLO_H_ 104 | -------------------------------------------------------------------------------- /yolo/nvdsinfer_custom_impl_Yolo/yoloPlugins.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | */ 22 | 23 | #include "yoloPlugins.h" 24 | #include "NvInferPlugin.h" 25 | #include 26 | #include 27 | #include 28 | 29 | namespace { 30 | template 31 | void write(char*& buffer, const T& val) 32 | { 33 | *reinterpret_cast(buffer) = val; 34 | buffer += sizeof(T); 35 | } 36 | 37 | template 38 | void read(const char*& buffer, T& val) 39 | { 40 | val = *reinterpret_cast(buffer); 41 | buffer += sizeof(T); 42 | } 43 | } //namespace 44 | 45 | // Forward declaration of cuda kernels 46 | cudaError_t cudaYoloLayerV3 ( 47 | const void* input, void* output, const uint& batchSize, 48 | const uint& gridSize, const uint& numOutputClasses, 49 | const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream); 50 | 51 | YoloLayerV3::YoloLayerV3 (const void* data, size_t length) 52 | { 53 | const char *d = static_cast(data); 54 | read(d, m_NumBoxes); 55 | read(d, m_NumClasses); 56 | read(d, m_GridSize); 57 | read(d, m_OutputSize); 58 | }; 59 | 60 | YoloLayerV3::YoloLayerV3 ( 61 | const uint& numBoxes, const uint& numClasses, const uint& gridSize) : 62 | m_NumBoxes(numBoxes), 63 | m_NumClasses(numClasses), 64 | m_GridSize(gridSize) 65 | { 66 | assert(m_NumBoxes > 0); 67 | assert(m_NumClasses > 0); 68 | assert(m_GridSize > 0); 69 | m_OutputSize = m_GridSize * m_GridSize * (m_NumBoxes * (4 + 1 + m_NumClasses)); 70 | }; 71 | 72 | nvinfer1::Dims 73 | YoloLayerV3::getOutputDimensions( 74 | int index, const nvinfer1::Dims* inputs, int nbInputDims) 75 | { 76 | assert(index == 0); 77 | assert(nbInputDims == 1); 78 | return inputs[0]; 79 | } 80 | 81 | bool YoloLayerV3::supportsFormat ( 82 | nvinfer1::DataType type, nvinfer1::PluginFormat format) const { 83 | return (type == nvinfer1::DataType::kFLOAT && 84 | format == nvinfer1::PluginFormat::kNCHW); 85 | } 86 | 87 | void 88 | YoloLayerV3::configureWithFormat ( 89 | const nvinfer1::Dims* inputDims, int nbInputs, 90 | const nvinfer1::Dims* outputDims, int nbOutputs, 91 | nvinfer1::DataType type, nvinfer1::PluginFormat format, int maxBatchSize) 92 | { 93 | assert(nbInputs == 1); 94 | assert (format == nvinfer1::PluginFormat::kNCHW); 95 | assert(inputDims != nullptr); 96 | } 97 | 98 | int YoloLayerV3::enqueue( 99 | int batchSize, const void* const* inputs, void** outputs, void* workspace, 100 | cudaStream_t stream) 101 | { 102 | CHECK(cudaYoloLayerV3( 103 | inputs[0], outputs[0], batchSize, m_GridSize, m_NumClasses, m_NumBoxes, 104 | m_OutputSize, stream)); 105 | return 0; 106 | } 107 | 108 | size_t YoloLayerV3::getSerializationSize() const 109 | { 110 | return sizeof(m_NumBoxes) + sizeof(m_NumClasses) + sizeof(m_GridSize) + sizeof(m_OutputSize); 111 | } 112 | 113 | void YoloLayerV3::serialize(void* buffer) const 114 | { 115 | char *d = static_cast(buffer); 116 | write(d, m_NumBoxes); 117 | write(d, m_NumClasses); 118 | write(d, m_GridSize); 119 | write(d, m_OutputSize); 120 | } 121 | 122 | nvinfer1::IPluginV2* YoloLayerV3::clone() const 123 | { 124 | return new YoloLayerV3 (m_NumBoxes, m_NumClasses, m_GridSize); 125 | } 126 | 127 | REGISTER_TENSORRT_PLUGIN(YoloLayerV3PluginCreator); 128 | -------------------------------------------------------------------------------- /yolo/nvdsinfer_custom_impl_Yolo/yoloPlugins.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | */ 22 | 23 | #ifndef __YOLO_PLUGINS__ 24 | #define __YOLO_PLUGINS__ 25 | 26 | #include 27 | #include 28 | #include 29 | #include 30 | #include 31 | 32 | #include "NvInferPlugin.h" 33 | 34 | #define CHECK(status) \ 35 | { \ 36 | if (status != 0) \ 37 | { \ 38 | std::cout << "Cuda failure: " << cudaGetErrorString(status) << " in file " << __FILE__ \ 39 | << " at line " << __LINE__ << std::endl; \ 40 | abort(); \ 41 | } \ 42 | } 43 | 44 | namespace 45 | { 46 | const char* YOLOV3LAYER_PLUGIN_VERSION {"1"}; 47 | const char* YOLOV3LAYER_PLUGIN_NAME {"YoloLayerV3_TRT"}; 48 | } // namespace 49 | 50 | class YoloLayerV3 : public nvinfer1::IPluginV2 51 | { 52 | public: 53 | YoloLayerV3 (const void* data, size_t length); 54 | YoloLayerV3 (const uint& numBoxes, const uint& numClasses, const uint& gridSize); 55 | const char* getPluginType () const override { return YOLOV3LAYER_PLUGIN_NAME; } 56 | const char* getPluginVersion () const override { return YOLOV3LAYER_PLUGIN_VERSION; } 57 | int getNbOutputs () const override { return 1; } 58 | 59 | nvinfer1::Dims getOutputDimensions ( 60 | int index, const nvinfer1::Dims* inputs, 61 | int nbInputDims) override; 62 | 63 | bool supportsFormat ( 64 | nvinfer1::DataType type, nvinfer1::PluginFormat format) const override; 65 | 66 | void configureWithFormat ( 67 | const nvinfer1::Dims* inputDims, int nbInputs, 68 | const nvinfer1::Dims* outputDims, int nbOutputs, 69 | nvinfer1::DataType type, nvinfer1::PluginFormat format, int maxBatchSize) override; 70 | 71 | int initialize () override { return 0; } 72 | void terminate () override {} 73 | size_t getWorkspaceSize (int maxBatchSize) const override { return 0; } 74 | int enqueue ( 75 | int batchSize, const void* const* inputs, void** outputs, 76 | void* workspace, cudaStream_t stream) override; 77 | size_t getSerializationSize() const override; 78 | void serialize (void* buffer) const override; 79 | void destroy () override { delete this; } 80 | nvinfer1::IPluginV2* clone() const override; 81 | 82 | void setPluginNamespace (const char* pluginNamespace)override { 83 | m_Namespace = pluginNamespace; 84 | } 85 | virtual const char* getPluginNamespace () const override { 86 | return m_Namespace.c_str(); 87 | } 88 | 89 | private: 90 | uint m_NumBoxes {0}; 91 | uint m_NumClasses {0}; 92 | uint m_GridSize {0}; 93 | uint64_t m_OutputSize {0}; 94 | std::string m_Namespace {""}; 95 | }; 96 | 97 | class YoloLayerV3PluginCreator : public nvinfer1::IPluginCreator 98 | { 99 | public: 100 | YoloLayerV3PluginCreator () {} 101 | ~YoloLayerV3PluginCreator () {} 102 | 103 | const char* getPluginName () const override { return YOLOV3LAYER_PLUGIN_NAME; } 104 | const char* getPluginVersion () const override { return YOLOV3LAYER_PLUGIN_VERSION; } 105 | 106 | const nvinfer1::PluginFieldCollection* getFieldNames() override { 107 | std::cerr<< "YoloLayerV3PluginCreator::getFieldNames is not implemented" << std::endl; 108 | return nullptr; 109 | } 110 | 111 | nvinfer1::IPluginV2* createPlugin ( 112 | const char* name, const nvinfer1::PluginFieldCollection* fc) override 113 | { 114 | std::cerr<< "YoloLayerV3PluginCreator::getFieldNames is not implemented.\n"; 115 | return nullptr; 116 | } 117 | 118 | nvinfer1::IPluginV2* deserializePlugin ( 119 | const char* name, const void* serialData, size_t serialLength) override 120 | { 121 | std::cout << "Deserialize yoloLayerV3 plugin: " << name << std::endl; 122 | return new YoloLayerV3(serialData, serialLength); 123 | } 124 | 125 | void setPluginNamespace(const char* libNamespace) override { 126 | m_Namespace = libNamespace; 127 | } 128 | const char* getPluginNamespace() const override { 129 | return m_Namespace.c_str(); 130 | } 131 | 132 | private: 133 | std::string m_Namespace {""}; 134 | }; 135 | 136 | #endif // __YOLO_PLUGINS__ --------------------------------------------------------------------------------