├── .gitignore ├── README.md ├── __pycache__ └── utils.cpython-36.pyc ├── configs ├── config_infer_primary.txt ├── config_infer_primary_1.txt ├── config_infer_primary_2.txt ├── config_nvdsanalytics.txt ├── config_tracker.txt └── config_tracker_NvDCF_perf.yml ├── pose.py ├── pose_test01.py ├── pose_test02.py ├── pose_test03.py ├── test_decode.py └── utils ├── display.py └── utils.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.avi 2 | *.zip 3 | __pycache__/ 4 | utils/__pycache__/ 5 | weights/ -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # DeepStream Pose Estimation 2 | ## 1. 代码内容 3 | * utils/utils.py: 功能函数，包含后处理相关函数 4 | * utils/display.py: 功能函数, 包含把displayMeta加入到batchMeta中的函数 5 | * pose.py: 只包含姿态估计推理 6 | * pose_test01.py: 只包含姿态估计推理 7 | * pose_test02.py: 包含姿态估计+跟踪+可视化, NMS放在后处理部分 8 | * pose_test03.py: 包含姿态估计+跟踪+可视化, NMS放在模型中 9 | 10 | ## 2. 对比 11 |   如果要对比NMS在模型中与在模型外的DeepStream运行速度，只需分别运行pose_test02.py与pose_test03.py即可，二者的区别仅在于NMS处理部分, 运行前需要设置好视频路径 INPUT_STREAM , 如果要在 jeson nx 本地可视化结果, 需在运行上述python脚本时设置环境变量export DISPLAY=:0 指定显示设备为本地设备，将nx接个显示器即可与代码运行同步显示，或者设置代码中的DEBUG为True, 保持推理结果视频到本地 12 |
13 | 展开 14 | 15 | step 1 16 | ```python 17 | ... 18 | # 在代码中指定推理视频路径 19 | INPUT_STREAM = ["file:///media/nvidia/SD/project/test/merge.mp4",] 20 | ... 21 | ``` 22 | 23 | step 2 24 | ```bash 25 | # 设置本地显示 26 | export DISPLAY=:0 27 | python3 pose_test02.py 28 | python3 pose_test03.py 29 | ``` 30 |
31 | 32 | ## 3. 后处理函数(utils/utils.py) 33 | preprocessNoNMS: NMS融入模型后的后处理函数 34 | preprocess: NMS未融入模型的后处理函数 -------------------------------------------------------------------------------- /__pycache__/utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Gwencong/deepstream-pose/578fed7e94b9ea9fba1eb6f472053d3c6a5d7b40/__pycache__/utils.cpython-36.pyc -------------------------------------------------------------------------------- /configs/config_infer_primary.txt: -------------------------------------------------------------------------------- 1 | [property] 2 | gpu-id=0 3 | net-scale-factor=0.0039215697906911373 4 | 5 | #Integer 0: RGB 1: BGR 2: GRAY 6 | model-color-format=0 7 | #Infer Processing Mode 1=Primary Mode 2=Secondary Mode 8 | process-mode=1 9 | #model-engine-file=/home/nvidia/project/edgeai-yolov5/weights/Yolov5s6_pose_640_ti_lite.engine 10 | model-engine-file=/home/nvidia/project/yolo-pose/trt_models/yolov5s6_pose_960_no_concat.trt 11 | #model-engine-file=/home/nvidia/project/yolo-pose/trt_models/yolov5s6_pose_640_ti_lite_static.trt 12 | 13 | gie-unique-id=1 14 | batch-size=1 15 | ## 0=Detector, 1=Classifier, 2=Segmentation, 100=Other 16 | network-type=100 17 | #Integer 0: FP32 1: INT8 2: FP16 18 | network-mode=2 19 | 20 | #When a network supports both implicit batch dimension and full dimension, force the implicit batch dimension mode. Boolean 21 | #force-implicit-batch-dim=1 22 | 23 | #Indicates whether to maintain aspect ratio while scaling input. 24 | maintain-aspect-ratio=1 25 | 26 | #Integer 0: OpenCV groupRectangles() 1: DBSCAN 2: Non Maximum Suppression 3: DBSCAN + NMS Hybrid 4: No clustering 27 | #cluster-mode=2 28 | 29 | infer-dims=3;640;640 30 | #output-blob-names=output 31 | 32 | #if batch-size!=1 33 | #tensor-meta-pool-size=30 34 | scaling-compute-hw=1 35 | 36 | symmetric-padding=1 37 | 38 | output-tensor-meta=1 39 | -------------------------------------------------------------------------------- /configs/config_infer_primary_1.txt: -------------------------------------------------------------------------------- 1 | [property] 2 | gpu-id=0 3 | net-scale-factor=0.0039215697906911373 4 | 5 | #Integer 0: RGB 1: BGR 2: GRAY 6 | model-color-format=0 7 | #Infer Processing Mode 1=Primary Mode 2=Secondary Mode 8 | process-mode=1 9 | 10 | #model-engine-file=../weights/yolov5s6_pose_640_ti_lite.trt 11 | model-engine-file=/media/nvidia/SD/project/yolo-pose-escalator/weights/yolov5l6_pose-FP16.trt 12 | 13 | gie-unique-id=1 14 | batch-size=1 15 | ## 0=Detector, 1=Classifier, 2=Segmentation, 100=Other 16 | network-type=100 17 | #Integer 0: FP32 1: INT8 2: FP16 18 | network-mode=2 19 | 20 | #When a network supports both implicit batch dimension and full dimension, force the implicit batch dimension mode. Boolean 21 | #force-implicit-batch-dim=1 22 | 23 | #Indicates whether to maintain aspect ratio while scaling input. 24 | maintain-aspect-ratio=1 25 | 26 | #Integer 0: OpenCV groupRectangles() 1: DBSCAN 2: Non Maximum Suppression 3: DBSCAN + NMS Hybrid 4: No clustering 27 | #cluster-mode=2 28 | 29 | infer-dims=3;832;832 30 | #output-blob-names=output 31 | 32 | #if batch-size!=1 33 | #tensor-meta-pool-size=30 34 | scaling-compute-hw=1 35 | 36 | symmetric-padding=1 37 | 38 | output-tensor-meta=1 39 | -------------------------------------------------------------------------------- /configs/config_infer_primary_2.txt: -------------------------------------------------------------------------------- 1 | [property] 2 | gpu-id=0 3 | net-scale-factor=0.0039215697906911373 4 | 5 | #Integer 0: RGB 1: BGR 2: GRAY 6 | model-color-format=0 7 | #Infer Processing Mode 1=Primary Mode 2=Secondary Mode 8 | process-mode=1 9 | 10 | #model-engine-file=../weights/yolov5s6_pose_640_ti_lite_EfficientNMS.trt 11 | model-engine-file=/home/nvidia/project/yolo-pose-escalator/weights/yolov5l6_pose-NMS-INT8.trt 12 | 13 | gie-unique-id=1 14 | batch-size=1 15 | ## 0=Detector, 1=Classifier, 2=Segmentation, 100=Other 16 | network-type=100 17 | #Integer 0: FP32 1: INT8 2: FP16 18 | network-mode=2 19 | 20 | #When a network supports both implicit batch dimension and full dimension, force the implicit batch dimension mode. Boolean 21 | #force-implicit-batch-dim=1 22 | 23 | #Indicates whether to maintain aspect ratio while scaling input. 24 | maintain-aspect-ratio=1 25 | 26 | #Integer 0: OpenCV groupRectangles() 1: DBSCAN 2: Non Maximum Suppression 3: DBSCAN + NMS Hybrid 4: No clustering 27 | #cluster-mode=2 28 | 29 | infer-dims=3;832;832 30 | #output-blob-names=output 31 | 32 | #if batch-size!=1 33 | #tensor-meta-pool-size=30 34 | scaling-compute-hw=1 35 | 36 | symmetric-padding=1 37 | 38 | output-tensor-meta=1 39 | -------------------------------------------------------------------------------- /configs/config_nvdsanalytics.txt: -------------------------------------------------------------------------------- 1 | ################################################################################ 2 | # SPDX-FileCopyrightText: Copyright (c) 2020-2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | # SPDX-License-Identifier: Apache-2.0 4 | # 5 | # Licensed under the Apache License, Version 2.0 (the "License"); 6 | # you may not use this file except in compliance with the License. 7 | # You may obtain a copy of the License at 8 | # 9 | # http://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, software 12 | # distributed under the License is distributed on an "AS IS" BASIS, 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | # See the License for the specific language governing permissions and 15 | # limitations under the License. 16 | ################################################################################ 17 | 18 | # The values in the config file are overridden by values set through GObject 19 | # properties. 20 | 21 | [property] 22 | enable=1 23 | #Width height used for configuration to which below configs are configured 24 | config-width=1280 25 | config-height=720 26 | #osd-mode 0: Dont display any lines, rois and text 27 | # 1: Display only lines, rois and static text i.e. labels 28 | # 2: Display all info from 1 plus information about counts 29 | osd-mode=2 30 | #Set OSD font size that has to be displayed 31 | display-font-size=12 32 | 33 | ## Per stream configuration 34 | [roi-filtering-stream-0] 35 | #enable or disable following feature 36 | enable=0 37 | #ROI to filter select objects, and remove from meta data 38 | roi-RF=612;77;1104;77;1104;720;612;720 39 | #remove objects in the ROI 40 | inverse-roi=0 41 | class-id=0 42 | 43 | ## Per stream configuration 44 | [roi-filtering-stream-1] 45 | #enable or disable following feature 46 | enable=0 47 | #ROI to filter select objects, and remove from meta data 48 | roi-RF=295;643;579;634;642;913;56;828 49 | #remove objects in the ROI 50 | inverse-roi=0 51 | class-id=0 52 | 53 | [overcrowding-stream-0] 54 | enable=0 55 | #roi-OC=295;643;579;634;642;913;56;828 56 | #roi-OC=1000;500;1000;700;300;700;300;500 57 | roi-OC=612;77;1104;77;1104;720;612;720 58 | #no of objects that will trigger OC 59 | object-threshold=3 60 | class-id=0 61 | 62 | [overcrowding-stream-1] 63 | enable=0 64 | #roi-OC=295;643;579;634;642;913;56;828 65 | #roi-OC=1000;500;1000;700;300;700;300;500 66 | roi-OC=900;500;900;700;400;700;400;500 67 | #no of objects that will trigger OC 68 | object-threshold=2 69 | class-id=0 70 | 71 | [line-crossing-stream-0] 72 | enable=0 73 | #Label;direction;lc 74 | #line-crossing-Entry=1072;911;1143;1058;944;1020;1297;1020; 75 | line-crossing-Exit=789;672;1084;900;851;773;1203;732 76 | class-id=0 77 | #extended when 0- only counts crossing on the configured Line 78 | # 1- assumes extended Line crossing counts all the crossing 79 | extended=0 80 | #LC modes supported: 81 | #loose : counts all crossing without strong adherence to direction 82 | #balanced: Strict direction adherence expected compared to mode=loose 83 | #strict : Strict direction adherence expected compared to mode=balanced 84 | mode=loose 85 | 86 | [line-crossing-stream-1] 87 | enable=0 88 | #Label;direction;lc 89 | #line-crossing-Entry=1072;911;1143;1058;944;1020;1297;1020; 90 | line-crossing-Exit=789;672;1084;900;851;773;1203;732 91 | class-id=0 92 | #extended when 0- only counts crossing on the configured Line 93 | # 1- assumes extended Line crossing counts all the crossing 94 | extended=0 95 | #LC modes supported: 96 | #loose : counts all crossing without strong adherence to direction 97 | #balanced: Strict direction adherence expected compared to mode=loose 98 | #strict : Strict direction adherence expected compared to mode=balanced 99 | mode=loose 100 | 101 | [direction-detection-stream-0] 102 | enable=1 103 | #Label;direction; 104 | #direction-South=284;840;360;662; 105 | #direction-North=1106;622;1312;701; 106 | direction-Forward=800;838;826;708; 107 | direction-Backward=826;708;800;838; 108 | class-id=0 109 | 110 | [direction-detection-stream-1] 111 | enable=0 112 | #Label;direction; 113 | direction-South=284;840;360;662; 114 | direction-North=1106;622;1312;701; 115 | class-id=0 116 | 117 | -------------------------------------------------------------------------------- /configs/config_tracker.txt: -------------------------------------------------------------------------------- 1 | ################################################################################ 2 | # SPDX-FileCopyrightText: Copyright (c) 2020-2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | # SPDX-License-Identifier: Apache-2.0 4 | # 5 | # Licensed under the Apache License, Version 2.0 (the "License"); 6 | # you may not use this file except in compliance with the License. 7 | # You may obtain a copy of the License at 8 | # 9 | # http://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, software 12 | # distributed under the License is distributed on an "AS IS" BASIS, 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | # See the License for the specific language governing permissions and 15 | # limitations under the License. 16 | ################################################################################ 17 | 18 | # Mandatory properties for the tracker: 19 | # tracker-width 20 | # tracker-height: needs to be multiple of 6 for NvDCF 21 | # gpu-id 22 | # ll-lib-file: path to low-level tracker lib 23 | # ll-config-file: required for NvDCF, optional for KLT and IOU 24 | # 25 | [tracker] 26 | tracker-width=640 27 | tracker-height=384 28 | gpu-id=0 29 | ll-lib-file=/opt/nvidia/deepstream/deepstream/lib/libnvds_nvmultiobjecttracker.so 30 | 31 | #ll-lib-file=/opt/nvidia/deepstream/deepstream/lib/libnvds_mot_iou.so 32 | #ll-lib-file=/opt/nvidia/deepstream/deepstream/lib/libnvds_mot_klt.so 33 | #ll-lib-file=/opt/nvidia/deepstream/deepstream/lib/libnvds_nvdcf.so 34 | ll-config-file=configs/config_tracker_NvDCF_perf.yml 35 | #ll-config-file=config_tracker_NvDCF_accuracy.yml 36 | #enable-past-frame=1 37 | enable-batch-process=1 38 | tracking-id-reset-mode=3 39 | display-tracking-id=1 40 | -------------------------------------------------------------------------------- /configs/config_tracker_NvDCF_perf.yml: -------------------------------------------------------------------------------- 1 | %YAML:1.0 2 | ################################################################################ 3 | # SPDX-FileCopyrightText: Copyright (c) 2020-2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 4 | # SPDX-License-Identifier: Apache-2.0 5 | # 6 | # Licensed under the Apache License, Version 2.0 (the "License"); 7 | # you may not use this file except in compliance with the License. 8 | # You may obtain a copy of the License at 9 | # 10 | # http://www.apache.org/licenses/LICENSE-2.0 11 | # 12 | # Unless required by applicable law or agreed to in writing, software 13 | # distributed under the License is distributed on an "AS IS" BASIS, 14 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | # See the License for the specific language governing permissions and 16 | # limitations under the License. 17 | ################################################################################ 18 | 19 | BaseConfig: 20 | minDetectorConfidence: 0 # If the confidence of a detector bbox is lower than this, then it won't be considered for tracking 21 | 22 | TargetManagement: 23 | enableBboxUnClipping: 1 # In case the bbox is likely to be clipped by image border, unclip bbox 24 | maxTargetsPerStream: 150 # Max number of targets to track per stream. Recommended to set >10. Note: this value should account for the targets being tracked in shadow mode as well. Max value depends on the GPU memory capacity 25 | 26 | # [Creation & Termination Policy] 27 | minIouDiff4NewTarget: 0.5 # If the IOU between the newly detected object and any of the existing targets is higher than this threshold, this newly detected object will be discarded. 28 | minTrackerConfidence: 0.2 # If the confidence of an object tracker is lower than this on the fly, then it will be tracked in shadow mode. Valid Range: [0.0, 1.0] 29 | probationAge: 3 # If the target's age exceeds this, the target will be considered to be valid. 30 | maxShadowTrackingAge: 30 # Max length of shadow tracking. If the shadowTrackingAge exceeds this limit, the tracker will be terminated. 31 | earlyTerminationAge: 1 # If the shadowTrackingAge reaches this threshold while in TENTATIVE period, the target will be terminated prematurely. 32 | 33 | TrajectoryManagement: 34 | useUniqueID: 0 # Use 64-bit long Unique ID when assignining tracker ID. Default is [true] 35 | 36 | DataAssociator: 37 | dataAssociatorType: 0 # the type of data associator among { DEFAULT= 0 } 38 | associationMatcherType: 0 # the type of matching algorithm among { GREEDY=0, GLOBAL=1 } 39 | checkClassMatch: 1 # If checked, only the same-class objects are associated with each other. Default: true 40 | 41 | # [Association Metric: Thresholds for valid candidates] 42 | minMatchingScore4Overall: 0.0 # Min total score 43 | minMatchingScore4SizeSimilarity: 0.6 # Min bbox size similarity score 44 | minMatchingScore4Iou: 0.0 # Min IOU score 45 | minMatchingScore4VisualSimilarity: 0.7 # Min visual similarity score 46 | 47 | # [Association Metric: Weights] 48 | matchingScoreWeight4VisualSimilarity: 0.6 # Weight for the visual similarity (in terms of correlation response ratio) 49 | matchingScoreWeight4SizeSimilarity: 0.0 # Weight for the Size-similarity score 50 | matchingScoreWeight4Iou: 0.4 # Weight for the IOU score 51 | 52 | StateEstimator: 53 | stateEstimatorType: 1 # the type of state estimator among { DUMMY=0, SIMPLE=1, REGULAR=2 } 54 | 55 | # [Dynamics Modeling] 56 | processNoiseVar4Loc: 2.0 # Process noise variance for bbox center 57 | processNoiseVar4Size: 1.0 # Process noise variance for bbox size 58 | processNoiseVar4Vel: 0.1 # Process noise variance for velocity 59 | measurementNoiseVar4Detector: 4.0 # Measurement noise variance for detector's detection 60 | measurementNoiseVar4Tracker: 16.0 # Measurement noise variance for tracker's localization 61 | 62 | VisualTracker: 63 | visualTrackerType: 1 # the type of visual tracker among { DUMMY=0, NvDCF=1 } 64 | 65 | # [NvDCF: Feature Extraction] 66 | useColorNames: 1 # Use ColorNames feature 67 | useHog: 0 # Use Histogram-of-Oriented-Gradient (HOG) feature 68 | featureImgSizeLevel: 2 # Size of a feature image. Valid range: {1, 2, 3, 4, 5}, from the smallest to the largest 69 | featureFocusOffsetFactor_y: -0.2 # The offset for the center of hanning window relative to the feature height. The center of hanning window would move by (featureFocusOffsetFactor_y*featureMatSize.height) in vertical direction 70 | 71 | # [NvDCF: Correlation Filter] 72 | filterLr: 0.075 # learning rate for DCF filter in exponential moving average. Valid Range: [0.0, 1.0] 73 | filterChannelWeightsLr: 0.1 # learning rate for the channel weights among feature channels. Valid Range: [0.0, 1.0] 74 | gaussianSigma: 0.75 # Standard deviation for Gaussian for desired response when creating DCF filter [pixels] 75 | -------------------------------------------------------------------------------- /pose.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # cython: language_level=3 3 | # coding=utf-8 4 | import sys 5 | import cv2 6 | import math 7 | import time 8 | import pyds 9 | import ctypes 10 | import platform 11 | import numpy as np 12 | import configparser 13 | from datetime import datetime 14 | 15 | import gi 16 | gi.require_version("Gst", "1.0") 17 | gi.require_version("GstRtspServer", "1.0") 18 | from gi.repository import Gst, GstRtspServer, GLib 19 | 20 | 21 | from utils.utils import make_element, is_aarch64, create_source_bin, bus_call 22 | from utils.utils import get_outshape, decode, postprocess 23 | from utils.display import add_obj_meta 24 | 25 | MUX_OUTPUT_WIDTH = 1280 26 | MUX_OUTPUT_HEIGHT = 720 27 | INFER_SHAPE = (1,3,640,640) 28 | 29 | 30 | start_time = time.time() 31 | start_time2 = 0 32 | vid_writer = None 33 | save_video = False 34 | aspect_ratio = float(192 / 256) # w/h 35 | DEBUG = True 36 | codec = "H264" 37 | INPUT_STREAM = [ 38 | # "rtsp://admin:zc62683949@172.16.240.91:554/h264/ch1/main/av_stream", 39 | # "file:////home/nvidia/project/lightweight-human-pose-estimation.pytorch/17人1080p.mp4" 40 | "file:///home/nvidia/project/yolo-pose/merge.mp4" 41 | ] 42 | 43 | vid_writer = None 44 | 45 | 46 | def tracker_sink_pad_buffer_probe(pad, info, u_data): 47 | t = time.time() 48 | global vid_writer 49 | frame_number = 0 50 | num_rects = 0 51 | 52 | gst_buffer = info.get_buffer() 53 | if not gst_buffer: 54 | print("Unable to get GstBuffer ") 55 | return 56 | 57 | # Retrieve batch metadata from the gst_buffer 58 | # Note that pyds.gst_buffer_get_nvds_batch_meta() expects the 59 | # C address of gst_buffer as input, which is obtained with hash(gst_buffer) 60 | batch_meta = pyds.gst_buffer_get_nvds_batch_meta(hash(gst_buffer)) 61 | l_frame = batch_meta.frame_meta_list 62 | 63 | while l_frame is not None: 64 | try: 65 | # Note that l_frame.data needs a cast to pyds.NvDsFrameMeta 66 | # The casting is done by pyds.NvDsFrameMeta.cast() 67 | # The casting also keeps ownership of the underlying memory 68 | # in the C code, so the Python garbage collector will leave 69 | # it alone. 70 | frame_meta = pyds.NvDsFrameMeta.cast(l_frame.data) 71 | except StopIteration: 72 | break 73 | 74 | if DEBUG: 75 | # Getting Image data using nvbufsurface 76 | # the input should be address of buffer and batch_id 77 | n_frame = pyds.get_nvds_buf_surface(hash(gst_buffer), frame_meta.batch_id) 78 | # convert python array into numy array format. 79 | frame_image = np.array(n_frame, copy=True, order="C") 80 | # covert the array into cv2 default color format 81 | frame_image = cv2.cvtColor(frame_image, cv2.COLOR_RGBA2BGR) 82 | 83 | frame_number = frame_meta.frame_num 84 | num_rects = frame_meta.num_obj_meta 85 | 86 | l_usr = frame_meta.frame_user_meta_list 87 | 88 | while l_usr is not None: 89 | try: 90 | # Casting l_obj.data to pyds.NvDsUserMeta 91 | user_meta = pyds.NvDsUserMeta.cast(l_usr.data) 92 | except StopIteration: 93 | break 94 | 95 | # get tensor output 96 | if (user_meta.base_meta.meta_type != pyds.NvDsMetaType.NVDSINFER_TENSOR_OUTPUT_META): # NVDSINFER_TENSOR_OUTPUT_META 97 | try: 98 | l_usr = l_usr.next 99 | except StopIteration: 100 | break 101 | # continue 102 | 103 | try: 104 | tensor_meta = pyds.NvDsInferTensorMeta.cast(user_meta.user_meta_data) 105 | infer_out = [] 106 | # translate pointer to numpy array 107 | for num in range(tensor_meta.num_output_layers): 108 | layer = pyds.get_nvds_LayerInfo(tensor_meta, num) 109 | # print(f'output layer: {layer.layerName}') 110 | # load float* buffer to python 111 | stride = int(layer.layerName.replace('stride_','')) 112 | ptr = ctypes.cast(pyds.get_ptr(layer.buffer), ctypes.POINTER(ctypes.c_float)) 113 | out_shape = get_outshape(INFER_SHAPE,stride) 114 | # print(out_shape) 115 | infer_out.append(np.ctypeslib.as_array(ptr, shape=out_shape)) 116 | decode_out = decode(infer_out,INFER_SHAPE[2:],strides=[8,16,32,64]) 117 | if DEBUG: 118 | pred = postprocess(decode_out,(MUX_OUTPUT_HEIGHT,MUX_OUTPUT_WIDTH),INFER_SHAPE[2:],frame_image) 119 | boxes, confs, kpts = pred 120 | else: 121 | pred = postprocess(decode_out,(MUX_OUTPUT_HEIGHT,MUX_OUTPUT_WIDTH),INFER_SHAPE[2:]) 122 | boxes, confs, kpts = pred 123 | if len(boxes)>0 and len(confs)>0 and len(kpts)>0: 124 | add_obj_meta(frame_meta,batch_meta,boxes[0],confs[0]) 125 | except StopIteration: 126 | break 127 | 128 | try: 129 | l_usr = l_usr.next 130 | except StopIteration: 131 | break 132 | 133 | if DEBUG: 134 | if vid_writer == None: # new video 135 | vid_writer = cv2.VideoWriter( 136 | "record.avi", 137 | cv2.VideoWriter_fourcc("X", "V", "I", "D"), 138 | 25, 139 | (frame_image.shape[1], frame_image.shape[0]), 140 | ) 141 | 142 | vid_writer.write(frame_image.copy()) 143 | 144 | 145 | global start_time, start_time2 146 | cur_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") 147 | CurFPS = 1 / (time.time() - start_time) 148 | AvgFPS = frame_number / (time.time() - start_time2) 149 | # print(boxes ) 150 | num_person = sum([len(box) for box in boxes]) 151 | display_text = f'{cur_time} Person={num_person} Frames={frame_number} FPS={CurFPS:.0f} AvgFPS={AvgFPS:.1f}' 152 | print(display_text) 153 | 154 | start_time = time.time() 155 | if int(start_time2) == 0: 156 | start_time2 = time.time() 157 | 158 | try: 159 | l_frame = l_frame.next 160 | except StopIteration: 161 | break 162 | # Indicating inference is performed on given frame. 163 | pyds.nvds_acquire_meta_lock(batch_meta) 164 | frame_meta.bInferDone=True 165 | pyds.nvds_release_meta_lock(batch_meta) 166 | 167 | cost = time.time()-t 168 | print(f'probe function time cost {cost*1000:.2f}ms') 169 | 170 | 171 | 172 | return Gst.PadProbeReturn.OK 173 | 174 | def osd_sink_pad_buffer_probe(pad, info, u_data): 175 | buffer = info.get_buffer() 176 | batch = pyds.gst_buffer_get_nvds_batch_meta(hash(buffer)) 177 | 178 | l_frame = batch.frame_meta_list 179 | while l_frame: 180 | frame_meta = pyds.NvDsFrameMeta.cast(l_frame.data) 181 | l_obj = frame_meta.obj_meta_list 182 | while l_obj: 183 | obj_meta = pyds.NvDsObjectMeta.cast(l_obj.data) 184 | obj_meta.text_params.display_text = "person{}: {:.2f}".format(obj_meta.object_id ,obj_meta.tracker_confidence) 185 | track_box = obj_meta.tracker_bbox_info.org_bbox_coords 186 | print(track_box.left,track_box.top,track_box.height,track_box.width) 187 | # rect_params = obj_meta.rect_params 188 | # rect_params.left = track_box.left 189 | # rect_params.top = track_box.top 190 | # rect_params.width = track_box.width 191 | # rect_params.height = track_box.height 192 | l_obj = l_obj.next 193 | l_frame = l_frame.next 194 | return Gst.PadProbeReturn.OK 195 | 196 | 197 | 198 | def main(args): 199 | # Standard GStreamer initialization 200 | # Since version 3.11, calling threads_init is no longer needed. See: https://wiki.gnome.org/PyGObject/Threading 201 | # GObject.threads_init() 202 | Gst.init(None) 203 | 204 | # Create gstreamer elements 205 | # Create Pipeline element that will form a connection of other elements 206 | print("Creating Pipeline \n ") 207 | pipeline = Gst.Pipeline() 208 | 209 | if not pipeline: 210 | sys.stderr.write(" Unable to create Pipeline \n") 211 | 212 | # Create nvstreammux instance to form batches from one or more sources. 213 | streammux = make_element("nvstreammux", "Stream-muxer-left") 214 | streammux.set_property("width", MUX_OUTPUT_WIDTH) 215 | streammux.set_property("height", MUX_OUTPUT_HEIGHT) 216 | streammux.set_property("batch-size", 1) 217 | streammux.set_property("batched-push-timeout", 4000) 218 | streammux.set_property("live-source", 1) # rtsp 219 | pipeline.add(streammux) 220 | 221 | number_src_left = len(INPUT_STREAM) 222 | for i in range(number_src_left): 223 | print("Creating source_bin ", i, " \n ") 224 | uri_name = INPUT_STREAM[i] 225 | print(uri_name) 226 | 227 | source_bin = create_source_bin(i, uri_name) 228 | if not source_bin: 229 | sys.stderr.write("Unable to create source bin \n") 230 | pipeline.add(source_bin) 231 | 232 | padname = "sink_%u" % i 233 | sinkpad = streammux.get_request_pad(padname) 234 | if not sinkpad: 235 | sys.stderr.write("Unable to create sink pad bin \n") 236 | srcpad = source_bin.get_static_pad("src") 237 | if not srcpad: 238 | sys.stderr.write("Unable to create src pad bin \n") 239 | srcpad.link(sinkpad) 240 | 241 | 242 | tracker = make_element("nvtracker", "tracker") 243 | 244 | # Use nvdsanalytics to perform analytics on object 245 | nvdsanalytics = make_element("nvdsanalytics", "nvdsanalytics") 246 | nvdsanalytics.set_property("config-file", "configs/config_nvdsanalytics.txt") 247 | 248 | nvvidconv = make_element("nvvideoconvert", "convertor") 249 | 250 | nvvidconv_postosd = make_element("nvvideoconvert", "convertor_postosd") 251 | 252 | nvvidconv1 = make_element("nvvideoconvert", "convertor_pre") 253 | 254 | caps1 = Gst.Caps.from_string("video/x-raw(memory:NVMM), format=RGBA") 255 | filter1 = make_element("capsfilter", "filter1") 256 | filter1.set_property("caps", caps1) 257 | 258 | # Create OSD to draw on the converted RGBA buffer 259 | nvosd = make_element("nvdsosd", "onscreendisplay") 260 | 261 | # Create a caps filter 262 | caps = make_element("capsfilter", "filter") 263 | caps.set_property("caps", Gst.Caps.from_string("video/x-raw(memory:NVMM), format=I420")) 264 | 265 | # Make the encoder 266 | encoder = make_element("nvv4l2h264enc", "encoder") 267 | encoder.set_property("bitrate", 4000000) 268 | encoder.set_property("preset-level", 1) 269 | encoder.set_property("insert-sps-pps", 1) 270 | encoder.set_property("bufapi-version", 1) 271 | 272 | # Make the payload-encode video into RTP packets 273 | rtppay = make_element("rtph264pay", "rtppay") 274 | 275 | # Make the UDP sink 276 | updsink_port_num = 5401 277 | sink = make_element("udpsink", "udpsink") 278 | sink.set_property("host", "224.224.255.255") 279 | sink.set_property("port", updsink_port_num) 280 | sink.set_property("async", False) 281 | sink.set_property("sync", 1) 282 | sink.set_property("qos", 0) 283 | 284 | # Set properties of tracker 285 | config = configparser.ConfigParser() 286 | config.read("configs/config_tracker.txt") 287 | config.sections() 288 | 289 | for key in config["tracker"]: 290 | if key == "tracker-width": 291 | tracker_width = config.getint("tracker", key) 292 | tracker.set_property("tracker-width", tracker_width) 293 | if key == "tracker-height": 294 | tracker_height = config.getint("tracker", key) 295 | tracker.set_property("tracker-height", tracker_height) 296 | if key == "gpu-id": 297 | tracker_gpu_id = config.getint("tracker", key) 298 | tracker.set_property("gpu_id", tracker_gpu_id) 299 | if key == "ll-lib-file": 300 | tracker_ll_lib_file = config.get("tracker", key) 301 | tracker.set_property("ll-lib-file", tracker_ll_lib_file) 302 | if key == "ll-config-file": 303 | tracker_ll_config_file = config.get("tracker", key) 304 | tracker.set_property("ll-config-file", tracker_ll_config_file) 305 | if key == "enable-batch-process": 306 | tracker_enable_batch_process = config.getint("tracker", key) 307 | tracker.set_property("enable_batch_process", tracker_enable_batch_process) 308 | 309 | # Use nvinfer to run inferencing on decoder's output, 310 | # behaviour of inferencing is set through config file 311 | pgie = make_element("nvinfer", "primary-inference-left") 312 | pgie.set_property("config-file-path", "configs/config_infer_primary.txt") 313 | 314 | 315 | if not is_aarch64(): 316 | # Use CUDA unified memory in the pipeline so frames 317 | # can be easily accessed on CPU in Python. 318 | mem_type = int(pyds.NVBUF_MEM_CUDA_UNIFIED) 319 | streammux.set_property("nvbuf-memory-type", mem_type) 320 | nvvidconv.set_property("nvbuf-memory-type", mem_type) 321 | nvvidconv1.set_property("nvbuf-memory-type", mem_type) 322 | #tiler.set_property("nvbuf-memory-type", mem_type) 323 | 324 | print("Adding elements to Pipeline \n") 325 | pipeline.add(pgie) 326 | pipeline.add(nvvidconv1) 327 | pipeline.add(filter1) 328 | pipeline.add(tracker) 329 | pipeline.add(nvdsanalytics) 330 | pipeline.add(nvvidconv) 331 | pipeline.add(nvosd) 332 | pipeline.add(nvvidconv_postosd) 333 | pipeline.add(caps) 334 | pipeline.add(encoder) 335 | pipeline.add(rtppay) 336 | pipeline.add(sink) 337 | # Link the elements together: 338 | print("Linking elements in the Pipeline \n") 339 | streammux.link(pgie) 340 | pgie.link(nvvidconv1) 341 | nvvidconv1.link(filter1) 342 | filter1.link(tracker) 343 | 344 | # nvosd -> nvvidconv -> caps -> encoder -> rtppay -> udpsink 345 | tracker.link(nvdsanalytics) 346 | nvdsanalytics.link(nvvidconv) 347 | nvvidconv.link(nvosd) 348 | 349 | nvosd.link(nvvidconv_postosd) 350 | nvvidconv_postosd.link(caps) 351 | caps.link(encoder) 352 | encoder.link(rtppay) 353 | rtppay.link(sink) 354 | 355 | # create and event loop and feed gstreamer bus mesages to it 356 | loop = GLib.MainLoop() 357 | 358 | bus = pipeline.get_bus() 359 | bus.add_signal_watch() 360 | bus.connect("message", bus_call, loop) 361 | 362 | # Start rtsp streaming 363 | rtsp_port_num = 9554 364 | 365 | server = GstRtspServer.RTSPServer.new() 366 | server.props.service = "%d" % rtsp_port_num 367 | server.attach(None) 368 | 369 | factory = GstRtspServer.RTSPMediaFactory.new() 370 | factory.set_launch( 371 | '( udpsrc name=pay0 port=%d buffer-size=524288 caps="application/x-rtp, media=video, clock-rate=90000, \ 372 | encoding-name=(string)%s, payload=96 " )' % (updsink_port_num, codec) 373 | ) 374 | factory.set_shared(True) 375 | server.get_mount_points().add_factory("/review", factory) 376 | 377 | # Lets add probe to get informed of the meta data generated, we add probe to 378 | # the sink pad of the osd element, since by that time, the buffer would have 379 | # had got all the metadata. 380 | 381 | trackerpad = tracker.get_static_pad("sink") 382 | if not trackerpad: 383 | sys.stderr.write(" Unable to get sink pad of tracker \n") 384 | trackerpad.add_probe(Gst.PadProbeType.BUFFER, tracker_sink_pad_buffer_probe, 0) 385 | ''' 386 | nvanalytics_src_pad = nvdsanalytics.get_static_pad("src") 387 | if not nvanalytics_src_pad: 388 | sys.stderr.write(" Unable to get src pad of analytics\n") 389 | nvanalytics_src_pad.add_probe( 390 | Gst.PadProbeType.BUFFER, nvanalytics_src_pad_buffer_probe, 0 391 | ) 392 | ''' 393 | print("Starting pipeline \n") 394 | 395 | # start play back and listed to events 396 | pipeline.set_state(Gst.State.PLAYING) 397 | try: 398 | loop.run() 399 | except: 400 | if save_video: 401 | vid_writer.release() 402 | 403 | # cleanup 404 | pipeline.set_state(Gst.State.NULL) 405 | print("End pipeline \n") 406 | 407 | 408 | if __name__ == "__main__": 409 | sys.exit(main(sys.argv)) 410 | -------------------------------------------------------------------------------- /pose_test01.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # cython: language_level=3 3 | # coding=utf-8 4 | import os 5 | import sys 6 | import cv2 7 | import time 8 | import pyds 9 | import ctypes 10 | import platform 11 | import numpy as np 12 | import configparser 13 | from datetime import datetime 14 | os.environ['CUDA_LAUNCH_BLOCKING'] = '1' # 设置cuda同步，方便debug 15 | 16 | import gi 17 | gi.require_version("Gst", "1.0") 18 | gi.require_version("GstRtspServer", "1.0") 19 | from gi.repository import Gst, GstRtspServer, GLib 20 | 21 | from utils.utils import make_element, is_aarch64, create_source_bin, bus_call 22 | from utils.utils import get_total_outshape, postprocess, postprocessNoNMS 23 | 24 | 25 | MUX_OUTPUT_WIDTH = 640 26 | MUX_OUTPUT_HEIGHT =640 27 | INFER_SHAPE = (1,3,640,640) 28 | OUT_SHAPE = get_total_outshape(INFER_SHAPE) 29 | INPUT_STREAM = ["file:///media/nvidia/SD/project/test/2in1_2.mp4"] 30 | DEBUG = False 31 | 32 | codec = "H264" 33 | start_time = time.time() 34 | start_time2 = 0 35 | vid_writer = None 36 | data = dict() 37 | data['conf_thres'] = 0.1 38 | data['iou_thres'] = 0.45 39 | 40 | 41 | def tracker_sink_pad_buffer_probe(pad, info, u_data): 42 | t = time.time() 43 | global data,vid_writer 44 | frame_number = 0 45 | num_rects = 0 46 | 47 | gst_buffer = info.get_buffer() 48 | if not gst_buffer: 49 | print("Unable to get GstBuffer ") 50 | return 51 | 52 | batch_meta = pyds.gst_buffer_get_nvds_batch_meta(hash(gst_buffer)) 53 | 54 | l_frame = batch_meta.frame_meta_list 55 | while l_frame is not None: 56 | try: 57 | frame_meta = pyds.NvDsFrameMeta.cast(l_frame.data) 58 | except StopIteration: 59 | break 60 | 61 | if DEBUG: 62 | # Getting Image data using nvbufsurface 63 | # the input should be address of buffer and batch_id 64 | n_frame = pyds.get_nvds_buf_surface(hash(gst_buffer), frame_meta.batch_id) 65 | frame_image = np.array(n_frame, copy=True, order="C") 66 | frame_image = cv2.cvtColor(frame_image, cv2.COLOR_RGBA2BGR) 67 | 68 | frame_number = frame_meta.frame_num 69 | num_rects = frame_meta.num_obj_meta 70 | pad_index = frame_meta.pad_index 71 | l_usr = frame_meta.frame_user_meta_list 72 | 73 | while l_usr is not None: 74 | try: 75 | # Casting l_obj.data to pyds.NvDsUserMeta 76 | user_meta = pyds.NvDsUserMeta.cast(l_usr.data) 77 | except StopIteration: 78 | break 79 | 80 | # get tensor output 81 | if (user_meta.base_meta.meta_type != 82 | pyds.NvDsMetaType.NVDSINFER_TENSOR_OUTPUT_META): # NVDSINFER_TENSOR_OUTPUT_META 83 | try: 84 | l_usr = l_usr.next 85 | except StopIteration: 86 | break 87 | continue 88 | 89 | try: 90 | tensor_meta = pyds.NvDsInferTensorMeta.cast(user_meta.user_meta_data) 91 | layer = pyds.get_nvds_LayerInfo(tensor_meta, 0) 92 | # load float* buffer to python 93 | ptr = ctypes.cast(pyds.get_ptr(layer.buffer), ctypes.POINTER(ctypes.c_float)) 94 | out = np.ctypeslib.as_array(ptr, shape=OUT_SHAPE) 95 | 96 | if DEBUG: 97 | pred = postprocess(out,(MUX_OUTPUT_HEIGHT,MUX_OUTPUT_WIDTH),INFER_SHAPE[2:],frame_image) 98 | boxes, confs, kpts = pred 99 | else: 100 | pred = postprocess(out,(MUX_OUTPUT_HEIGHT,MUX_OUTPUT_WIDTH),INFER_SHAPE[2:]) 101 | boxes, confs, kpts = pred 102 | # tensor_meta = pyds.NvDsInferTensorMeta.cast(user_meta.user_meta_data) 103 | # model_outs = [] 104 | # shapes = [(1,100,3),(1,100,4),(1,100,51),(1,100,1)] 105 | # for idx in range(tensor_meta.num_output_layers): 106 | # layer = pyds.get_nvds_LayerInfo(tensor_meta, idx) 107 | # # load float* buffer to python 108 | # ptr = ctypes.cast(pyds.get_ptr(layer.buffer), ctypes.POINTER(ctypes.c_float)) 109 | # out = np.ctypeslib.as_array(ptr, shape=shapes[idx]) 110 | # model_outs.append(out) 111 | # if DEBUG: 112 | # pred = postprocess_nonms(model_outs,(MUX_OUTPUT_HEIGHT,MUX_OUTPUT_WIDTH),INFER_SHAPE[2:],frame_image) 113 | # boxes, confs, kpts = pred 114 | # else: 115 | # pred = postprocess_nonms(model_outs,(MUX_OUTPUT_HEIGHT,MUX_OUTPUT_WIDTH),INFER_SHAPE[2:]) 116 | # boxes, confs, kpts = pred 117 | except StopIteration: 118 | break 119 | 120 | try: 121 | l_usr = l_usr.next 122 | except StopIteration: 123 | break 124 | 125 | if DEBUG: 126 | if vid_writer == None: # new video 127 | vid_writer = cv2.VideoWriter( 128 | "record.avi", 129 | cv2.VideoWriter_fourcc("X", "V", "I", "D"), 130 | 25, 131 | (frame_image.shape[1], frame_image.shape[0]), 132 | ) 133 | 134 | vid_writer.write(frame_image.copy()) 135 | 136 | 137 | global start_time, start_time2 138 | cur_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") 139 | CurFPS = 1 / (time.time() - start_time) 140 | AvgFPS = frame_number / (time.time() - start_time2) 141 | # print(boxes ) 142 | try: 143 | num_person = sum([len(box) for box in boxes]) 144 | except Exception: 145 | num_person = 0 146 | display_text = f'{cur_time} Person={num_person} Frames={frame_number} FPS={CurFPS:.0f} AvgFPS={AvgFPS:.1f}' 147 | print(display_text) 148 | 149 | start_time = time.time() 150 | if int(start_time2) == 0: 151 | start_time2 = time.time() 152 | try: 153 | l_frame = l_frame.next 154 | except StopIteration: 155 | break 156 | print(f'probe function time cost:{(time.time()-t)*1000:.2f}ms') 157 | return Gst.PadProbeReturn.OK 158 | 159 | 160 | def main(args): 161 | # Standard GStreamer initialization 162 | # Since version 3.11, calling threads_init is no longer needed. See: https://wiki.gnome.org/PyGObject/Threading 163 | # GObject.threads_init() 164 | Gst.init(None) 165 | 166 | # Create gstreamer elements 167 | # Create Pipeline element that will form a connection of other elements 168 | print("Creating Pipeline \n ") 169 | pipeline = Gst.Pipeline() 170 | if not pipeline: 171 | sys.stderr.write(" Unable to create Pipeline \n") 172 | 173 | nvvidconv_postosd = make_element("nvvideoconvert", "convertor_postosd") 174 | 175 | nvvidconv1 = make_element("nvvideoconvert", "convertor_pre") 176 | 177 | caps1 = Gst.Caps.from_string("video/x-raw(memory:NVMM), format=RGBA") 178 | filter1 = make_element("capsfilter", "filter1") 179 | filter1.set_property("caps", caps1) 180 | 181 | # Create a caps filter 182 | caps = make_element("capsfilter", "filter") 183 | caps.set_property("caps", Gst.Caps.from_string("video/x-raw(memory:NVMM), format=I420")) 184 | 185 | fakesink = make_element('fakesink','sink') 186 | fakesink.set_property('silent', 1) 187 | 188 | # Create nvstreammux instance to form batches from one or more sources. 189 | streammux = make_element("nvstreammux", "Stream-muxer-left") 190 | streammux.set_property("enable-padding", True) 191 | streammux.set_property("width", MUX_OUTPUT_WIDTH) 192 | streammux.set_property("height", MUX_OUTPUT_HEIGHT) 193 | streammux.set_property("batch-size", 1) 194 | streammux.set_property("batched-push-timeout", 40000) 195 | streammux.set_property("live-source", 1) # rtsp 196 | pipeline.add(streammux) 197 | 198 | pgie = make_element("nvinfer", "primary-inference-left") 199 | pgie.set_property("config-file-path", "configs/config_infer_primary_1.txt") 200 | 201 | number_src = len(INPUT_STREAM) 202 | for i in range(number_src): 203 | print("Creating source_bin ", i, " \n ") 204 | uri_name = INPUT_STREAM[i] 205 | print(uri_name) 206 | 207 | source_bin = create_source_bin(i, uri_name) 208 | if not source_bin: 209 | sys.stderr.write("Unable to create source bin \n") 210 | pipeline.add(source_bin) 211 | 212 | padname = "sink_%u" % i 213 | sinkpad = streammux.get_request_pad(padname) 214 | if not sinkpad: 215 | sys.stderr.write("Unable to create sink pad bin \n") 216 | srcpad = source_bin.get_static_pad("src") 217 | if not srcpad: 218 | sys.stderr.write("Unable to create src pad bin \n") 219 | srcpad.link(sinkpad) 220 | 221 | 222 | print("Adding elements to Pipeline \n") 223 | pipeline.add(pgie) 224 | pipeline.add(nvvidconv_postosd) 225 | pipeline.add(nvvidconv1) 226 | pipeline.add(filter1) 227 | pipeline.add(caps) 228 | pipeline.add(fakesink) 229 | 230 | # Link the elements together: 231 | print("Linking elements in the Pipeline \n") 232 | streammux.link(nvvidconv1) 233 | nvvidconv1.link(filter1) 234 | filter1.link(pgie) 235 | pgie.link(nvvidconv_postosd) 236 | nvvidconv_postosd.link(caps) 237 | caps.link(fakesink) 238 | 239 | # create and event loop and feed gstreamer bus mesages to it 240 | loop = GLib.MainLoop() 241 | 242 | bus = pipeline.get_bus() 243 | bus.add_signal_watch() 244 | bus.connect("message", bus_call, loop) 245 | 246 | pgiepad = pgie.get_static_pad("src") 247 | if not pgiepad: 248 | sys.stderr.write(" Unable to get sink pad of tracker \n") 249 | pgiepad.add_probe(Gst.PadProbeType.BUFFER, tracker_sink_pad_buffer_probe, 0) 250 | 251 | print("Starting pipeline \n") 252 | 253 | # start play back and listed to events 254 | pipeline.set_state(Gst.State.PLAYING) 255 | try: 256 | loop.run() 257 | except: 258 | if vid_writer: 259 | vid_writer.release() 260 | # cleanup 261 | pipeline.set_state(Gst.State.NULL) 262 | print("End pipeline \n") 263 | 264 | 265 | if __name__ == "__main__": 266 | sys.exit(main(sys.argv)) 267 | -------------------------------------------------------------------------------- /pose_test02.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # cython: language_level=3 3 | # coding=utf-8 4 | import os 5 | import sys 6 | import cv2 7 | import time 8 | import pyds 9 | import ctypes 10 | import numpy as np 11 | from datetime import datetime 12 | os.environ['CUDA_LAUNCH_BLOCKING'] = '1' # 设置cuda同步，方便debug 13 | 14 | import gi 15 | gi.require_version("Gst", "1.0") 16 | gi.require_version("GstRtspServer", "1.0") 17 | from gi.repository import Gst, GstRtspServer, GLib 18 | 19 | from utils.utils import make_element, is_aarch64, create_source_bin, bus_call, set_tracker_config 20 | from utils.utils import get_total_outshape, postprocess 21 | from utils.display import dispaly_frame_pose,add_obj_meta 22 | 23 | 24 | MUX_OUTPUT_WIDTH = 832 25 | MUX_OUTPUT_HEIGHT= 832 26 | INFER_SHAPE = (1,3,832,832) 27 | OUT_SHAPE = get_total_outshape(INFER_SHAPE) 28 | INPUT_STREAM = [ 29 | "file:///media/nvidia/SD/project/test/2in1_2.mp4", 30 | # "file:///media/nvidia/SD/project/test/merge.mp4", 31 | ] 32 | DEBUG = False 33 | 34 | codec = "H264" 35 | start_time = time.time() 36 | start_time2 = 0 37 | vid_writer = None 38 | 39 | 40 | 41 | def pose_src_pad_buffer_probe(pad, info, u_data): 42 | t = time.time() 43 | global data,vid_writer 44 | frame_number = 0 45 | num_rects = 0 46 | 47 | gst_buffer = info.get_buffer() 48 | if not gst_buffer: 49 | print("Unable to get GstBuffer ") 50 | return 51 | 52 | batch_meta = pyds.gst_buffer_get_nvds_batch_meta(hash(gst_buffer)) 53 | 54 | l_frame = batch_meta.frame_meta_list 55 | while l_frame is not None: 56 | try: 57 | frame_meta = pyds.NvDsFrameMeta.cast(l_frame.data) 58 | except StopIteration: 59 | break 60 | 61 | if DEBUG: 62 | # Getting Image data using nvbufsurface 63 | # the input should be address of buffer and batch_id 64 | n_frame = pyds.get_nvds_buf_surface(hash(gst_buffer), frame_meta.batch_id) 65 | frame_image = np.array(n_frame, copy=True, order="C") 66 | frame_image = cv2.cvtColor(frame_image, cv2.COLOR_RGBA2BGR) 67 | 68 | frame_number = frame_meta.frame_num 69 | num_rects = frame_meta.num_obj_meta 70 | pad_index = frame_meta.pad_index 71 | l_usr = frame_meta.frame_user_meta_list 72 | 73 | while l_usr is not None: 74 | try: 75 | # Casting l_obj.data to pyds.NvDsUserMeta 76 | user_meta = pyds.NvDsUserMeta.cast(l_usr.data) 77 | except StopIteration: 78 | break 79 | 80 | # get tensor output 81 | if (user_meta.base_meta.meta_type != 82 | pyds.NvDsMetaType.NVDSINFER_TENSOR_OUTPUT_META): # NVDSINFER_TENSOR_OUTPUT_META 83 | try: 84 | l_usr = l_usr.next 85 | except StopIteration: 86 | break 87 | continue 88 | 89 | try: 90 | tensor_meta = pyds.NvDsInferTensorMeta.cast(user_meta.user_meta_data) 91 | layer = pyds.get_nvds_LayerInfo(tensor_meta, 0) 92 | # load float* buffer to python 93 | ptr = ctypes.cast(pyds.get_ptr(layer.buffer), ctypes.POINTER(ctypes.c_float)) 94 | out = np.ctypeslib.as_array(ptr, shape=OUT_SHAPE) 95 | 96 | if DEBUG: 97 | pred = postprocess(out,(MUX_OUTPUT_HEIGHT,MUX_OUTPUT_WIDTH),INFER_SHAPE[2:],frame_image) 98 | boxes, confs, kpts = pred 99 | else: 100 | pred = postprocess(out,(MUX_OUTPUT_HEIGHT,MUX_OUTPUT_WIDTH),INFER_SHAPE[2:]) 101 | boxes, confs, kpts = pred 102 | if len(boxes)>0 and len(confs)>0 and len(kpts)>0: 103 | add_obj_meta(frame_meta,batch_meta,boxes[0],confs[0]) 104 | dispaly_frame_pose(frame_meta,batch_meta,boxes[0],confs[0],kpts[0]) 105 | except StopIteration: 106 | break 107 | 108 | try: 109 | l_usr = l_usr.next 110 | except StopIteration: 111 | break 112 | 113 | if DEBUG: 114 | if vid_writer == None: # new video 115 | vid_writer = cv2.VideoWriter( 116 | "record.avi", 117 | cv2.VideoWriter_fourcc("X", "V", "I", "D"), 118 | 25, 119 | (frame_image.shape[1], frame_image.shape[0]), 120 | ) 121 | 122 | vid_writer.write(frame_image.copy()) 123 | 124 | 125 | global start_time, start_time2 126 | cur_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") 127 | CurFPS = 1 / (time.time() - start_time) 128 | AvgFPS = frame_number / (time.time() - start_time2) 129 | # print(boxes ) 130 | num_person = sum([len(box) for box in boxes]) 131 | display_text = f'{cur_time} Person={num_person} Frames={frame_number} FPS={CurFPS:.0f} AvgFPS={AvgFPS:.1f}' 132 | print(display_text) 133 | 134 | start_time = time.time() 135 | if int(start_time2) == 0: 136 | start_time2 = time.time() 137 | try: 138 | l_frame = l_frame.next 139 | except StopIteration: 140 | break 141 | pyds.nvds_acquire_meta_lock(batch_meta) 142 | frame_meta.bInferDone=True 143 | pyds.nvds_release_meta_lock(batch_meta) 144 | 145 | print(f'probe function time cost:{(time.time()-t)*1000:.2f}ms') 146 | return Gst.PadProbeReturn.OK 147 | 148 | 149 | def osd_sink_pad_buffer_probe(pad, info, u_data): 150 | buffer = info.get_buffer() 151 | batch = pyds.gst_buffer_get_nvds_batch_meta(hash(buffer)) 152 | 153 | l_frame = batch.frame_meta_list 154 | while l_frame: 155 | frame_meta = pyds.NvDsFrameMeta.cast(l_frame.data) 156 | l_obj = frame_meta.obj_meta_list 157 | while l_obj: 158 | obj_meta = pyds.NvDsObjectMeta.cast(l_obj.data) 159 | obj_meta.text_params.display_text = "person{}: {:.2f}".format(obj_meta.object_id ,obj_meta.tracker_confidence) 160 | track_box = obj_meta.tracker_bbox_info.org_bbox_coords 161 | # print(track_box.left,track_box.top,track_box.height,track_box.width) 162 | # rect_params = obj_meta.rect_params 163 | # rect_params.left = track_box.left 164 | # rect_params.top = track_box.top 165 | # rect_params.width = track_box.width 166 | # rect_params.height = track_box.height 167 | l_obj = l_obj.next 168 | l_frame = l_frame.next 169 | return Gst.PadProbeReturn.OK 170 | 171 | 172 | def main(args): 173 | # Standard GStreamer initialization 174 | # Since version 3.11, calling threads_init is no longer needed. See: https://wiki.gnome.org/PyGObject/Threading 175 | # GObject.threads_init() 176 | Gst.init(None) 177 | 178 | # Create gstreamer elements 179 | # Create Pipeline element that will form a connection of other elements 180 | print("Creating Pipeline \n ") 181 | pipeline = Gst.Pipeline() 182 | if not pipeline: 183 | sys.stderr.write(" Unable to create Pipeline \n") 184 | 185 | nvvidconv_postosd = make_element("nvvideoconvert", "convertor_postosd") 186 | 187 | nvvidconv1 = make_element("nvvideoconvert", "convertor_pre") 188 | 189 | caps1 = Gst.Caps.from_string("video/x-raw(memory:NVMM), format=RGBA") 190 | filter1 = make_element("capsfilter", "filter1") 191 | filter1.set_property("caps", caps1) 192 | 193 | nvosd = make_element("nvdsosd", "onscreendisplay") 194 | nvosd.set_property('display-bbox',1) 195 | nvosd.set_property('display-text',1) 196 | 197 | # Create a caps filter 198 | caps = make_element("capsfilter", "filter") 199 | caps.set_property("caps", Gst.Caps.from_string("video/x-raw(memory:NVMM), format=I420")) 200 | 201 | transform = Gst.ElementFactory.make("nvegltransform", "nvegl-transform") 202 | 203 | sink = make_element("nveglglessink", "nvvideo-renderer") 204 | 205 | sink.set_property('sync', False) 206 | 207 | # Create nvstreammux instance to form batches from one or more sources. 208 | streammux = make_element("nvstreammux", "Stream-muxer-left") 209 | #streammux.set_property("enable-padding", True) 210 | streammux.set_property("width", MUX_OUTPUT_WIDTH) 211 | streammux.set_property("height", MUX_OUTPUT_HEIGHT) 212 | streammux.set_property("batch-size", 1) 213 | streammux.set_property("batched-push-timeout", 40000) 214 | streammux.set_property("live-source", 1) # rtsp 215 | pipeline.add(streammux) 216 | 217 | pgie = make_element("nvinfer", "primary-inference-left") 218 | pgie.set_property("config-file-path", "configs/config_infer_primary_1.txt") 219 | 220 | tracker = make_element("nvtracker", "tracker") 221 | set_tracker_config("configs/config_tracker.txt",tracker) 222 | 223 | 224 | number_src = len(INPUT_STREAM) 225 | for i in range(number_src): 226 | print("Creating source_bin ", i, " \n ") 227 | uri_name = INPUT_STREAM[i] 228 | print(uri_name) 229 | 230 | source_bin = create_source_bin(i, uri_name) 231 | if not source_bin: 232 | sys.stderr.write("Unable to create source bin \n") 233 | pipeline.add(source_bin) 234 | 235 | padname = "sink_%u" % i 236 | sinkpad = streammux.get_request_pad(padname) 237 | if not sinkpad: 238 | sys.stderr.write("Unable to create sink pad bin \n") 239 | srcpad = source_bin.get_static_pad("src") 240 | if not srcpad: 241 | sys.stderr.write("Unable to create src pad bin \n") 242 | srcpad.link(sinkpad) 243 | 244 | 245 | print("Adding elements to Pipeline \n") 246 | pipeline.add(pgie) 247 | pipeline.add(tracker) 248 | pipeline.add(nvvidconv_postosd) 249 | pipeline.add(nvvidconv1) 250 | pipeline.add(nvosd) 251 | pipeline.add(filter1) 252 | pipeline.add(caps) 253 | pipeline.add(transform) 254 | pipeline.add(sink) 255 | 256 | # Link the elements together: 257 | print("Linking elements in the Pipeline \n") 258 | streammux.link(nvvidconv1) 259 | nvvidconv1.link(filter1) 260 | filter1.link(pgie) 261 | pgie.link(tracker) 262 | tracker.link(nvosd) 263 | nvosd.link(nvvidconv_postosd) 264 | nvvidconv_postosd.link(caps) 265 | caps.link(transform) 266 | transform.link(sink) 267 | 268 | # create and event loop and feed gstreamer bus mesages to it 269 | loop = GLib.MainLoop() 270 | 271 | bus = pipeline.get_bus() 272 | bus.add_signal_watch() 273 | bus.connect("message", bus_call, loop) 274 | 275 | pgiepad = pgie.get_static_pad("src") 276 | if not pgiepad: 277 | sys.stderr.write(" Unable to get src pad of tracker \n") 278 | pgiepad.add_probe(Gst.PadProbeType.BUFFER, pose_src_pad_buffer_probe, 0) 279 | 280 | osdpad = nvosd.get_static_pad("sink") 281 | if not osdpad: 282 | sys.stderr.write(" Unable to get sink pad of tracker \n") 283 | osdpad.add_probe(Gst.PadProbeType.BUFFER, osd_sink_pad_buffer_probe, 0) 284 | 285 | print("Starting pipeline \n") 286 | 287 | # start play back and listed to events 288 | pipeline.set_state(Gst.State.PLAYING) 289 | try: 290 | loop.run() 291 | except: 292 | if vid_writer: 293 | vid_writer.release() 294 | # cleanup 295 | pipeline.set_state(Gst.State.NULL) 296 | print("End pipeline \n") 297 | 298 | 299 | if __name__ == "__main__": 300 | sys.exit(main(sys.argv)) 301 | -------------------------------------------------------------------------------- /pose_test03.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # cython: language_level=3 3 | # coding=utf-8 4 | import os 5 | import sys 6 | import cv2 7 | import time 8 | import pyds 9 | import ctypes 10 | import numpy as np 11 | from datetime import datetime 12 | os.environ['CUDA_LAUNCH_BLOCKING'] = '1' # 设置cuda同步，方便debug 13 | 14 | import gi 15 | gi.require_version("Gst", "1.0") 16 | gi.require_version("GstRtspServer", "1.0") 17 | from gi.repository import Gst, GstRtspServer, GLib 18 | 19 | from utils.utils import make_element, is_aarch64, create_source_bin, bus_call, postprocessNoNMS, set_tracker_config 20 | from utils.utils import get_total_outshape, postprocess 21 | from utils.display import dispaly_frame_pose,add_obj_meta 22 | 23 | 24 | MUX_OUTPUT_WIDTH = 832 25 | MUX_OUTPUT_HEIGHT =832 26 | INFER_SHAPE = (1,3,832,832) 27 | OUT_SHAPE = get_total_outshape(INFER_SHAPE) 28 | INPUT_STREAM = [ 29 | # "file:///media/nvidia/SD/project/test/2in1_2.mp4", 30 | "file:///media/nvidia/SD/project/test/no-people.mp4", 31 | ] 32 | DEBUG = False 33 | 34 | codec = "H264" 35 | start_time = time.time() 36 | start_time2 = 0 37 | vid_writer = None 38 | 39 | 40 | 41 | def pose_src_pad_buffer_probe(pad, info, u_data): 42 | t = time.time() 43 | global data,vid_writer 44 | frame_number = 0 45 | num_rects = 0 46 | 47 | gst_buffer = info.get_buffer() 48 | if not gst_buffer: 49 | print("Unable to get GstBuffer ") 50 | return 51 | 52 | batch_meta = pyds.gst_buffer_get_nvds_batch_meta(hash(gst_buffer)) 53 | 54 | l_frame = batch_meta.frame_meta_list 55 | while l_frame is not None: 56 | try: 57 | frame_meta = pyds.NvDsFrameMeta.cast(l_frame.data) 58 | except StopIteration: 59 | break 60 | 61 | if DEBUG: 62 | # Getting Image data using nvbufsurface 63 | # the input should be address of buffer and batch_id 64 | n_frame = pyds.get_nvds_buf_surface(hash(gst_buffer), frame_meta.batch_id) 65 | frame_image = np.array(n_frame, copy=True, order="C") 66 | frame_image = cv2.cvtColor(frame_image, cv2.COLOR_RGBA2BGR) 67 | 68 | frame_number = frame_meta.frame_num 69 | num_rects = frame_meta.num_obj_meta 70 | pad_index = frame_meta.pad_index 71 | l_usr = frame_meta.frame_user_meta_list 72 | 73 | while l_usr is not None: 74 | try: 75 | # Casting l_obj.data to pyds.NvDsUserMeta 76 | user_meta = pyds.NvDsUserMeta.cast(l_usr.data) 77 | except StopIteration: 78 | break 79 | 80 | # get tensor output 81 | if (user_meta.base_meta.meta_type != 82 | pyds.NvDsMetaType.NVDSINFER_TENSOR_OUTPUT_META): # NVDSINFER_TENSOR_OUTPUT_META 83 | try: 84 | l_usr = l_usr.next 85 | except StopIteration: 86 | break 87 | continue 88 | 89 | try: 90 | tensor_meta = pyds.NvDsInferTensorMeta.cast(user_meta.user_meta_data) 91 | model_outs = [] 92 | shapes = [(1,100,3),(1,100,4),(1,100,51),(1,100,1)] 93 | for idx in range(tensor_meta.num_output_layers): 94 | layer = pyds.get_nvds_LayerInfo(tensor_meta, idx) 95 | # load float* buffer to python 96 | ptr = ctypes.cast(pyds.get_ptr(layer.buffer), ctypes.POINTER(ctypes.c_float)) 97 | out = np.ctypeslib.as_array(ptr, shape=shapes[idx]) 98 | model_outs.append(out) 99 | if DEBUG: 100 | pred = postprocessNoNMS(model_outs,(MUX_OUTPUT_HEIGHT,MUX_OUTPUT_WIDTH),INFER_SHAPE[2:],frame_image) 101 | boxes, confs, kpts = pred 102 | else: 103 | pred = postprocessNoNMS(model_outs,(MUX_OUTPUT_HEIGHT,MUX_OUTPUT_WIDTH),INFER_SHAPE[2:]) 104 | boxes, confs, kpts = pred 105 | if len(boxes)>0 and len(confs)>0 and len(kpts)>0: 106 | add_obj_meta(frame_meta,batch_meta,boxes[0],confs[0]) 107 | dispaly_frame_pose(frame_meta,batch_meta,boxes[0],confs[0],kpts[0]) 108 | except StopIteration: 109 | break 110 | 111 | try: 112 | l_usr = l_usr.next 113 | except StopIteration: 114 | break 115 | 116 | if DEBUG: 117 | if vid_writer == None: # new video 118 | vid_writer = cv2.VideoWriter( 119 | "record01.avi", 120 | cv2.VideoWriter_fourcc("X", "V", "I", "D"), 121 | 25, 122 | (frame_image.shape[1], frame_image.shape[0]), 123 | ) 124 | 125 | vid_writer.write(frame_image.copy()) 126 | 127 | 128 | global start_time, start_time2 129 | cur_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") 130 | CurFPS = 1 / (time.time() - start_time) 131 | AvgFPS = frame_number / (time.time() - start_time2) 132 | # print(boxes ) 133 | num_person = sum([len(box) for box in boxes]) 134 | display_text = f'{cur_time} Person={num_person} Frames={frame_number} FPS={CurFPS:.0f} AvgFPS={AvgFPS:.1f}' 135 | print(display_text) 136 | 137 | start_time = time.time() 138 | if int(start_time2) == 0: 139 | start_time2 = time.time() 140 | try: 141 | l_frame = l_frame.next 142 | except StopIteration: 143 | break 144 | pyds.nvds_acquire_meta_lock(batch_meta) 145 | frame_meta.bInferDone=True 146 | pyds.nvds_release_meta_lock(batch_meta) 147 | 148 | print(f'probe function time cost:{(time.time()-t)*1000:.2f}ms') 149 | return Gst.PadProbeReturn.OK 150 | 151 | 152 | def osd_sink_pad_buffer_probe(pad, info, u_data): 153 | buffer = info.get_buffer() 154 | batch = pyds.gst_buffer_get_nvds_batch_meta(hash(buffer)) 155 | 156 | l_frame = batch.frame_meta_list 157 | while l_frame: 158 | frame_meta = pyds.NvDsFrameMeta.cast(l_frame.data) 159 | l_obj = frame_meta.obj_meta_list 160 | while l_obj: 161 | obj_meta = pyds.NvDsObjectMeta.cast(l_obj.data) 162 | obj_meta.text_params.display_text = "person{}: {:.2f}".format(obj_meta.object_id ,obj_meta.tracker_confidence) 163 | track_box = obj_meta.tracker_bbox_info.org_bbox_coords 164 | # print(track_box.left,track_box.top,track_box.height,track_box.width) 165 | # rect_params = obj_meta.rect_params 166 | # rect_params.left = track_box.left 167 | # rect_params.top = track_box.top 168 | # rect_params.width = track_box.width 169 | # rect_params.height = track_box.height 170 | l_obj = l_obj.next 171 | l_frame = l_frame.next 172 | return Gst.PadProbeReturn.OK 173 | 174 | 175 | def main(args): 176 | # Standard GStreamer initialization 177 | # Since version 3.11, calling threads_init is no longer needed. See: https://wiki.gnome.org/PyGObject/Threading 178 | # GObject.threads_init() 179 | Gst.init(None) 180 | 181 | # Create gstreamer elements 182 | # Create Pipeline element that will form a connection of other elements 183 | print("Creating Pipeline \n ") 184 | pipeline = Gst.Pipeline() 185 | if not pipeline: 186 | sys.stderr.write(" Unable to create Pipeline \n") 187 | 188 | nvvidconv_postosd = make_element("nvvideoconvert", "convertor_postosd") 189 | 190 | nvvidconv1 = make_element("nvvideoconvert", "convertor_pre") 191 | 192 | caps1 = Gst.Caps.from_string("video/x-raw(memory:NVMM), format=RGBA") 193 | filter1 = make_element("capsfilter", "filter1") 194 | filter1.set_property("caps", caps1) 195 | 196 | nvosd = make_element("nvdsosd", "onscreendisplay") 197 | nvosd.set_property('display-bbox',1) 198 | nvosd.set_property('display-text',1) 199 | 200 | # Create a caps filter 201 | caps = make_element("capsfilter", "filter") 202 | caps.set_property("caps", Gst.Caps.from_string("video/x-raw(memory:NVMM), format=I420")) 203 | 204 | transform = Gst.ElementFactory.make("nvegltransform", "nvegl-transform") 205 | 206 | sink = make_element("nveglglessink", "nvvideo-renderer") 207 | 208 | sink.set_property('sync', False) 209 | 210 | # Create nvstreammux instance to form batches from one or more sources. 211 | streammux = make_element("nvstreammux", "Stream-muxer-left") 212 | #streammux.set_property("enable-padding", True) 213 | streammux.set_property("width", MUX_OUTPUT_WIDTH) 214 | streammux.set_property("height", MUX_OUTPUT_HEIGHT) 215 | streammux.set_property("batch-size", 1) 216 | streammux.set_property("batched-push-timeout", 40000) 217 | streammux.set_property("live-source", 1) # rtsp 218 | pipeline.add(streammux) 219 | 220 | pgie = make_element("nvinfer", "primary-inference-left") 221 | pgie.set_property("config-file-path", "configs/config_infer_primary_2.txt") 222 | 223 | tracker = make_element("nvtracker", "tracker") 224 | set_tracker_config("configs/config_tracker.txt",tracker) 225 | 226 | 227 | number_src = len(INPUT_STREAM) 228 | for i in range(number_src): 229 | print("Creating source_bin ", i, " \n ") 230 | uri_name = INPUT_STREAM[i] 231 | print(uri_name) 232 | 233 | source_bin = create_source_bin(i, uri_name) 234 | if not source_bin: 235 | sys.stderr.write("Unable to create source bin \n") 236 | pipeline.add(source_bin) 237 | 238 | padname = "sink_%u" % i 239 | sinkpad = streammux.get_request_pad(padname) 240 | if not sinkpad: 241 | sys.stderr.write("Unable to create sink pad bin \n") 242 | srcpad = source_bin.get_static_pad("src") 243 | if not srcpad: 244 | sys.stderr.write("Unable to create src pad bin \n") 245 | srcpad.link(sinkpad) 246 | 247 | 248 | print("Adding elements to Pipeline \n") 249 | pipeline.add(pgie) 250 | pipeline.add(tracker) 251 | pipeline.add(nvvidconv_postosd) 252 | pipeline.add(nvvidconv1) 253 | pipeline.add(nvosd) 254 | pipeline.add(filter1) 255 | pipeline.add(caps) 256 | pipeline.add(transform) 257 | pipeline.add(sink) 258 | 259 | # Link the elements together: 260 | print("Linking elements in the Pipeline \n") 261 | streammux.link(nvvidconv1) 262 | nvvidconv1.link(filter1) 263 | filter1.link(pgie) 264 | pgie.link(tracker) 265 | tracker.link(nvosd) 266 | nvosd.link(nvvidconv_postosd) 267 | nvvidconv_postosd.link(caps) 268 | caps.link(transform) 269 | transform.link(sink) 270 | 271 | # create and event loop and feed gstreamer bus mesages to it 272 | loop = GLib.MainLoop() 273 | 274 | bus = pipeline.get_bus() 275 | bus.add_signal_watch() 276 | bus.connect("message", bus_call, loop) 277 | 278 | pgiepad = pgie.get_static_pad("src") 279 | if not pgiepad: 280 | sys.stderr.write(" Unable to get src pad of tracker \n") 281 | pgiepad.add_probe(Gst.PadProbeType.BUFFER, pose_src_pad_buffer_probe, 0) 282 | 283 | osdpad = nvosd.get_static_pad("sink") 284 | if not osdpad: 285 | sys.stderr.write(" Unable to get sink pad of tracker \n") 286 | osdpad.add_probe(Gst.PadProbeType.BUFFER, osd_sink_pad_buffer_probe, 0) 287 | 288 | print("Starting pipeline \n") 289 | 290 | # start play back and listed to events 291 | pipeline.set_state(Gst.State.PLAYING) 292 | try: 293 | loop.run() 294 | except: 295 | if vid_writer: 296 | vid_writer.release() 297 | # cleanup 298 | pipeline.set_state(Gst.State.NULL) 299 | print("End pipeline \n") 300 | 301 | 302 | if __name__ == "__main__": 303 | sys.exit(main(sys.argv)) 304 | -------------------------------------------------------------------------------- /test_decode.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # cython: language_level=3 3 | # coding=utf-8 4 | import os 5 | import sys 6 | import cv2 7 | import time 8 | import pyds 9 | import ctypes 10 | import platform 11 | import numpy as np 12 | import configparser 13 | from datetime import datetime 14 | os.environ['CUDA_LAUNCH_BLOCKING'] = '1' # 设置cuda同步，方便debug 15 | 16 | import gi 17 | gi.require_version("Gst", "1.0") 18 | gi.require_version("GstRtspServer", "1.0") 19 | from gi.repository import Gst, GstRtspServer, GLib 20 | 21 | from utils.utils import make_element, is_aarch64, create_source_bin, bus_call 22 | from utils.utils import get_total_outshape, postprocess 23 | 24 | 25 | MUX_OUTPUT_WIDTH = 640 26 | MUX_OUTPUT_HEIGHT =640 27 | INFER_SHAPE = (1,3,640,640) 28 | OUT_SHAPE = get_total_outshape(INFER_SHAPE) 29 | INPUT_STREAM = ["file:///media/nvidia/SD/project/test/2in1_2.mp4"] 30 | DEBUG = False 31 | 32 | codec = "H264" 33 | start_time = time.time() 34 | start_time2 = 0 35 | vid_writer = None 36 | data = dict() 37 | data['conf_thres'] = 0.1 38 | data['iou_thres'] = 0.45 39 | 40 | 41 | def tracker_sink_pad_buffer_probe(pad, info, u_data): 42 | t = time.time() 43 | global data,vid_writer 44 | frame_number = 0 45 | num_rects = 0 46 | 47 | gst_buffer = info.get_buffer() 48 | if not gst_buffer: 49 | print("Unable to get GstBuffer ") 50 | return 51 | 52 | batch_meta = pyds.gst_buffer_get_nvds_batch_meta(hash(gst_buffer)) 53 | 54 | l_frame = batch_meta.frame_meta_list 55 | while l_frame is not None: 56 | try: 57 | frame_meta = pyds.NvDsFrameMeta.cast(l_frame.data) 58 | except StopIteration: 59 | break 60 | 61 | if DEBUG: 62 | # Getting Image data using nvbufsurface 63 | # the input should be address of buffer and batch_id 64 | n_frame = pyds.get_nvds_buf_surface(hash(gst_buffer), frame_meta.batch_id) 65 | frame_image = np.array(n_frame, copy=True, order="C") 66 | frame_image = cv2.cvtColor(frame_image, cv2.COLOR_RGBA2BGR) 67 | 68 | frame_number = frame_meta.frame_num 69 | 70 | if DEBUG: 71 | if vid_writer == None: # new video 72 | vid_writer = cv2.VideoWriter( 73 | "record.avi", 74 | cv2.VideoWriter_fourcc("X", "V", "I", "D"), 75 | 25, 76 | (frame_image.shape[1], frame_image.shape[0]), 77 | ) 78 | 79 | vid_writer.write(frame_image.copy()) 80 | 81 | 82 | global start_time, start_time2 83 | cur_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") 84 | CurFPS = 1 / (time.time() - start_time) 85 | AvgFPS = frame_number / (time.time() - start_time2) 86 | 87 | display_text = f'{cur_time} Frames={frame_number} FPS={CurFPS:.0f} AvgFPS={AvgFPS:.1f}' 88 | print(display_text) 89 | 90 | start_time = time.time() 91 | if int(start_time2) == 0: 92 | start_time2 = time.time() 93 | try: 94 | l_frame = l_frame.next 95 | except StopIteration: 96 | break 97 | print(f'probe function time cost:{(time.time()-t)*1000:.2f}ms') 98 | return Gst.PadProbeReturn.OK 99 | 100 | 101 | def main(args): 102 | # Standard GStreamer initialization 103 | # Since version 3.11, calling threads_init is no longer needed. See: https://wiki.gnome.org/PyGObject/Threading 104 | # GObject.threads_init() 105 | Gst.init(None) 106 | 107 | # Create gstreamer elements 108 | # Create Pipeline element that will form a connection of other elements 109 | print("Creating Pipeline \n ") 110 | pipeline = Gst.Pipeline() 111 | if not pipeline: 112 | sys.stderr.write(" Unable to create Pipeline \n") 113 | 114 | nvvidconv_postosd = make_element("nvvideoconvert", "convertor_postosd") 115 | 116 | nvvidconv1 = make_element("nvvideoconvert", "convertor_pre") 117 | 118 | caps1 = Gst.Caps.from_string("video/x-raw(memory:NVMM), format=RGBA") 119 | filter1 = make_element("capsfilter", "filter1") 120 | filter1.set_property("caps", caps1) 121 | 122 | # Create a caps filter 123 | caps = make_element("capsfilter", "filter") 124 | caps.set_property("caps", Gst.Caps.from_string("video/x-raw(memory:NVMM), format=I420")) 125 | 126 | fakesink = make_element('fakesink','sink') 127 | fakesink.set_property('silent', 1) 128 | 129 | # Create nvstreammux instance to form batches from one or more sources. 130 | streammux = make_element("nvstreammux", "Stream-muxer-left") 131 | streammux.set_property("width", MUX_OUTPUT_WIDTH) 132 | streammux.set_property("height", MUX_OUTPUT_HEIGHT) 133 | streammux.set_property("batch-size", 1) 134 | streammux.set_property("batched-push-timeout", 40000) 135 | streammux.set_property("live-source", 1) # rtsp 136 | pipeline.add(streammux) 137 | 138 | number_src = len(INPUT_STREAM) 139 | for i in range(number_src): 140 | print("Creating source_bin ", i, " \n ") 141 | uri_name = INPUT_STREAM[i] 142 | print(uri_name) 143 | 144 | source_bin = create_source_bin(i, uri_name) 145 | if not source_bin: 146 | sys.stderr.write("Unable to create source bin \n") 147 | pipeline.add(source_bin) 148 | 149 | padname = "sink_%u" % i 150 | sinkpad = streammux.get_request_pad(padname) 151 | if not sinkpad: 152 | sys.stderr.write("Unable to create sink pad bin \n") 153 | srcpad = source_bin.get_static_pad("src") 154 | if not srcpad: 155 | sys.stderr.write("Unable to create src pad bin \n") 156 | srcpad.link(sinkpad) 157 | 158 | 159 | print("Adding elements to Pipeline \n") 160 | pipeline.add(nvvidconv_postosd) 161 | pipeline.add(nvvidconv1) 162 | pipeline.add(filter1) 163 | pipeline.add(caps) 164 | pipeline.add(fakesink) 165 | 166 | # Link the elements together: 167 | print("Linking elements in the Pipeline \n") 168 | streammux.link(nvvidconv1) 169 | nvvidconv1.link(filter1) 170 | filter1.link(nvvidconv_postosd) 171 | nvvidconv_postosd.link(caps) 172 | caps.link(fakesink) 173 | 174 | # create and event loop and feed gstreamer bus mesages to it 175 | loop = GLib.MainLoop() 176 | 177 | bus = pipeline.get_bus() 178 | bus.add_signal_watch() 179 | bus.connect("message", bus_call, loop) 180 | 181 | pgiepad = nvvidconv_postosd.get_static_pad("src") 182 | if not pgiepad: 183 | sys.stderr.write(" Unable to get sink pad of tracker \n") 184 | pgiepad.add_probe(Gst.PadProbeType.BUFFER, tracker_sink_pad_buffer_probe, 0) 185 | 186 | print("Starting pipeline \n") 187 | 188 | # start play back and listed to events 189 | pipeline.set_state(Gst.State.PLAYING) 190 | try: 191 | loop.run() 192 | except: 193 | if vid_writer: 194 | vid_writer.release() 195 | # cleanup 196 | pipeline.set_state(Gst.State.NULL) 197 | print("End pipeline \n") 198 | 199 | 200 | if __name__ == "__main__": 201 | sys.exit(main(sys.argv)) 202 | -------------------------------------------------------------------------------- /utils/display.py: -------------------------------------------------------------------------------- 1 | 2 | import cv2 3 | import sys 4 | import pyds 5 | import random 6 | import platform 7 | import matplotlib 8 | import numpy as np 9 | from math import ceil 10 | 11 | import gi 12 | from gi.repository import Gst 13 | gi.require_version("Gst", "1.0") 14 | gi.require_version("GstRtspServer", "1.0") 15 | 16 | 17 | def dispaly_frame_pose(frame_meta, batch_meta, boxes, confs, kpts, step=3): 18 | # add_box_display_meta(frame_meta, batch_meta, boxes, confs,draw_boxes=True, draw_confs=False) 19 | add_pose_display_meta(frame_meta, batch_meta, kpts, draw_kpts=False, draw_limbs=True) 20 | 21 | 22 | 23 | def add_box_display_meta(frame_meta, batch_meta, boxes, confs, draw_boxes=True, draw_confs=True): 24 | # Please try to acquire another display meta if the number is bigger than 16. 25 | need_num = ceil(len(boxes)/16) 26 | display_metas = [pyds.nvds_acquire_display_meta_from_pool(batch_meta) for i in range(need_num)] 27 | display_meta = None 28 | for i, (box,conf) in enumerate(zip(boxes,confs)): 29 | if i % 16 == 0: 30 | display_meta = display_metas[i//16] 31 | display_meta.num_rects = 1 32 | display_meta.num_labels = 1 33 | # boxes 34 | if draw_boxes: 35 | rect_params = display_meta.rect_params[i] 36 | rect_params.left = int(box[0]) 37 | rect_params.top =int(box[1]) 38 | rect_params.width = int(box[2]-box[0]) 39 | rect_params.height = int(box[3]-box[1]) 40 | rect_params.border_width = 2 41 | rect_params.border_color.set(0,0,1.0,1.0) 42 | rect_params.has_bg_color = 1 43 | rect_params.bg_color.set(0.5,0.5,0.5,0.1) 44 | display_meta.num_rects += 1 45 | 46 | # confs 47 | if draw_confs: 48 | conf_text_params = display_meta.text_params[i] 49 | conf_text_params.display_text = "person:{:.2f}".format(conf.squeeze()) 50 | conf_text_params.x_offset = int(box[0]) 51 | conf_text_params.y_offset = max(int(box[1])-20,0) 52 | conf_text_params.font_params.font_name = "Serif" 53 | conf_text_params.font_params.font_color.set(1.0, 1.0, 1.0, 1.0) 54 | conf_text_params.font_params.font_size = 10 55 | conf_text_params.set_bg_clr = 1 56 | conf_text_params.text_bg_clr.set(0.0, 0.0, 0.0, 1.0) 57 | display_meta.num_labels += 1 58 | for display_meta in display_metas: 59 | pyds.nvds_add_display_meta_to_frame(frame_meta, display_meta) 60 | 61 | def add_pose_display_meta(frame_meta, batch_meta, kpts,draw_kpts=True,draw_limbs=True): 62 | palette = np.array([[255, 128, 0], [255, 153, 51], [255, 178, 102], 63 | [230, 230, 0], [255, 153, 255], [153, 204, 255], 64 | [255, 102, 255], [255, 51, 255], [102, 178, 255], 65 | [51, 153, 255], [255, 153, 153], [255, 102, 102], 66 | [255, 51, 51], [153, 255, 153], [102, 255, 102], 67 | [51, 255, 51], [0, 255, 0], [0, 0, 255], [255, 0, 0], 68 | [255, 255, 255]])/255 69 | 70 | skeleton = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12], 71 | [7, 13], [6, 7], [6, 8], [7, 9], [8, 10], [9, 11], [2, 3], 72 | [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]] 73 | 74 | pose_limb_color = palette[[9, 9, 9, 9, 7, 7, 7, 0, 0, 0, 0, 0, 16, 16, 16, 16, 16, 16, 16]] 75 | pose_kpt_color = palette[[16, 16, 16, 16, 16, 0, 0, 0, 0, 0, 0, 9, 9, 9, 9, 9, 9]] 76 | radius = 2 77 | 78 | 79 | for kpt in kpts: 80 | # keypoints 81 | if draw_kpts: 82 | display_meta = pyds.nvds_acquire_display_meta_from_pool(batch_meta) 83 | display_meta.num_circles = 1 84 | for j,pt in enumerate(kpt): 85 | if pt[2]<0.5: 86 | continue 87 | r,g,b = pose_kpt_color[j] 88 | cnt = display_meta.num_circles-1 89 | circle_params = display_meta.circle_params[cnt] 90 | circle_params.xc = int(pt[0]) 91 | circle_params.yc = int(pt[1]) 92 | circle_params.radius = radius 93 | circle_params.circle_color.set(r,g,b,1.0) 94 | circle_params.has_bg_color = 1 95 | circle_params.bg_color.set(r,g,b,1.0) 96 | display_meta.num_circles += 1 97 | if display_meta.num_circles>=16: 98 | display_meta = pyds.nvds_acquire_display_meta_from_pool(batch_meta) 99 | display_meta.num_circles = 1 100 | pyds.nvds_add_display_meta_to_frame(frame_meta, display_meta) 101 | 102 | # limbs 103 | if draw_limbs: 104 | display_meta = pyds.nvds_acquire_display_meta_from_pool(batch_meta) 105 | display_meta.num_lines = 1 106 | for i,(s1,s2) in enumerate(skeleton): 107 | pt1,pt2 = kpt[s1-1],kpt[s2-1] 108 | if pt1[2]<0.5 or pt2[2]<0.5: 109 | continue 110 | r,g,b = pose_limb_color[i] 111 | cnt = display_meta.num_lines-1 112 | line_params = display_meta.line_params[cnt] 113 | line_params.x1 = int(pt1[0]) 114 | line_params.y1 = int(pt1[1]) 115 | line_params.x2 = int(pt2[0]) 116 | line_params.y2 = int(pt2[1]) 117 | line_params.line_width = 2 118 | line_params.line_color.set(r, g, b, 1.0) 119 | display_meta.num_lines += 1 120 | 121 | if display_meta.num_lines>=16: 122 | pyds.nvds_add_display_meta_to_frame(frame_meta, display_meta) 123 | display_meta = pyds.nvds_acquire_display_meta_from_pool(batch_meta) 124 | display_meta.num_lines = 1 125 | pyds.nvds_add_display_meta_to_frame(frame_meta, display_meta) 126 | 127 | def add_obj_meta(frame_meta,batch_meta,boxes,confs): 128 | pyds.nvds_acquire_meta_lock(batch_meta) 129 | for i,(box,conf) in enumerate(zip(boxes,confs)): 130 | new_object = pyds.nvds_acquire_obj_meta_from_pool(batch_meta) 131 | new_object.unique_component_id = 1 132 | new_object.class_id = 0 133 | new_object.confidence = conf 134 | new_object.obj_label = 'person' 135 | new_object.parent = None 136 | 137 | rect_params = new_object.rect_params 138 | rect_params.left = int(box[0]) 139 | rect_params.top = int(box[1]) 140 | rect_params.width = int(box[2]-box[0]) 141 | rect_params.height = int(box[3]-box[1]) 142 | rect_params.border_width = 2 143 | rect_params.border_color.set(0,0,1.0,1.0) 144 | rect_params.has_bg_color = 1 145 | rect_params.bg_color.set(0.5,0.5,0.5,0.1) 146 | 147 | text_params = new_object.text_params 148 | # text_params.display_text = "person{}: {:.2f}".format(new_object.object_id ,conf.squeeze()) 149 | text_params.x_offset = int(box[0]) 150 | text_params.y_offset = max(int(box[1])-20,0) 151 | text_params.font_params.font_name = "Serif" 152 | text_params.font_params.font_color.set(1.0, 1.0, 1.0, 1.0) 153 | text_params.font_params.font_size = 10 154 | text_params.set_bg_clr = 1 155 | text_params.text_bg_clr.set(0.0, 0.0, 0.0, 1.0) 156 | 157 | raw_box = new_object.detector_bbox_info.org_bbox_coords 158 | raw_box.left = int(box[0]) 159 | raw_box.top = int(box[1]) 160 | raw_box.width = int(box[2]-box[0]) 161 | raw_box.height = int(box[3]-box[1]) 162 | pyds.nvds_add_obj_meta_to_frame(frame_meta, new_object, None) 163 | pyds.nvds_release_meta_lock(batch_meta) -------------------------------------------------------------------------------- /utils/utils.py: -------------------------------------------------------------------------------- 1 | 2 | import cv2 3 | import sys 4 | import random 5 | import platform 6 | import matplotlib 7 | import numpy as np 8 | import configparser 9 | 10 | import gi 11 | from gi.repository import Gst 12 | gi.require_version("Gst", "1.0") 13 | gi.require_version("GstRtspServer", "1.0") 14 | 15 | 16 | class Colors: 17 | # Ultralytics color palette https://ultralytics.com/ 18 | def __init__(self): 19 | self.palette = [self.hex2rgb(c) for c in matplotlib.colors.TABLEAU_COLORS.values()] 20 | self.n = len(self.palette) 21 | 22 | def __call__(self, i, bgr=False): 23 | c = self.palette[int(i) % self.n] 24 | return (c[2], c[1], c[0]) if bgr else c 25 | 26 | @staticmethod 27 | def hex2rgb(h): # rgb order (PIL) 28 | return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4)) 29 | 30 | colors = Colors() 31 | 32 | categories = ['person'] 33 | 34 | ## gst 35 | 36 | def bus_call(bus, message, loop): 37 | t = message.type 38 | if t == Gst.MessageType.EOS: 39 | sys.stdout.write("End-of-stream\n") 40 | loop.quit() 41 | elif t == Gst.MessageType.WARNING: 42 | err, debug = message.parse_warning() 43 | sys.stderr.write("Warning: %s: %s\n" % (err, debug)) 44 | elif t == Gst.MessageType.ERROR: 45 | err, debug = message.parse_error() 46 | sys.stderr.write("Error: %s: %s\n" % (err, debug)) 47 | loop.quit() 48 | return True 49 | 50 | def is_aarch64(): 51 | return platform.uname()[4] == "aarch64" 52 | 53 | def make_element(factoryname, name, detail=""): 54 | """ Creates an element with Gst Element Factory make. 55 | Return the element if successfully created, otherwise print 56 | to stderr and return None. 57 | """ 58 | print("Creating {}({}) \n".format(name, factoryname)) 59 | elm = Gst.ElementFactory.make(factoryname, name) 60 | if not elm: 61 | sys.stderr.write("Unable to create {}({}) \n".format(name, factoryname)) 62 | if detail: 63 | sys.stderr.write(detail) 64 | return elm 65 | 66 | def cb_newpad(decodebin, decoder_src_pad, data): 67 | print("In cb_newpad\n") 68 | caps = decoder_src_pad.get_current_caps() 69 | gststruct = caps.get_structure(0) 70 | gstname = gststruct.get_name() 71 | source_bin = data 72 | features = caps.get_features(0) 73 | 74 | # Need to check if the pad created by the decodebin is for video and not 75 | # audio. 76 | print("gstname=", gstname) 77 | if gstname.find("video") != -1: 78 | # Link the decodebin pad only if decodebin has picked nvidia 79 | # decoder plugin nvdec_*. We do this by checking if the pad caps contain 80 | # NVMM memory features. 81 | # print("features=", features) 82 | if features.contains("memory:NVMM"): 83 | # Get the source bin ghost pad 84 | bin_ghost_pad = source_bin.get_static_pad("src") 85 | if not bin_ghost_pad.set_target(decoder_src_pad): 86 | sys.stderr.write("Failed to link decoder src pad to source bin ghost pad\n") 87 | else: 88 | sys.stderr.write(" Error: Decodebin did not pick nvidia decoder plugin.\n") 89 | 90 | def decodebin_child_added(child_proxy, Object, name, user_data): 91 | print("Decodebin child added:", name, "\n") 92 | if name.find("decodebin") != -1: 93 | Object.connect("child-added", decodebin_child_added, user_data) 94 | if is_aarch64() and name.find("nvv4l2decoder") != -1: 95 | print("Seting bufapi_version\n") 96 | Object.set_property("bufapi-version", True) 97 | 98 | def create_source_bin(index, uri): 99 | print("Creating source bin") 100 | 101 | # Create a source GstBin to abstract this bin's content from the rest of the 102 | # pipeline 103 | bin_name = "source-bin-%02d" % index 104 | print(bin_name) 105 | nbin = Gst.Bin.new(bin_name) 106 | if not nbin: 107 | sys.stderr.write(" Unable to create source bin \n") 108 | 109 | # Source element for reading from the uri. 110 | # We will use decodebin and let it figure out the container format of the 111 | # stream and the codec and plug the appropriate demux and decode plugins. 112 | uri_decode_bin = Gst.ElementFactory.make("uridecodebin", "uri-decode-bin") 113 | if not uri_decode_bin: 114 | sys.stderr.write(" Unable to create uri decode bin \n") 115 | # We set the input uri to the source element 116 | uri_decode_bin.set_property("uri", uri) 117 | # Connect to the "pad-added" signal of the decodebin which generates a 118 | # callback once a new pad for raw data has beed created by the decodebin 119 | uri_decode_bin.connect("pad-added", cb_newpad, nbin) 120 | uri_decode_bin.connect("child-added", decodebin_child_added, nbin) 121 | 122 | # We need to create a ghost pad for the source bin which will act as a proxy 123 | # for the video decoder src pad. The ghost pad will not have a target right 124 | # now. Once the decode bin creates the video decoder and generates the 125 | # cb_newpad callback, we will set the ghost pad target to the video decoder 126 | # src pad. 127 | Gst.Bin.add(nbin, uri_decode_bin) 128 | bin_pad = nbin.add_pad(Gst.GhostPad.new_no_target("src", Gst.PadDirection.SRC)) 129 | if not bin_pad: 130 | sys.stderr.write(" Failed to add ghost pad in source bin \n") 131 | return None 132 | return nbin 133 | 134 | def set_tracker_config(cfg_path,tracker): 135 | config = configparser.ConfigParser() 136 | config.read(cfg_path) 137 | config.sections() 138 | 139 | for key in config['tracker']: 140 | if key == 'tracker-width' : 141 | tracker_width = config.getint('tracker', key) 142 | tracker.set_property('tracker-width', tracker_width) 143 | if key == 'tracker-height' : 144 | tracker_height = config.getint('tracker', key) 145 | tracker.set_property('tracker-height', tracker_height) 146 | if key == 'gpu-id' : 147 | tracker_gpu_id = config.getint('tracker', key) 148 | tracker.set_property('gpu_id', tracker_gpu_id) 149 | if key == 'll-lib-file' : 150 | tracker_ll_lib_file = config.get('tracker', key) 151 | tracker.set_property('ll-lib-file', tracker_ll_lib_file) 152 | if key == 'll-config-file' : 153 | tracker_ll_config_file = config.get('tracker', key) 154 | tracker.set_property('ll-config-file', tracker_ll_config_file) 155 | if key == 'enable-batch-process' : 156 | tracker_enable_batch_process = config.getint('tracker', key) 157 | tracker.set_property('enable_batch_process', tracker_enable_batch_process) 158 | if key == 'enable-past-frame' : 159 | tracker_enable_past_frame = config.getint('tracker', key) 160 | tracker.set_property('enable_past_frame', tracker_enable_past_frame) 161 | if key == 'tracking-id-reset-mode' : 162 | tracking_id_reset_mode = config.getint('tracker', key) 163 | tracker.set_property('tracking-id-reset-mode', tracking_id_reset_mode) 164 | if key == 'display-tracking-id' : 165 | display_tracking_id = config.getint('tracker', key) 166 | tracker.set_property('display-tracking-id', display_tracking_id) 167 | 168 | ## postprocess 169 | 170 | def get_outshape(input_shape,stride=8,num_calsses=57,anchor_num=3): 171 | b,c,h,w = input_shape 172 | s = stride 173 | height_out = int(h/s) 174 | width_out = int(w/s) 175 | out_shape = (b,anchor_num,width_out,height_out,num_calsses) 176 | return out_shape 177 | 178 | def get_total_outshape(input_shape,stride=(8,16,32,64),num_calsses=57,anchor_num=3): 179 | b,c,h,w = input_shape 180 | height_out = [int(h/s) for s in stride] 181 | width_out = [int(w/s) for s in stride] 182 | num_grids = sum([i*j for i,j in zip(height_out,width_out)]) 183 | num_preds = anchor_num*num_grids 184 | out_shape = (b,num_preds,num_calsses) 185 | return out_shape 186 | 187 | def postprocess(model_out,im0_shape,im1_shape,im0=None,conf_thres=0.5,iou_thres=0.4,draw=True,line_thickness=3): 188 | pred = non_max_suppression(model_out, conf_thres, iou_thres) 189 | boxes,confs,kpts = [], [], [] 190 | for i, det in enumerate(pred): # detections per image 191 | if len(det): 192 | # Rescale boxes from img_size to im0 size 193 | scale_coords(im1_shape, det[:, :4], im0_shape, kpt_label=False) 194 | scale_coords(im1_shape, det[:, 6:], im0_shape, kpt_label=True, step=3) 195 | 196 | boxes.append(det[:,:4]) 197 | confs.append(det[:,4:5]) 198 | kpts.append(det[:,6:].reshape(det.shape[0],-1,3)) 199 | # Write results 200 | if not draw or im0 is None: 201 | continue 202 | for det_index, (*xyxy, conf, cls) in enumerate(reversed(det[:,:6])): 203 | c = int(cls) # integer class 204 | label = f'{categories[c]}:{conf:.2f}' 205 | kpt = det[det_index, 6:] 206 | plot_one_box(xyxy, im0, colors(c, True), label, line_thickness, kpt_label=True, kpts=kpt, steps=3, orig_shape=im0.shape[:2]) 207 | return boxes, confs, kpts 208 | 209 | def postprocessNoNMS(model_out,im0_shape,im1_shape,im0=None,conf_thres=0.5,iou_thres=0.4,draw=True,line_thickness=3): 210 | # pred = non_max_suppression(model_out, conf_thres, iou_thres) 211 | nmsed_indices,nmsed_boxes,nmsed_poses,nmsed_scores = model_out 212 | batch_index, class_index, box_index = nmsed_indices[...,0],nmsed_indices[...,1],nmsed_indices[...,2] 213 | if np.any(np.isnan(box_index)) or np.all(box_index < 0): 214 | return [], [], [] 215 | keep_num = np.unique(box_index).size 216 | # print(box_index) 217 | assert (box_index[...,keep_num:]-box_index[...,keep_num-1]).sum()==0,f'{box_index}' 218 | nmsed_boxes = nmsed_boxes[:,:keep_num,:] 219 | nmsed_poses = nmsed_poses[:,:keep_num,:] 220 | nmsed_scores = nmsed_scores[:,:keep_num,:] 221 | boxes,confs,kpts = [], [], [] 222 | batch_size = nmsed_boxes.shape[0] 223 | for batch_id in range(batch_size): # detections per image 224 | bboxes = nmsed_boxes[batch_id] 225 | bboxes = xywh2xyxy(bboxes) 226 | pposes = nmsed_poses[batch_id] 227 | scores = nmsed_scores[batch_id] 228 | num_person = len(bboxes) 229 | if num_person>0: 230 | # Rescale boxes from img_size to im0 size 231 | scale_coords(im1_shape, bboxes, im0_shape, kpt_label=False) 232 | scale_coords(im1_shape, pposes, im0_shape, kpt_label=True, step=3) 233 | 234 | boxes.append(bboxes) 235 | confs.append(scores) 236 | kpts.append(pposes.reshape(num_person,-1,3)) 237 | # Write results 238 | if not draw or im0 is None: 239 | continue 240 | for det_index in range(num_person): 241 | xyxy = bboxes[det_index] 242 | conf = scores[det_index] 243 | kpt = pposes[det_index] 244 | c = 0 # yolo pose just has one class `person` 245 | label = f'{categories[c]}:{conf.item():.2f}' 246 | plot_one_box(xyxy, im0, colors(c, True), label, line_thickness, kpt_label=True, kpts=kpt, steps=3, orig_shape=im0.shape[:2]) 247 | return boxes, confs, kpts 248 | 249 | 250 | def non_max_suppression(predictions, conf_thres=0.5, nms_thres=0.4): 251 | """ 252 | description: Removes detections with lower object confidence score than 'conf_thres' and performs 253 | Non-Maximum Suppression to further filter detections. 254 | param: 255 | prediction: detections, (x1, y1, x2, y2, conf, cls_id) 256 | origin_h: original image height 257 | origin_w: original image width 258 | conf_thres: a confidence threshold to filter detections 259 | nms_thres: a iou threshold to filter detections 260 | return: 261 | boxes: output after nms with the shape (x1, y1, x2, y2, conf, cls_id) 262 | """ 263 | output = [np.zeros((0, 57))] * predictions.shape[0] 264 | for i,prediction in enumerate(predictions): 265 | # Get the boxes that score > CONF_THRESH 266 | boxes = prediction[prediction[:, 4] >= conf_thres] 267 | # Trandform bbox from [center_x, center_y, w, h] to [x1, y1, x2, y2] 268 | boxes[:, :4] = xywh2xyxy(boxes[:, :4]) 269 | # Object confidence 270 | confs = boxes[:, 4] * boxes[:, 5] 271 | # Sort by the confs 272 | boxes = boxes[np.argsort(-confs)] 273 | # Perform non-maximum suppression 274 | keep_boxes = [] 275 | while boxes.shape[0]: 276 | large_overlap = bbox_iou(np.expand_dims(boxes[0, :4], 0), boxes[:, :4]) > nms_thres 277 | # label_match = boxes[0, -1] == boxes[:, -1] 278 | # Indices of boxes with lower confidence scores, large IOUs and matching labels 279 | # invalid = large_overlap & label_match 280 | invalid = large_overlap 281 | keep_boxes += [boxes[0]] 282 | boxes = boxes[~invalid] 283 | boxes = np.stack(keep_boxes, 0) if len(keep_boxes) else np.array([]) 284 | output[i] = boxes 285 | return output 286 | 287 | def bbox_iou(box1, box2, x1y1x2y2=True): 288 | """ 289 | description: compute the IoU of two bounding boxes 290 | param: 291 | box1: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h)) 292 | box2: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h)) 293 | x1y1x2y2: select the coordinate format 294 | return: 295 | iou: computed iou 296 | """ 297 | if not x1y1x2y2: 298 | # Transform from center and width to exact coordinates 299 | b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2 300 | b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2 301 | b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2 302 | b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2 303 | else: 304 | # Get the coordinates of bounding boxes 305 | b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3] 306 | b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3] 307 | 308 | # Get the coordinates of the intersection rectangle 309 | inter_rect_x1 = np.maximum(b1_x1, b2_x1) 310 | inter_rect_y1 = np.maximum(b1_y1, b2_y1) 311 | inter_rect_x2 = np.minimum(b1_x2, b2_x2) 312 | inter_rect_y2 = np.minimum(b1_y2, b2_y2) 313 | # Intersection area 314 | inter_area = np.clip(inter_rect_x2 - inter_rect_x1 + 1, 0, None) * \ 315 | np.clip(inter_rect_y2 - inter_rect_y1 + 1, 0, None) 316 | # Union Area 317 | b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1) 318 | b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1) 319 | 320 | iou = inter_area / (b1_area + b2_area - inter_area + 1e-16) 321 | 322 | return iou 323 | 324 | def xywh2xyxy(x): 325 | # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right 326 | y = np.copy(x) 327 | y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x 328 | y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y 329 | y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x 330 | y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y 331 | return y 332 | 333 | def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None, kpt_label=False, step=2): 334 | # Rescale coords (xyxy) from img1_shape to img0_shape 335 | if ratio_pad is None: # calculate from img0_shape 336 | gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new 337 | pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding 338 | else: 339 | gain = ratio_pad[0] 340 | pad = ratio_pad[1] 341 | if isinstance(gain, (list, tuple)): 342 | gain = gain[0] 343 | if not kpt_label: 344 | coords[:, [0, 2]] -= pad[0] # x padding 345 | coords[:, [1, 3]] -= pad[1] # y padding 346 | coords[:, [0, 2]] /= gain 347 | coords[:, [1, 3]] /= gain 348 | clip_coords(coords[: ,0:4], img0_shape) 349 | #coords[:, 0:4] = coords[:, 0:4].round() 350 | else: 351 | coords[:, 0::step] -= pad[0] # x padding 352 | coords[:, 1::step] -= pad[1] # y padding 353 | coords[:, 0::step] /= gain 354 | coords[:, 1::step] /= gain 355 | clip_coords(coords, img0_shape, step=step) 356 | #coords = coords.round() 357 | return coords 358 | 359 | def clip_coords(boxes, img_shape, step=2): 360 | # Clip bounding xyxy bounding boxes to image shape (height, width) 361 | boxes[:, 0::step] = np.clip(boxes[:, 0::step], 0, img_shape[1]) 362 | boxes[:, 1::step] = np.clip(boxes[:, 1::step], 0, img_shape[0]) 363 | 364 | def decode(model_outs,input_shape=None,strides=[8,16,32,64]): 365 | anchors = { 8:[ 19,27, 44,40, 38,94 ], 366 | 16:[ 96,68, 86,152, 180,137 ], 367 | 32:[ 140,301, 303,264, 238,542 ], 368 | 64:[ 436,615, 739,380, 925,792 ] } 369 | decode_outs = [] 370 | for i,out in enumerate(model_outs): 371 | b,na,nx,ny,nkpt = out.shape 372 | if input_shape is not None: 373 | s_x, s_y = input_shape[0]/ny,input_shape[1]/nx 374 | assert s_x == s_y, f'stride not equal.' 375 | stride = int(s_x) 376 | else: 377 | stride = strides[i] 378 | kpt_grid_x,kpt_grid_y = np.meshgrid(np.arange(nx),np.arange(ny)) 379 | kpt_grid_x = kpt_grid_x[np.newaxis,np.newaxis,:,:,np.newaxis] 380 | kpt_grid_y = kpt_grid_y[np.newaxis,np.newaxis,:,:,np.newaxis] 381 | grid = np.concatenate([kpt_grid_x,kpt_grid_y],axis=-1) 382 | anchor_grid = np.array(anchors[stride]).reshape(na,-1) 383 | anchor_grid = anchor_grid[np.newaxis,:,np.newaxis,np.newaxis,:] 384 | det = out[..., :6] 385 | kpt = out[..., 6:] 386 | det = 1/(1 + np.exp(-det)) # sigmoid 387 | 388 | det[..., 0:2] = (det[..., 0:2] * 2. - 0.5 + grid) * stride # xy 389 | det[..., 2:4] = (det[..., 2:4] * 2) ** 2 * anchor_grid.reshape(1, na, 1, 1, 2) # wh 390 | 391 | kpt[..., 0::3] = (kpt[..., ::3] * 2. - 0.5 + kpt_grid_x.repeat(17,axis=-1)) * stride # xy 392 | kpt[..., 1::3] = (kpt[..., 1::3] * 2. - 0.5 + kpt_grid_y.repeat(17,axis=-1)) * stride # xy 393 | kpt[..., 2::3] = 1/(1 + np.exp(-kpt[..., 2::3])) 394 | 395 | y = np.concatenate((det, kpt), axis=-1) 396 | decode_outs.append(y.reshape(b,-1,nkpt)) 397 | decode_outs = np.concatenate(decode_outs,axis=1) 398 | return decode_outs 399 | 400 | 401 | ## plot 402 | 403 | def plot_one_box(x, im, color=None, label=None, line_thickness=3, kpt_label=False, kpts=None, steps=2, orig_shape=None): 404 | # Plots one bounding box on image 'im' using OpenCV 405 | assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to plot_on_box() input image.' 406 | tl = line_thickness or round(0.002 * (im.shape[0] + im.shape[1]) / 2) + 1 # line/font thickness 407 | color = color or [random.randint(0, 255) for _ in range(3)] 408 | c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3])) 409 | cv2.rectangle(im, c1, c2, (255,0,0), thickness=tl*1//3, lineType=cv2.LINE_AA) 410 | if label: 411 | tf = max(tl - 1, 1) # font thickness 412 | t_size = cv2.getTextSize(label, 0, fontScale=tl / 6, thickness=tf)[0] 413 | c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 414 | cv2.rectangle(im, c1, c2, color, -1, cv2.LINE_AA) # filled 415 | cv2.putText(im, label, (c1[0], c1[1] - 2), 0, tl / 6, [225, 255, 255], thickness=tf//2, lineType=cv2.LINE_AA) 416 | if kpt_label: 417 | plot_skeleton_kpts(im, kpts, steps, orig_shape=orig_shape) 418 | 419 | def plot_skeleton_kpts(im, kpts, steps, orig_shape=None): 420 | #Plot the skeleton and keypointsfor coco datatset 421 | palette = np.array([[255, 128, 0], [255, 153, 51], [255, 178, 102], 422 | [230, 230, 0], [255, 153, 255], [153, 204, 255], 423 | [255, 102, 255], [255, 51, 255], [102, 178, 255], 424 | [51, 153, 255], [255, 153, 153], [255, 102, 102], 425 | [255, 51, 51], [153, 255, 153], [102, 255, 102], 426 | [51, 255, 51], [0, 255, 0], [0, 0, 255], [255, 0, 0], 427 | [255, 255, 255]]) 428 | 429 | skeleton = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12], 430 | [7, 13], [6, 7], [6, 8], [7, 9], [8, 10], [9, 11], [2, 3], 431 | [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]] 432 | 433 | pose_limb_color = palette[[9, 9, 9, 9, 7, 7, 7, 0, 0, 0, 0, 0, 16, 16, 16, 16, 16, 16, 16]] 434 | pose_kpt_color = palette[[16, 16, 16, 16, 16, 0, 0, 0, 0, 0, 0, 9, 9, 9, 9, 9, 9]] 435 | radius = 5 436 | num_kpts = len(kpts) // steps 437 | 438 | for kid in range(num_kpts): 439 | r, g, b = pose_kpt_color[kid] 440 | x_coord, y_coord = kpts[steps * kid], kpts[steps * kid + 1] 441 | if not (x_coord % 640 == 0 or y_coord % 640 == 0): 442 | if steps == 3: 443 | conf = kpts[steps * kid + 2] 444 | if conf < 0.5: 445 | continue 446 | cv2.circle(im, (int(x_coord), int(y_coord)), radius, (int(r), int(g), int(b)), -1) 447 | 448 | for sk_id, sk in enumerate(skeleton): 449 | r, g, b = pose_limb_color[sk_id] 450 | pos1 = (int(kpts[(sk[0]-1)*steps]), int(kpts[(sk[0]-1)*steps+1])) 451 | pos2 = (int(kpts[(sk[1]-1)*steps]), int(kpts[(sk[1]-1)*steps+1])) 452 | if steps == 3: 453 | conf1 = kpts[(sk[0]-1)*steps+2] 454 | conf2 = kpts[(sk[1]-1)*steps+2] 455 | if conf1<0.5 or conf2<0.5: 456 | continue 457 | if pos1[0]%640 == 0 or pos1[1]%640==0 or pos1[0]<0 or pos1[1]<0: 458 | continue 459 | if pos2[0] % 640 == 0 or pos2[1] % 640 == 0 or pos2[0]<0 or pos2[1]<0: 460 | continue 461 | cv2.line(im, pos1, pos2, (int(r), int(g), int(b)), thickness=2) 462 | 463 | --------------------------------------------------------------------------------