mirror of https://github.com/open-mmlab/mmyolo.git
[Happy New Year] TensorRT + DeepStream Support. (#485)
* [Happy New Years] TensorRT + DeepStream Support. * Add deepstream config * pre-commit fix * Fix name * [Happy New Years] TensorRT + DeepStream Support. * Add deepstream config * pre-commit fix * Fix name * Add config * Add rtmdet deepstreampull/531/head
tripleMu
GitHub
parent
4a7228a5cc
commit
031e7450bc
projects/easydeploy/deepstream
configs
custom_mmyolo_bbox_parser
|
|
@ -0,0 +1,35 @@
|
|||
cmake_minimum_required(VERSION 2.8.12)
|
||||
|
||||
set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86)
|
||||
set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc)
|
||||
|
||||
project(nvdsparsebbox_mmyolo LANGUAGES CXX)
|
||||
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3 -g -Wall -Werror -shared -fPIC")
|
||||
set(CMAKE_CXX_STANDARD 14)
|
||||
set(CMAKE_BUILD_TYPE Release)
|
||||
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
|
||||
|
||||
# CUDA
|
||||
find_package(CUDA REQUIRED)
|
||||
|
||||
# TensorRT
|
||||
set(TensorRT_INCLUDE_DIRS "/usr/include/x86_64-linux-gnu" CACHE STRING "TensorRT headers path")
|
||||
set(TensorRT_LIBRARIES "/usr/lib/x86_64-linux-gnu" CACHE STRING "TensorRT libs path")
|
||||
|
||||
# DeepStream
|
||||
set(DEEPSTREAM "/opt/nvidia/deepstream/deepstream" CACHE STRING "DeepStream root path")
|
||||
set(DS_LIBRARIES ${DEEPSTREAM}/lib)
|
||||
set(DS_INCLUDE_DIRS ${DEEPSTREAM}/sources/includes)
|
||||
|
||||
include_directories(
|
||||
${CUDA_INCLUDE_DIRS}
|
||||
${TensorRT_INCLUDE_DIRS}
|
||||
${DS_INCLUDE_DIRS})
|
||||
|
||||
add_library(
|
||||
${PROJECT_NAME}
|
||||
SHARED
|
||||
custom_mmyolo_bbox_parser/nvdsparsebbox_mmyolo.cpp)
|
||||
|
||||
target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin)
|
||||
|
|
@ -0,0 +1,80 @@
|
|||
person
|
||||
bicycle
|
||||
car
|
||||
motorbike
|
||||
aeroplane
|
||||
bus
|
||||
train
|
||||
truck
|
||||
boat
|
||||
traffic light
|
||||
fire hydrant
|
||||
stop sign
|
||||
parking meter
|
||||
bench
|
||||
bird
|
||||
cat
|
||||
dog
|
||||
horse
|
||||
sheep
|
||||
cow
|
||||
elephant
|
||||
bear
|
||||
zebra
|
||||
giraffe
|
||||
backpack
|
||||
umbrella
|
||||
handbag
|
||||
tie
|
||||
suitcase
|
||||
frisbee
|
||||
skis
|
||||
snowboard
|
||||
sports ball
|
||||
kite
|
||||
baseball bat
|
||||
baseball glove
|
||||
skateboard
|
||||
surfboard
|
||||
tennis racket
|
||||
bottle
|
||||
wine glass
|
||||
cup
|
||||
fork
|
||||
knife
|
||||
spoon
|
||||
bowl
|
||||
banana
|
||||
apple
|
||||
sandwich
|
||||
orange
|
||||
broccoli
|
||||
carrot
|
||||
hot dog
|
||||
pizza
|
||||
donut
|
||||
cake
|
||||
chair
|
||||
sofa
|
||||
pottedplant
|
||||
bed
|
||||
diningtable
|
||||
toilet
|
||||
tvmonitor
|
||||
laptop
|
||||
mouse
|
||||
remote
|
||||
keyboard
|
||||
cell phone
|
||||
microwave
|
||||
oven
|
||||
toaster
|
||||
sink
|
||||
refrigerator
|
||||
book
|
||||
clock
|
||||
vase
|
||||
scissors
|
||||
teddy bear
|
||||
hair drier
|
||||
toothbrush
|
||||
|
|
@ -0,0 +1,22 @@
|
|||
[property]
|
||||
gpu-id=0
|
||||
net-scale-factor=0.01735207357279195
|
||||
offsets=57.375;57.12;58.395
|
||||
model-color-format=1
|
||||
model-engine-file=../end2end.engine
|
||||
labelfile-path=../coco_labels.txt
|
||||
batch-size=1
|
||||
network-mode=0
|
||||
num-detected-classes=80
|
||||
interval=0
|
||||
gie-unique-id=1
|
||||
process-mode=1
|
||||
network-type=0
|
||||
cluster-mode=2
|
||||
maintain-aspect-ratio=1
|
||||
parse-bbox-func-name=NvDsInferParseCustomMMYOLO
|
||||
custom-lib-path=../build/libnvdsparsebbox_mmyolo.so
|
||||
|
||||
[class-attrs-all]
|
||||
pre-cluster-threshold=0.45
|
||||
topk=100
|
||||
|
|
@ -0,0 +1,118 @@
|
|||
#include "nvdsinfer_custom_impl.h"
|
||||
#include <cassert>
|
||||
#include <iostream>
|
||||
|
||||
/**
|
||||
* Function expected by DeepStream for decoding the MMYOLO output.
|
||||
*
|
||||
* C-linkage [extern "C"] was written to prevent name-mangling. This function must return true after
|
||||
* adding all bounding boxes to the objectList vector.
|
||||
*
|
||||
* @param [outputLayersInfo] std::vector of NvDsInferLayerInfo objects with information about the output layer.
|
||||
* @param [networkInfo] NvDsInferNetworkInfo object with information about the MMYOLO network.
|
||||
* @param [detectionParams] NvDsInferParseDetectionParams with information about some config params.
|
||||
* @param [objectList] std::vector of NvDsInferParseObjectInfo objects to which bounding box information must
|
||||
* be stored.
|
||||
*
|
||||
* @return true
|
||||
*/
|
||||
|
||||
// This is just the function prototype. The definition is written at the end of the file.
|
||||
extern "C" bool NvDsInferParseCustomMMYOLO(
|
||||
std::vector<NvDsInferLayerInfo> const& outputLayersInfo,
|
||||
NvDsInferNetworkInfo const& networkInfo,
|
||||
NvDsInferParseDetectionParams const& detectionParams,
|
||||
std::vector<NvDsInferParseObjectInfo>& objectList);
|
||||
|
||||
static __inline__ float clamp(float& val, float min, float max)
|
||||
{
|
||||
return val > min ? (val < max ? val : max) : min;
|
||||
}
|
||||
|
||||
static std::vector<NvDsInferParseObjectInfo> decodeMMYoloTensor(
|
||||
const int* num_dets,
|
||||
const float* bboxes,
|
||||
const float* scores,
|
||||
const int* labels,
|
||||
const float& conf_thres,
|
||||
const unsigned int& img_w,
|
||||
const unsigned int& img_h
|
||||
)
|
||||
{
|
||||
std::vector<NvDsInferParseObjectInfo> bboxInfo;
|
||||
size_t nums = num_dets[0];
|
||||
for (size_t i = 0; i < nums; i++)
|
||||
{
|
||||
float score = scores[i];
|
||||
if (score < conf_thres)continue;
|
||||
float x0 = (bboxes[i * 4]);
|
||||
float y0 = (bboxes[i * 4 + 1]);
|
||||
float x1 = (bboxes[i * 4 + 2]);
|
||||
float y1 = (bboxes[i * 4 + 3]);
|
||||
x0 = clamp(x0, 0.f, img_w);
|
||||
y0 = clamp(y0, 0.f, img_h);
|
||||
x1 = clamp(x1, 0.f, img_w);
|
||||
y1 = clamp(y1, 0.f, img_h);
|
||||
NvDsInferParseObjectInfo obj;
|
||||
obj.left = x0;
|
||||
obj.top = y0;
|
||||
obj.width = x1 - x0;
|
||||
obj.height = y1 - y0;
|
||||
obj.detectionConfidence = score;
|
||||
obj.classId = labels[i];
|
||||
bboxInfo.push_back(obj);
|
||||
}
|
||||
|
||||
return bboxInfo;
|
||||
}
|
||||
|
||||
/* C-linkage to prevent name-mangling */
|
||||
extern "C" bool NvDsInferParseCustomMMYOLO(
|
||||
std::vector<NvDsInferLayerInfo> const& outputLayersInfo,
|
||||
NvDsInferNetworkInfo const& networkInfo,
|
||||
NvDsInferParseDetectionParams const& detectionParams,
|
||||
std::vector<NvDsInferParseObjectInfo>& objectList)
|
||||
{
|
||||
|
||||
// Some assertions and error checking.
|
||||
if (outputLayersInfo.empty() || outputLayersInfo.size() != 4)
|
||||
{
|
||||
std::cerr << "Could not find output layer in bbox parsing" << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
// Score threshold of bboxes.
|
||||
const float conf_thres = detectionParams.perClassThreshold[0];
|
||||
|
||||
// Obtaining the output layer.
|
||||
const NvDsInferLayerInfo& num_dets = outputLayersInfo[0];
|
||||
const NvDsInferLayerInfo& bboxes = outputLayersInfo[1];
|
||||
const NvDsInferLayerInfo& scores = outputLayersInfo[2];
|
||||
const NvDsInferLayerInfo& labels = outputLayersInfo[3];
|
||||
|
||||
// num_dets(int) bboxes(float) scores(float) labels(int)
|
||||
assert (num_dets.dims.numDims == 2);
|
||||
assert (bboxes.dims.numDims == 3);
|
||||
assert (scores.dims.numDims == 2);
|
||||
assert (labels.dims.numDims == 2);
|
||||
|
||||
|
||||
// Decoding the output tensor of MMYOLO to the NvDsInferParseObjectInfo format.
|
||||
std::vector<NvDsInferParseObjectInfo> objects =
|
||||
decodeMMYoloTensor(
|
||||
(const int*)(num_dets.buffer),
|
||||
(const float*)(bboxes.buffer),
|
||||
(const float*)(scores.buffer),
|
||||
(const int*)(labels.buffer),
|
||||
conf_thres,
|
||||
networkInfo.width,
|
||||
networkInfo.height
|
||||
);
|
||||
|
||||
objectList.clear();
|
||||
objectList = objects;
|
||||
return true;
|
||||
}
|
||||
|
||||
/* Check that the custom function has been defined correctly */
|
||||
CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomMMYOLO);
|
||||
|
|
@ -0,0 +1,62 @@
|
|||
[application]
|
||||
enable-perf-measurement=1
|
||||
perf-measurement-interval-sec=5
|
||||
|
||||
[tiled-display]
|
||||
enable=1
|
||||
rows=1
|
||||
columns=1
|
||||
width=1280
|
||||
height=720
|
||||
gpu-id=0
|
||||
nvbuf-memory-type=0
|
||||
|
||||
[source0]
|
||||
enable=1
|
||||
type=3
|
||||
uri=file:///opt/nvidia/deepstream/deepstream/samples/streams/sample_1080p_h264.mp4
|
||||
num-sources=1
|
||||
gpu-id=0
|
||||
cudadec-memtype=0
|
||||
|
||||
[sink0]
|
||||
enable=1
|
||||
type=2
|
||||
sync=0
|
||||
gpu-id=0
|
||||
nvbuf-memory-type=0
|
||||
|
||||
[osd]
|
||||
enable=1
|
||||
gpu-id=0
|
||||
border-width=5
|
||||
text-size=15
|
||||
text-color=1;1;1;1;
|
||||
text-bg-color=0.3;0.3;0.3;1
|
||||
font=Serif
|
||||
show-clock=0
|
||||
clock-x-offset=800
|
||||
clock-y-offset=820
|
||||
clock-text-size=12
|
||||
clock-color=1;0;0;0
|
||||
nvbuf-memory-type=0
|
||||
|
||||
[streammux]
|
||||
gpu-id=0
|
||||
live-source=0
|
||||
batch-size=1
|
||||
batched-push-timeout=40000
|
||||
width=1920
|
||||
height=1080
|
||||
enable-padding=0
|
||||
nvbuf-memory-type=0
|
||||
|
||||
[primary-gie]
|
||||
enable=1
|
||||
gpu-id=0
|
||||
gie-unique-id=1
|
||||
nvbuf-memory-type=0
|
||||
config-file=configs/config_infer_rtmdet.txt
|
||||
|
||||
[tests]
|
||||
file-loop=0
|
||||
Loading…
Reference in New Issue