├── .github └── workflows │ └── blossom-ci.yml ├── .gitignore ├── License.md ├── README.md ├── docker └── Dockerfile ├── docs └── configuring_the_client.md ├── images └── pose_cls_e2e_inference.png ├── misc └── cla.md ├── model_repository ├── centerpose_tao │ └── config.pbtxt ├── dashcamnet_tao │ └── config.pbtxt ├── foundationpose_refiner_tao │ └── config.pbtxt ├── foundationpose_scorer_tao │ └── config.pbtxt ├── lprnet_tao │ ├── characters_list.txt │ └── config.pbtxt ├── multitask_classification_tao │ └── config.pbtxt ├── peoplenet_tao │ └── config.pbtxt ├── peoplesegnet_tao │ └── config.pbtxt ├── pose_classification_tao │ ├── config.pbtxt │ └── labels.txt ├── re_identification_tao │ └── config.pbtxt ├── retinanet_tao │ └── config.pbtxt ├── vehicletypenet_tao │ ├── config.pbtxt │ └── labels.txt ├── visual_changenet_segmentation_tao │ └── config.pbtxt └── yolov3_tao │ └── config.pbtxt ├── scripts ├── config.sh ├── download_and_convert.sh ├── pose_cls_e2e_inference │ ├── demo.mp4 │ ├── download_and_convert.sh │ ├── plot_e2e_inference.py │ ├── start_client.sh │ └── start_server.sh ├── re_id_e2e_inference │ ├── download_and_convert.sh │ ├── plot_e2e_inference.py │ ├── re_ranking.py │ ├── sample_dataset.py │ ├── start_client.sh │ ├── start_server.sh │ └── utils.py └── start_server.sh └── tao_triton ├── python ├── __init__.py ├── clustering_specs │ ├── clustering_config_centerpose.prototxt │ ├── clustering_config_dashcamnet.prototxt │ └── clustering_config_peoplenet.prototxt ├── dataset_convert_specs │ └── dataset_convert_config_pose_classification.yaml ├── entrypoints │ ├── __init__.py │ └── tao_client.py ├── model │ ├── __init__.py │ ├── centerpose_model.py │ ├── classification_model.py │ ├── detectnet_model.py │ ├── foundationpose_model.py │ ├── lprnet_model.py │ ├── multitask_classification_model.py │ ├── peoplesegnet_model.py │ ├── pose_classification_model.py │ ├── re_identification_model.py │ ├── retinanet_model.py │ ├── triton_model.py │ ├── visual_changenet_model.py │ └── yolov3_model.py ├── postprocessing │ ├── __init__.py │ ├── centerpose_postprocessor.py │ ├── classification_postprocessor.py │ ├── detectnet_processor.py │ ├── foundationpose_postprocessor.py │ ├── lprnet_postprocessor.py │ ├── multitask_classification_postprocessor.py │ ├── peoplesegnet_postprocessor.py │ ├── pose_classification_postprocessor.py │ ├── postprocessor.py │ ├── re_identification_postprocessor.py │ ├── retinanet_postprocessor.py │ ├── utils.py │ ├── visual_changenet_postprocessor.py │ └── yolov3_postprocessor.py ├── proto │ ├── __init__.py │ ├── postprocessor_config.proto │ └── postprocessor_config_pb2.py ├── types │ ├── __init__.py │ ├── annotation.py │ ├── frame.py │ └── user_data.py └── utils │ ├── __init__.py │ ├── kitti.py │ ├── pose_cls_dataset_convert.py │ └── preprocess_input.py ├── requirements-pip.txt └── setup.py /.github/workflows/blossom-ci.yml: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2020-2021, NVIDIA CORPORATION. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # A workflow to trigger ci on hybrid infra (github + self hosted runner) 16 | name: Blossom-CI 17 | on: 18 | issue_comment: 19 | types: [created] 20 | workflow_dispatch: 21 | inputs: 22 | platform: 23 | description: 'runs-on argument' 24 | required: false 25 | args: 26 | description: 'argument' 27 | required: false 28 | jobs: 29 | Authorization: 30 | name: Authorization 31 | runs-on: blossom 32 | outputs: 33 | args: ${{ env.args }} 34 | 35 | # This job only runs for pull request comments 36 | if: contains('\ 37 | Arun-George-Zachariah,\ 38 | ', format('{0},', github.actor)) 39 | steps: 40 | - name: Check if comment is issued by authorized person 41 | run: blossom-ci 42 | env: 43 | OPERATION: 'AUTH' 44 | REPO_TOKEN: ${{ secrets.GITHUB_TOKEN }} 45 | REPO_KEY_DATA: ${{ secrets.BLOSSOM_KEY }} 46 | 47 | Vulnerability-scan: 48 | name: Vulnerability scan 49 | needs: [Authorization] 50 | runs-on: ubuntu-latest 51 | steps: 52 | - name: Checkout code 53 | uses: actions/checkout@v3 54 | with: 55 | repository: ${{ fromJson(needs.Authorization.outputs.args).repo }} 56 | ref: ${{ fromJson(needs.Authorization.outputs.args).ref }} 57 | lfs: 'true' 58 | 59 | - name: Setup blackduck properties 60 | run: | 61 | echo detect.excluded.detector.types=PIP >> application.properties 62 | 63 | - name: Run blossom action 64 | uses: NVIDIA/blossom-action@main 65 | env: 66 | REPO_TOKEN: ${{ secrets.GITHUB_TOKEN }} 67 | REPO_KEY_DATA: ${{ secrets.BLOSSOM_KEY }} 68 | with: 69 | args1: ${{ fromJson(needs.Authorization.outputs.args).args1 }} 70 | args2: ${{ fromJson(needs.Authorization.outputs.args).args2 }} 71 | args3: ${{ fromJson(needs.Authorization.outputs.args).args3 }} 72 | 73 | Job-trigger: 74 | name: Start ci job 75 | needs: [Vulnerability-scan] 76 | runs-on: blossom 77 | steps: 78 | - name: Start ci job 79 | run: blossom-ci 80 | env: 81 | OPERATION: 'START-CI-JOB' 82 | CI_SERVER: ${{ secrets.CI_SERVER }} 83 | REPO_TOKEN: ${{ secrets.GITHUB_TOKEN }} 84 | 85 | Upload-Log: 86 | name: Upload log 87 | runs-on: blossom 88 | if : github.event_name == 'workflow_dispatch' 89 | steps: 90 | - name: Jenkins log for pull request ${{ fromJson(github.event.inputs.args).pr }} (click here) 91 | run: blossom-ci 92 | env: 93 | OPERATION: 'POST-PROCESSING' 94 | CI_SERVER: ${{ secrets.CI_SERVER }} 95 | REPO_TOKEN: ${{ secrets.GITHUB_TOKEN }} 96 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Ignore pycache files. 2 | **pycache** 3 | *.pyc 4 | -------------------------------------------------------------------------------- /License.md: -------------------------------------------------------------------------------- 1 | Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining 4 | a copy of this software and associated documentation files (the 5 | "Software"), to deal in the Software without restriction, including 6 | without limitation the rights to use, copy, modify, merge, publish, 7 | distribute, sublicense, and/or sell copies of the Software, and to 8 | permit persons to whom the Software is furnished to do so, subject to 9 | the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be 12 | included in all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | -------------------------------------------------------------------------------- /docker/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM nvcr.io/nvidia/tritonserver:23.02-py3 2 | 3 | # Installing TRT OSS to the Triton Container. 4 | RUN apt-get update && \ 5 | apt-get install -y pkg-config && \ 6 | apt-get install -y git && \ 7 | apt-get install -y zlib1g && \ 8 | apt-get install -y zlib1g-dev 9 | RUN cd /tmp \ 10 | && wget https://github.com/Kitware/CMake/releases/download/v3.14.4/cmake-3.14.4-Linux-x86_64.sh \ 11 | && chmod +x cmake-3.14.4-Linux-x86_64.sh \ 12 | && ./cmake-3.14.4-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license \ 13 | && rm ./cmake-3.14.4-Linux-x86_64.sh \ 14 | && cd - 15 | RUN cd /opt 16 | RUN ln -s /usr/bin/python3 /usr/bin/python 17 | # Clone and checkout TensorRT OSS 18 | # Moving TensorRT to 8.5 branch. 19 | ENV TRT_TAG "release/8.5" 20 | ENV TRT_INCLUDE_DIR="/usr/include/x86_64-linux-gnu" 21 | # Install TRT OSS 22 | RUN mkdir trt_oss_src && \ 23 | cd trt_oss_src && \ 24 | echo "$PWD Building TRT OSS..." && \ 25 | git clone -b $TRT_TAG https://github.com/NVIDIA/TensorRT.git TensorRT && \ 26 | cd TensorRT && \ 27 | git submodule update --init --recursive && \ 28 | mkdir -p build && cd build && \ 29 | cmake .. -DGPU_ARCHS="53 60 61 70 75 80 86 90" -DTRT_LIB_DIR=/usr/lib/x86_64-linux-gnu -DTRT_BIN_DIR=`pwd`/out -DCUDA_VERSION=11.8 -DCUDNN_VERSION=8.7 && \ 30 | make -j32 && \ 31 | cp libnvinfer_plugin.so.8.5.3 /usr/lib/x86_64-linux-gnu/libnvinfer_plugin.so.8.5.3 && \ 32 | cp libnvinfer_plugin_static.a /usr/lib/x86_64-linux-gnu/libnvinfer_plugin_static.a && \ 33 | cp libnvonnxparser.so.8.5.3 /usr/lib/x86_64-linux-gnu/libnvonnxparser.so.8.5.3 && \ 34 | cp libnvcaffeparser.so.8.5.3 /usr/lib/x86_64-linux-gnu/libnvcaffeparser.so.8.5.3 && \ 35 | cp trtexec /usr/local/bin/ && \ 36 | cd ../../../ && \ 37 | rm -rf trt_oss_src 38 | 39 | # Setting up TensorRT Paths. 40 | ENV TRT_LIB_PATH=/usr/lib/x86_64-linux-gnu 41 | ENV TRT_INC_PATH=/usr/include/x86_64-linux-gnu 42 | 43 | # Download and install TAO Toolkit converter 44 | RUN mkdir /opt/tao-converter && \ 45 | wget --content-disposition https://api.ngc.nvidia.com/v2/resources/nvidia/tao/tao-converter/versions/v3.22.05_trt8.2_x86/zip -O /opt/tao-converter/tao-converter_v3.22.05_trt8.2_x86.zip && \ 46 | apt update && apt install unzip libssl-dev -y && \ 47 | unzip /opt/tao-converter/tao-converter_v3.22.05_trt8.2_x86.zip -d /opt/tao-converter && \ 48 | chmod +x /opt/tao-converter/tao-converter && \ 49 | rm -rf /opt/tao-converter/tao-converter_v3.22.05_trt8.2_x86.zip 50 | 51 | ENV PATH=/opt/tao-converter:$PATH 52 | 53 | CMD ["/bin/bash"] 54 | -------------------------------------------------------------------------------- /images/pose_cls_e2e_inference.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps/9a30f9692bf29fb728520e9dba1c79be2bf65e74/images/pose_cls_e2e_inference.png -------------------------------------------------------------------------------- /misc/cla.md: -------------------------------------------------------------------------------- 1 | ## Individual Contributor License Agreement (CLA) 2 | 3 | **Thank you for submitting your contributions to this project.** 4 | 5 | By signing this CLA, you agree that the following terms apply to all of your past, present and future contributions 6 | to the project. 7 | 8 | ### License. 9 | 10 | You hereby represent that all present, past and future contributions are governed by the 11 | [MIT License](https://opensource.org/licenses/MIT) 12 | copyright statement. 13 | 14 | This entails that to the extent possible under law, you transfer all copyright and related or neighboring rights 15 | of the code or documents you contribute to the project itself or its maintainers. 16 | Furthermore you also represent that you have the authority to perform the above waiver 17 | with respect to the entirety of you contributions. 18 | 19 | ### Moral Rights. 20 | 21 | To the fullest extent permitted under applicable law, you hereby waive, and agree not to 22 | assert, all of your “moral rights” in or relating to your contributions for the benefit of the project. 23 | 24 | ### Third Party Content. 25 | 26 | If your Contribution includes or is based on any source code, object code, bug fixes, configuration changes, tools, 27 | specifications, documentation, data, materials, feedback, information or other works of authorship that were not 28 | authored by you (“Third Party Content”) or if you are aware of any third party intellectual property or proprietary 29 | rights associated with your Contribution (“Third Party Rights”), 30 | then you agree to include with the submission of your Contribution full details respecting such Third Party 31 | Content and Third Party Rights, including, without limitation, identification of which aspects of your 32 | Contribution contain Third Party Content or are associated with Third Party Rights, the owner/author of the 33 | Third Party Content and Third Party Rights, where you obtained the Third Party Content, and any applicable 34 | third party license terms or restrictions respecting the Third Party Content and Third Party Rights. For greater 35 | certainty, the foregoing obligations respecting the identification of Third Party Content and Third Party Rights 36 | do not apply to any portion of a Project that is incorporated into your Contribution to that same Project. 37 | 38 | ### Representations. 39 | 40 | You represent that, other than the Third Party Content and Third Party Rights identified by 41 | you in accordance with this Agreement, you are the sole author of your Contributions and are legally entitled 42 | to grant the foregoing licenses and waivers in respect of your Contributions. If your Contributions were 43 | created in the course of your employment with your past or present employer(s), you represent that such 44 | employer(s) has authorized you to make your Contributions on behalf of such employer(s) or such employer 45 | (s) has waived all of their right, title or interest in or to your Contributions. 46 | 47 | ### Disclaimer. 48 | 49 | To the fullest extent permitted under applicable law, your Contributions are provided on an "as is" 50 | basis, without any warranties or conditions, express or implied, including, without limitation, any implied 51 | warranties or conditions of non-infringement, merchantability or fitness for a particular purpose. You are not 52 | required to provide support for your Contributions, except to the extent you desire to provide support. 53 | 54 | ### No Obligation. 55 | 56 | You acknowledge that the maintainers of this project are under no obligation to use or incorporate your contributions 57 | into the project. The decision to use or incorporate your contributions into the project will be made at the 58 | sole discretion of the maintainers or their authorized delegates. 59 | -------------------------------------------------------------------------------- /model_repository/centerpose_tao/config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "centerpose_tao" 2 | platform: "tensorrt_plan" 3 | max_batch_size: 16 4 | input [ 5 | { 6 | name: "input" 7 | data_type: TYPE_FP32 8 | format: FORMAT_NCHW 9 | dims: [ 3, 512, 512 ] 10 | } 11 | ] 12 | output [ 13 | { 14 | name: "bboxes" 15 | data_type: TYPE_FP32 16 | dims: [ 100, 4 ] 17 | }, 18 | { 19 | name: "scores" 20 | data_type: TYPE_FP32 21 | dims: [ 100, 1 ] 22 | }, 23 | { 24 | name: "kps" 25 | data_type: TYPE_FP32 26 | dims: [ 100, 16 ] 27 | }, 28 | { 29 | name: "clses" 30 | data_type: TYPE_FP32 31 | dims: [ 100, 1 ] 32 | }, 33 | { 34 | name: "obj_scale" 35 | data_type: TYPE_FP32 36 | dims: [ 100, 3 ] 37 | }, 38 | { 39 | name: "kps_displacement_mean" 40 | data_type: TYPE_FP32 41 | dims: [ 100, 16 ] 42 | }, 43 | { 44 | name: "kps_heatmap_mean" 45 | data_type: TYPE_FP32 46 | dims: [ 100, 16 ] 47 | } 48 | ] 49 | dynamic_batching { } -------------------------------------------------------------------------------- /model_repository/dashcamnet_tao/config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "dashcamnet_tao" 2 | platform: "tensorrt_plan" 3 | max_batch_size: 16 4 | input [ 5 | { 6 | name: "input_1:0" 7 | data_type: TYPE_FP32 8 | format: FORMAT_NCHW 9 | dims: [ 3, 544, 960 ] 10 | } 11 | ] 12 | output [ 13 | { 14 | name: "output_bbox/BiasAdd:0" 15 | data_type: TYPE_FP32 16 | dims: [ 16, 34, 60 ] 17 | }, 18 | { 19 | name: "output_cov/Sigmoid:0" 20 | data_type: TYPE_FP32 21 | dims: [ 4, 34, 60 ] 22 | } 23 | ] 24 | dynamic_batching { } 25 | -------------------------------------------------------------------------------- /model_repository/foundationpose_refiner_tao/config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "foundationpose_refiner_tao" 2 | platform: "tensorrt_plan" 3 | max_batch_size: 252 4 | 5 | input [ 6 | { 7 | name: "inputA" 8 | data_type: TYPE_FP32 9 | format: FORMAT_NCHW 10 | dims: [ 6, 160, 160 ] 11 | }, 12 | { 13 | name: "inputB" 14 | data_type: TYPE_FP32 15 | format: FORMAT_NCHW 16 | dims: [ 6, 160, 160 ] 17 | } 18 | ] 19 | output [ 20 | { 21 | name: "trans" 22 | data_type: TYPE_FP32 23 | dims: [ 3 ] 24 | }, 25 | { 26 | name: "rot" 27 | data_type: TYPE_FP32 28 | dims: [ 3 ] 29 | } 30 | ] 31 | dynamic_batching { } -------------------------------------------------------------------------------- /model_repository/foundationpose_scorer_tao/config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "foundationpose_scorer_tao" 2 | platform: "tensorrt_plan" 3 | max_batch_size: 252 4 | 5 | input [ 6 | { 7 | name: "inputA" 8 | data_type: TYPE_FP32 9 | format: FORMAT_NCHW 10 | dims: [ 6, 160, 160 ] 11 | }, 12 | { 13 | name: "inputB" 14 | data_type: TYPE_FP32 15 | format: FORMAT_NCHW 16 | dims: [ 6, 160, 160 ] 17 | } 18 | ] 19 | output [ 20 | { 21 | name: "score_logit" 22 | data_type: TYPE_FP32 23 | dims: [ 1 ] 24 | } 25 | ] 26 | dynamic_batching { } -------------------------------------------------------------------------------- /model_repository/lprnet_tao/characters_list.txt: -------------------------------------------------------------------------------- 1 | 0 2 | 1 3 | 2 4 | 3 5 | 4 6 | 5 7 | 6 8 | 7 9 | 8 10 | 9 11 | A 12 | B 13 | C 14 | D 15 | E 16 | F 17 | G 18 | H 19 | I 20 | J 21 | K 22 | L 23 | M 24 | N 25 | P 26 | Q 27 | R 28 | S 29 | T 30 | U 31 | V 32 | W 33 | X 34 | Y 35 | Z 36 | -------------------------------------------------------------------------------- /model_repository/lprnet_tao/config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "lprnet_tao" 2 | platform: "tensorrt_plan" 3 | max_batch_size: 16 4 | input [ 5 | { 6 | name: "image_input" 7 | data_type: TYPE_FP32 8 | format: FORMAT_NCHW 9 | dims: [ 3, 48, 96 ] 10 | } 11 | ] 12 | output [ 13 | { 14 | name: "tf_op_layer_ArgMax" 15 | data_type: TYPE_INT32 16 | dims: [ 24 ] 17 | }, 18 | { 19 | name: "tf_op_layer_Max" 20 | data_type: TYPE_FP32 21 | dims: [ 24 ] 22 | } 23 | ] 24 | dynamic_batching { } 25 | -------------------------------------------------------------------------------- /model_repository/multitask_classification_tao/config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "multitask_classification_tao" 2 | platform: "tensorrt_plan" 3 | max_batch_size: 16 4 | input [ 5 | { 6 | name: "input_1" 7 | data_type: TYPE_FP32 8 | format: FORMAT_NCHW 9 | dims: [ 3, 80, 60 ] 10 | } 11 | ] 12 | output [ 13 | { 14 | name: "season/Softmax" 15 | data_type: TYPE_FP32 16 | dims: [ 4, 1, 1 ] 17 | }, 18 | { 19 | name: "category/Softmax" 20 | data_type: TYPE_FP32 21 | dims: [ 10, 1, 1 ] 22 | }, 23 | { 24 | name: "base_color/Softmax" 25 | data_type: TYPE_FP32 26 | dims: [ 11, 1, 1 ] 27 | } 28 | ] 29 | dynamic_batching { } 30 | -------------------------------------------------------------------------------- /model_repository/peoplenet_tao/config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "peoplenet_tao" 2 | platform: "tensorrt_plan" 3 | max_batch_size: 16 4 | input [ 5 | { 6 | name: "input_1:0" 7 | data_type: TYPE_FP32 8 | format: FORMAT_NCHW 9 | dims: [ 3, 544, 960 ] 10 | } 11 | ] 12 | output [ 13 | { 14 | name: "output_bbox/BiasAdd:0" 15 | data_type: TYPE_FP32 16 | dims: [ 12, 34, 60 ] 17 | }, 18 | { 19 | name: "output_cov/Sigmoid:0" 20 | data_type: TYPE_FP32 21 | dims: [ 3, 34, 60 ] 22 | } 23 | ] 24 | dynamic_batching { } 25 | -------------------------------------------------------------------------------- /model_repository/peoplesegnet_tao/config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "peoplesegnet_tao" 2 | platform: "tensorrt_plan" 3 | max_batch_size: 16 4 | input [ 5 | { 6 | name: "Input" 7 | data_type: TYPE_FP32 8 | format: FORMAT_NCHW 9 | dims: [ 3, 576, 960 ] 10 | } 11 | ] 12 | output [ 13 | { 14 | name: "generate_detections" 15 | data_type: TYPE_FP32 16 | dims: [ 100, 6 ] 17 | }, 18 | { 19 | name: "mask_fcn_logits/BiasAdd" 20 | data_type: TYPE_FP32 21 | dims: [ 100, 2, 28, 28 ] 22 | } 23 | ] 24 | dynamic_batching { } 25 | -------------------------------------------------------------------------------- /model_repository/pose_classification_tao/config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "pose_classification_tao" 2 | platform: "tensorrt_plan" 3 | max_batch_size: 16 4 | input [ 5 | { 6 | name: "input" 7 | data_type: TYPE_FP32 8 | dims: [ 3, 300, 34, 1 ] 9 | } 10 | ] 11 | output [ 12 | { 13 | name: "fc_pred" 14 | data_type: TYPE_FP32 15 | dims: [ 6 ] 16 | label_filename: "labels.txt" 17 | } 18 | ] 19 | dynamic_batching { } 20 | -------------------------------------------------------------------------------- /model_repository/pose_classification_tao/labels.txt: -------------------------------------------------------------------------------- 1 | sitting_down 2 | getting_up 3 | sitting 4 | standing 5 | walking 6 | jumping -------------------------------------------------------------------------------- /model_repository/re_identification_tao/config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "re_identification_tao" 2 | platform: "tensorrt_plan" 3 | max_batch_size: 16 4 | input [ 5 | { 6 | name: "input" 7 | data_type: TYPE_FP32 8 | format: FORMAT_NCHW 9 | dims: [ 3, 256, 128 ] 10 | } 11 | ] 12 | output [ 13 | { 14 | name: "fc_pred" 15 | data_type: TYPE_FP32 16 | dims: [ 256 ] 17 | } 18 | ] 19 | dynamic_batching { } 20 | -------------------------------------------------------------------------------- /model_repository/retinanet_tao/config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "retinanet_tao" 2 | platform: "tensorrt_plan" 3 | max_batch_size: 16 4 | input [ 5 | { 6 | name: "Input" 7 | data_type: TYPE_FP32 8 | format: FORMAT_NCHW 9 | dims: [ 3, 544, 960 ] 10 | } 11 | ] 12 | output [ 13 | { 14 | name: "NMS" 15 | data_type: TYPE_FP32 16 | dims: [ 1, 250, 7 ] 17 | }, 18 | { 19 | name: "NMS_1" 20 | data_type: TYPE_FP32 21 | dims: [ 1, 1, 1 ] 22 | } 23 | ] 24 | dynamic_batching { } 25 | -------------------------------------------------------------------------------- /model_repository/vehicletypenet_tao/config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "vehicletypenet_tao" 2 | platform: "tensorrt_plan" 3 | max_batch_size : 1 4 | input [ 5 | { 6 | name: "input_1" 7 | data_type: TYPE_FP32 8 | format: FORMAT_NCHW 9 | dims: [ 3, 224, 224 ] 10 | } 11 | ] 12 | output [ 13 | { 14 | name: "predictions/Softmax" 15 | data_type: TYPE_FP32 16 | dims: [6, 1, 1] 17 | label_filename: "labels.txt" 18 | } 19 | ] 20 | dynamic_batching { } -------------------------------------------------------------------------------- /model_repository/vehicletypenet_tao/labels.txt: -------------------------------------------------------------------------------- 1 | coupe 2 | largevehicle 3 | sedan 4 | suv 5 | truck 6 | van -------------------------------------------------------------------------------- /model_repository/visual_changenet_segmentation_tao/config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "visual_changenet_segmentation_tao" 2 | platform: "tensorrt_plan" 3 | max_batch_size: 1 4 | input [ 5 | { 6 | name: "input0" 7 | data_type: TYPE_FP32 8 | format: FORMAT_NCHW 9 | dims: [ 3, 256, 256 ] 10 | }, 11 | { 12 | name: "input1" 13 | data_type: TYPE_FP32 14 | format: FORMAT_NCHW 15 | dims: [ 3, 256, 256 ] 16 | } 17 | 18 | ] 19 | output [ 20 | { 21 | name: "output_final" 22 | data_type: TYPE_FP32 23 | dims: [ 2, 256, 256 ] 24 | }, 25 | { 26 | name: "output0" 27 | data_type: TYPE_FP32 28 | dims: [ 2, 16, 16 ] 29 | }, 30 | { 31 | name: "output1" 32 | data_type: TYPE_FP32 33 | dims: [ 2, 16, 16 ] 34 | }, 35 | { 36 | name: "output2" 37 | data_type: TYPE_FP32 38 | dims: [ 2, 32, 32 ] 39 | }, 40 | { 41 | name: "output3" 42 | data_type: TYPE_FP32 43 | dims: [ 2, 64, 64 ] 44 | } 45 | ] 46 | dynamic_batching { } 47 | -------------------------------------------------------------------------------- /model_repository/yolov3_tao/config.pbtxt: -------------------------------------------------------------------------------- 1 | name: "yolov3_tao" 2 | platform: "tensorrt_plan" 3 | max_batch_size: 16 4 | input [ 5 | { 6 | name: "Input" 7 | data_type: TYPE_FP32 8 | format: FORMAT_NCHW 9 | dims: [ 3, 544, 960 ] 10 | } 11 | ] 12 | output [ 13 | { 14 | name: "BatchedNMS" 15 | data_type: TYPE_INT32 16 | dims: [ 1 ] 17 | }, 18 | { 19 | name: "BatchedNMS_1" 20 | data_type: TYPE_FP32 21 | dims: [ 200, 4 ] 22 | }, 23 | { 24 | name: "BatchedNMS_2" 25 | data_type: TYPE_FP32 26 | dims: [ 200 ] 27 | }, 28 | { 29 | name: "BatchedNMS_3" 30 | data_type: TYPE_FP32 31 | dims: [ 200 ] 32 | } 33 | ] 34 | dynamic_batching { } 35 | -------------------------------------------------------------------------------- /scripts/config.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining 6 | # a copy of this software and associated documentation files (the 7 | # "Software"), to deal in the Software without restriction, including 8 | # without limitation the rights to use, copy, modify, merge, publish, 9 | # distribute, sublicense, and/or sell copies of the Software, and to 10 | # permit persons to whom the Software is furnished to do so, subject to 11 | # the following conditions: 12 | # 13 | # The above copyright notice and this permission notice shall be 14 | # included in all copies or substantial portions of the Software. 15 | # 16 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 17 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 18 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 19 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 20 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 21 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 22 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 23 | 24 | tao_triton_root=$PWD 25 | gpu_id=0 26 | cuda_ver=12.0 27 | tao_triton_server_docker="nvcr.io/nvidia/tao/triton-apps" 28 | tao_triton_server_tag="23.02-py3" 29 | 30 | # Load key for the models. 31 | tlt_key_peoplenet="tlt_encode" 32 | tlt_key_dashcamnet="tlt_encode" 33 | tlt_key_vehicletypenet="tlt_encode" 34 | tlt_key_lprnet="nvidia_tlt" 35 | tlt_key_yolov3="nvidia_tlt" 36 | tlt_key_peoplesegnet="nvidia_tlt" 37 | tlt_key_retinanet="nvidia_tlt" 38 | tlt_key_multitask_classification="nvidia_tlt" 39 | tlt_key_pose_classification="nvidia_tao" 40 | tlt_key_re_identification="nvidia_tao" 41 | 42 | # Setting model version to run inference on. 43 | peoplenet_version="pruned_quantized_decrypted_v2.3.3" 44 | dashcamnet_version="pruned_v1.0.4" 45 | vehicletypenet_version="pruned_v1.0.1" 46 | visual_changenet_version="visual_changenet_levircd_deployable_v1.0" 47 | centerpose_version="deployable_bottle_fan_small_v1.0" 48 | foundationpose_version="deployable_v1.0" 49 | 50 | # Setting model version to run inference on for Pose Classification. 51 | pc_peoplenet_version="nvidia/tao/peoplenet:deployable_quantized_v2.5" 52 | pc_bodypose3dnet_version="nvidia/tao/bodypose3dnet:deployable_accuracy_v1.0" 53 | 54 | # NGC URL's to download the model. 55 | ngc_peoplenet="https://api.ngc.nvidia.com/v2/models/nvidia/tao/peoplenet/versions/${peoplenet_version}/zip" 56 | ngc_dashcamnet="https://api.ngc.nvidia.com/v2/models/nvidia/tao/dashcamnet/versions/${dashcamnet_version}/zip" 57 | ngc_vehicletypenet="https://api.ngc.nvidia.com/v2/models/nvidia/tao/vehicletypenet/versions/${vehicletypenet_version}/zip" 58 | ngc_lprnet="https://api.ngc.nvidia.com/v2/models/nvidia/tao/lprnet/versions/deployable_v1.0/zip" 59 | ngc_yolov3="https://nvidia.box.com/shared/static/3a00fdf8e1s2k3nezoxmfyykydxiyxy7" 60 | ngc_peoplesegnet="https://api.ngc.nvidia.com/v2/models/nvidia/tao/peoplesegnet/versions/deployable_v2.0/zip" 61 | ngc_retinanet="https://nvidia.box.com/shared/static/3a00fdf8e1s2k3nezoxmfyykydxiyxy7" 62 | ngc_mcls_classification="https://docs.google.com/uc?export=download&id=1blJQDQSlLPU6zX3yRmXODRwkcss6B3a3" 63 | ngc_pose_classification="https://api.ngc.nvidia.com/v2/models/nvidia/tao/poseclassificationnet/versions/deployable_v1.0/zip" 64 | ngc_re_identification="https://drive.google.com/uc?export=download&id=1jicWzrPgEgvHLoxS57XLwk3o2xRbXeN_" 65 | ngc_visual_changenet="nvidia/tao/visual_changenet_segmentation_levircd:${visual_changenet_version}" 66 | ngc_centerpose="https://api.ngc.nvidia.com/v2/models/nvidia/tao/centerpose_ros/versions/${centerpose_version}/zip" 67 | ngc_foundationpose="nvidia/tao/foundationpose:${foundationpose_version}" 68 | 69 | default_model_download_path="${tao_triton_root}/tao_models" 70 | -------------------------------------------------------------------------------- /scripts/download_and_convert.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Generate a peoplenet model. 4 | echo "Converting the PeopleNet model" 5 | mkdir -p /model_repository/peoplenet_tao/1 6 | trtexec --onnx=/tao_models/peoplenet_model/resnet34_peoplenet_int8.onnx \ 7 | --maxShapes="input_1:0":16x3x544x960 \ 8 | --minShapes="input_1:0":1x3x544x960 \ 9 | --optShapes="input_1:0":8x3x544x960 \ 10 | --calib=/tao_models/peoplenet_model/resnet34_peoplenet_int8.txt \ 11 | --int8 \ 12 | --saveEngine=/model_repository/peoplenet_tao/1/model.plan 13 | 14 | # Generate a dashcamnet model. 15 | echo "Converting the DashcamNet model" 16 | mkdir -p /model_repository/dashcamnet_tao/1 17 | trtexec --onnx=/tao_models/dashcamnet_model/resnet18_dashcamnet_pruned.onnx \ 18 | --maxShapes="input_1:0":16x3x544x960 \ 19 | --minShapes="input_1:0":1x3x544x960 \ 20 | --optShapes="input_1:0":8x3x544x960 \ 21 | --calib=/tao_models/dashcamnet_model/resnet18_dashcamnet_pruned_int8.txt \ 22 | --int8 \ 23 | --saveEngine=/model_repository/dashcamnet_tao/1/model.plan 24 | 25 | # Generate a vehicletypnet model. 26 | echo "Converting the VehicleTypeNet model" 27 | mkdir -p /model_repository/vehicletypenet_tao/1 28 | tao-converter /tao_models/vehicletypenet_model/resnet18_vehicletypenet_pruned.etlt \ 29 | -k tlt_encode \ 30 | -c /tao_models/vehicletypenet_model/vehicletypenet_int8.txt \ 31 | -d 3,224,224 \ 32 | -o predictions/Softmax \ 33 | -t int8 \ 34 | -m 16 \ 35 | -e /model_repository/vehicletypenet_tao/1/model.plan 36 | 37 | # Generate an LPRnet model. 38 | echo "Converting the LPRNet model" 39 | mkdir -p /model_repository/lprnet_tao/1 40 | tao-converter /tao_models/lprnet_model/us_lprnet_baseline18_deployable.etlt \ 41 | -k nvidia_tlt \ 42 | -p image_input,1x3x48x96,4x3x48x96,16x3x48x96 \ 43 | -t fp16 \ 44 | -e /model_repository/lprnet_tao/1/model.plan 45 | 46 | # Generate a YOLOv3 model. 47 | echo "Converting the YOLOv3 model" 48 | mkdir -p /model_repository/yolov3_tao/1 49 | tao-converter /tao_models/yolov3_model/yolov3_resnet18.etlt \ 50 | -k nvidia_tlt \ 51 | -p Input,1x3x544x960,4x3x544x960,16x3x544x960 \ 52 | -o BatchedNMS \ 53 | -t fp16 \ 54 | -e /model_repository/yolov3_tao/1/model.plan 55 | 56 | # Generate a peoplesegnet model. 57 | echo "Converting the peoplesegnet model" 58 | mkdir -p /model_repository/peoplesegnet_tao/1 59 | tao-converter /tao_models/peoplesegnet_model/peoplesegnet_resnet50.etlt \ 60 | -k nvidia_tlt \ 61 | -d 3,576,960 \ 62 | -p Input,1x3x576x960,4x3x576x960,16x3x576x960 \ 63 | -o generate_detections,mask_fcn_logits/BiasAdd \ 64 | -t fp16 \ 65 | -e /model_repository/peoplesegnet_tao/1/model.plan 66 | 67 | # Generate a retinanet model. 68 | echo "Converting the Retinanet model" 69 | mkdir -p /model_repository/retinanet_tao/1 70 | tao-converter /tao_models/retinanet_model/retinanet_resnet18_epoch_080_its_trt8.etlt \ 71 | -k nvidia_tlt \ 72 | -d 3,544,960 \ 73 | -o NMS \ 74 | -t fp16 \ 75 | -e /model_repository/retinanet_tao/1/model.plan 76 | 77 | # Generate a multitask_classification model. 78 | echo "Converting the multitask_classification model" 79 | mkdir -p /model_repository/multitask_classification_tao/1 80 | tao-converter /tao_models/multitask_cls_model/multitask_cls_resnet18.etlt \ 81 | -k nvidia_tlt \ 82 | -d 3,80,60 \ 83 | -o base_color/Softmax,category/Softmax,season/Softmax \ 84 | -t fp16 \ 85 | -m 16 \ 86 | -e /model_repository/multitask_classification_tao/1/model.plan 87 | 88 | # Generate a pose_classification model. 89 | echo "Converting the pose_classification model" 90 | mkdir -p /model_repository/pose_classification_tao/1 91 | tao-converter /tao_models/pose_cls_model/st-gcn_3dbp_nvidia.etlt \ 92 | -k nvidia_tao \ 93 | -d 3,300,34,1 \ 94 | -p input,1x3x300x34x1,4x3x300x34x1,16x3x300x34x1 \ 95 | -o fc_pred \ 96 | -t fp16 \ 97 | -m 16 \ 98 | -e /model_repository/pose_classification_tao/1/model.plan 99 | 100 | # Generate a re_identification model. 101 | echo "Converting the re_identification model" 102 | mkdir -p /model_repository/re_identification_tao/1 103 | tao-converter /tao_models/re_id_model/resnet50_market1501.etlt \ 104 | -k nvidia_tao \ 105 | -d 3,256,128 \ 106 | -p input,1x3x256x128,4x3x256x128,16x3x256x128 \ 107 | -o fc_pred \ 108 | -t fp16 \ 109 | -m 16 \ 110 | -e /model_repository/re_identification_tao/1/model.plan 111 | 112 | echo "Converting the changenet segmentation model" 113 | mkdir -p /model_repository/visual_changenet_segmentation_tao/1 114 | trtexec --onnx=/tao_models/visual_changenet_segmentation_tao/changenet_segment.onnx \ 115 | --saveEngine=/model_repository/visual_changenet_segmentation_tao/1/model.plan 116 | 117 | # Generate a CenterPose model 118 | echo "Converting the CenterPose model" 119 | mkdir -p /model_repository/centerpose_tao/1 120 | trtexec --onnx=/tao_models/centerpose_model/bottle_FAN_small.onnx \ 121 | --maxShapes="input":16x3x512x512 \ 122 | --minShapes="input":1x3x512x512 \ 123 | --optShapes="input":8x3x512x512 \ 124 | --saveEngine=/model_repository/centerpose_tao/1/model.plan 125 | 126 | # Generate a FoundationPose model, which includes two sub-models 127 | echo "Converting the FoundationPose refine model" 128 | mkdir -p /model_repository/foundationpose_refiner_tao/1 129 | trtexec --onnx=/tao_models/foundation_pose_vdeployable_v1.0/refiner_net.onnx \ 130 | --minShapes=inputA:1x6x160x160,inputB:1x6x160x160 \ 131 | --optShapes=inputA:252x6x160x160,inputB:252x6x160x160 \ 132 | --maxShapes=inputA:252x6x160x160,inputB:252x6x160x160 \ 133 | --saveEngine=/model_repository/foundationpose_refiner_tao/1/model.plan \ 134 | --preview=+fasterDynamicShapes0805 135 | 136 | # echo "Converting the FoundationPose score model" 137 | mkdir -p /model_repository/foundationpose_scorer_tao/1 138 | trtexec --onnx=/tao_models/foundation_pose_vdeployable_v1.0/score_net.onnx \ 139 | --minShapes=inputA:1x6x160x160,inputB:1x6x160x160 \ 140 | --optShapes=inputA:252x6x160x160,inputB:252x6x160x160 \ 141 | --maxShapes=inputA:252x6x160x160,inputB:252x6x160x160 \ 142 | --saveEngine=/model_repository/foundationpose_scorer_tao/1/model.plan 143 | --preview=+fasterDynamicShapes0805 144 | 145 | /opt/tritonserver/bin/tritonserver --model-store /model_repository 146 | -------------------------------------------------------------------------------- /scripts/pose_cls_e2e_inference/demo.mp4: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps/9a30f9692bf29fb728520e9dba1c79be2bf65e74/scripts/pose_cls_e2e_inference/demo.mp4 -------------------------------------------------------------------------------- /scripts/pose_cls_e2e_inference/download_and_convert.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Generate a pose_classification model. 4 | echo "Converting the pose_classification model" 5 | mkdir -p /model_repository/pose_classification_tao/1 6 | tao-converter /tao_models/pose_cls_model/st-gcn_3dbp_nvidia.etlt \ 7 | -k nvidia_tao \ 8 | -d 3,300,34,1 \ 9 | -p input,1x3x300x34x1,4x3x300x34x1,16x3x300x34x1 \ 10 | -o fc_pred \ 11 | -t fp16 \ 12 | -m 16 \ 13 | -e /model_repository/pose_classification_tao/1/model.plan 14 | 15 | /opt/tritonserver/bin/tritonserver --model-store /model_repository 16 | -------------------------------------------------------------------------------- /scripts/pose_cls_e2e_inference/plot_e2e_inference.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | import random 4 | import json 5 | import cv2 6 | 7 | 8 | num_joints = 34 9 | joint_links = { 10 | (255, 0, 0): [(0, 1), (1, 4), (4, 7), (7, 9), (9, 11), (7, 13), (1, 20), (6, 20), (20, 22), (22, 24), (24, 26), (24, 28), (24, 30), (24, 32)], 11 | (0, 255, 0): [(0, 3), (3, 6), (6, 15), (15, 16), (15, 17), (16, 18), (17, 19)], 12 | (0, 0, 255): [(0, 2), (2, 5), (5, 8), (8, 10), (10, 12), (8, 14), (2, 21), (6, 21), (21, 23), (23, 25), (25, 27), (25, 29), (25, 31), (25, 33)] 13 | } 14 | 15 | 16 | def main(): 17 | if sys.argv[3]: 18 | json_metadata_path = sys.argv[1] 19 | input_video_path = sys.argv[2] 20 | output_video_path = sys.argv[3] 21 | with open(json_metadata_path, 'r') as f: 22 | person_metadata = json.load(f) 23 | 24 | objects = {} 25 | 26 | for batch in person_metadata: 27 | for frame in batch["batches"]: 28 | frame_num = frame["frame_num"] 29 | objects[frame_num] = [] 30 | for person in frame["objects"]: 31 | objects[frame_num].append(person) 32 | 33 | video_capture = cv2.VideoCapture(input_video_path) 34 | frame_width = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH)) 35 | frame_height = int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) 36 | fps = video_capture.get(cv2.CAP_PROP_FPS) 37 | video_writer = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc('M', 'P', '4', 'V'), 38 | fps, (frame_width, frame_height)) 39 | success, image_frame = video_capture.read() 40 | frame_num = 0 41 | bbox_colors = {} 42 | while success: 43 | if frame_num in objects.keys(): 44 | print("Plotting results for frame %06d" % frame_num) 45 | for person in objects[frame_num]: 46 | # Use a random BGR color to represent the object ID 47 | object_id = person["object_id"] 48 | bbox_color = None 49 | if object_id in bbox_colors.keys(): 50 | bbox_color = bbox_colors[object_id] 51 | else: 52 | bbox_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) 53 | bbox_colors[object_id] = bbox_color 54 | 55 | # Plot 2D body pose 56 | joints = [] 57 | for j in range(num_joints): 58 | x = int(person["pose25d"][j*4+0]) 59 | y = int(person["pose25d"][j*4+1]) 60 | if person["pose25d"][j*4+3] == 0: 61 | x = -1 62 | y = -1 63 | joints.append((x, y)) 64 | for link_color in joint_links.keys(): 65 | for joint_pair in joint_links[link_color]: 66 | if joints[joint_pair[0]] == (-1, -1) or joints[joint_pair[1]] == (-1, -1): 67 | continue 68 | image_frame = cv2.line(image_frame, joints[joint_pair[0]], joints[joint_pair[1]], link_color, 2) 69 | for j in range(num_joints): 70 | image_frame = cv2.circle(image_frame, joints[j], 2, (255, 255, 255), 2) 71 | 72 | # Plot bounding box 73 | bbox = person["bbox"] 74 | bbox_top_left = [int(bbox[0]), int(bbox[1])] 75 | bbox_bottom_right = [int(bbox[0] + bbox[2] - 1), int(bbox[1] + bbox[3] - 1)] 76 | if bbox_top_left[0] < 0: 77 | bbox_top_left[0] = 0 78 | if bbox_top_left[1] < 0: 79 | bbox_top_left[1] = 0 80 | if bbox_bottom_right[0] > frame_width - 1: 81 | bbox_bottom_right[0] = frame_width - 1 82 | if bbox_bottom_right[1] > frame_height - 1: 83 | bbox_bottom_right[1] = frame_height - 1 84 | image_frame = cv2.rectangle(image_frame, bbox_top_left, bbox_bottom_right, bbox_color, 2) 85 | 86 | # Plot object ID and action 87 | image_frame = cv2.putText(image_frame, f"{str(object_id)}: {person['action']}", 88 | (bbox_top_left[0], bbox_top_left[1] - 10), 89 | cv2.FONT_HERSHEY_SIMPLEX, 0.8, bbox_color, 2, cv2.LINE_AA) 90 | 91 | # Plot frame number 92 | image_frame = cv2.putText(image_frame, "%06d" % frame_num, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 93 | 0.8, (0, 0, 255), 2, cv2.LINE_AA) 94 | 95 | video_writer.write(image_frame) 96 | success, image_frame = video_capture.read() 97 | frame_num += 1 98 | 99 | video_capture.release() 100 | video_writer.release() 101 | 102 | print("Output video saved at %s" % output_video_path) 103 | 104 | else: 105 | print("Usage: %s json_metadata_path input_video_path output_video_path" % __file__) 106 | 107 | 108 | if __name__ == '__main__': 109 | main() 110 | -------------------------------------------------------------------------------- /scripts/pose_cls_e2e_inference/start_client.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | function check_ngc_cli_installation { 4 | if ! command -v ngc > /dev/null; then 5 | echo "[ERROR] The NGC CLI tool not found on device in /usr/bin/ or PATH env var" 6 | echo "[ERROR] Please follow: https://ngc.nvidia.com/setup/installers/cli" 7 | exit 8 | fi 9 | } 10 | 11 | get_ngc_key_from_environment() { 12 | # first check the global NGC_API_KEY environment variable. 13 | local ngc_key=$NGC_API_KEY 14 | # if env variable was not set, and a ~/.ngc/config exists 15 | # try to get it from there. 16 | if [ -z "$ngc_key" ] && [[ -f "$HOME/.ngc/config" ]] 17 | then 18 | ngc_key=$(cat $HOME/.ngc/config | grep apikey -m1 | awk '{print $3}') 19 | fi 20 | echo $ngc_key 21 | } 22 | 23 | check_ngc_cli_installation 24 | NGC_API_KEY="$(get_ngc_key_from_environment)" 25 | if [ -z "$NGC_API_KEY" ]; then 26 | echo -e 'Did not find environment variable "$NGC_API_KEY"' 27 | read -sp 'Please enter API key for ngc.nvidia.com: ' NGC_API_KEY 28 | echo 29 | fi 30 | 31 | set -e 32 | 33 | # Docker login to Nvidia GPU Cloud (NGC). 34 | docker login nvcr.io -u \$oauthtoken -p ${NGC_API_KEY} 35 | 36 | # Load config file 37 | script_path="$( cd "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )" 38 | if [ -z "$1" ]; then 39 | config_path="${script_path}/../config.sh" 40 | else 41 | config_path=$(readlink -f $1) 42 | fi 43 | 44 | # Configure the required environment variables. 45 | if [[ ! -f $config_path ]]; then 46 | echo 'Unable to load configuration file. Override path to file with -c argument.' 47 | exit 1 48 | fi 49 | source $config_path 50 | 51 | # Clone DeepStream repo 52 | # git clone https://github.com/NVIDIA-AI-IOT/deepstream_reference_apps.git ${tao_triton_root}/deepstream_reference_apps 53 | git clone https://github.com/xunleiw/deepstream_reference_apps.git ${tao_triton_root}/deepstream_reference_apps 54 | export BODYPOSE3D_HOME=${tao_triton_root}/deepstream_reference_apps/deepstream-bodypose-3d 55 | 56 | # Download models using NGC 57 | mkdir -p $BODYPOSE3D_HOME/models 58 | cd $BODYPOSE3D_HOME/models 59 | check_ngc_cli_installation 60 | ngc registry model download-version ${pc_peoplenet_version} 61 | ngc registry model download-version ${pc_bodypose3dnet_version} 62 | apt-get install -y tree 63 | tree $BODYPOSE3D_HOME -d 64 | 65 | # Install Eigen 66 | cd $BODYPOSE3D_HOME 67 | wget https://gitlab.com/libeigen/eigen/-/archive/3.4.0/eigen-3.4.0.tar.gz 68 | tar xvzf eigen-3.4.0.tar.gz 69 | ln eigen-3.4.0 eigen -s 70 | rm eigen-3.4.0.tar.gz 71 | 72 | # Update event message payload of DeepStream 73 | cp $BODYPOSE3D_HOME/sources/deepstream-sdk/eventmsg_payload.cpp /opt/nvidia/deepstream/deepstream/sources/libs/nvmsgconv/deepstream_schema 74 | apt-get install -y libjson-glib-dev uuid-dev 75 | cd /opt/nvidia/deepstream/deepstream/sources/libs/nvmsgconv 76 | make; make install 77 | 78 | # Make 3D body pose sources 79 | export CUDA_VER=${cuda_ver} 80 | cd $BODYPOSE3D_HOME/sources/nvdsinfer_custom_impl_BodyPose3DNet 81 | make 82 | cd $BODYPOSE3D_HOME/sources 83 | make 84 | 85 | # Run 3D body pose 86 | ./deepstream-pose-estimation-app --input file://${tao_triton_root}/scripts/pose_cls_e2e_inference/demo.mp4 \ 87 | --output ${tao_triton_root}/scripts/pose_cls_e2e_inference/demo_3dbp.mp4 \ 88 | --focal 1200.0 \ 89 | --width 1920 \ 90 | --height 1080 \ 91 | --fps \ 92 | --save-pose ${tao_triton_root}/scripts/pose_cls_e2e_inference/demo_3dbp.json 93 | 94 | # Run the Triton client 95 | cd ${tao_triton_root} 96 | python3 -m tao_triton.python.entrypoints.tao_client ${tao_triton_root}/scripts/pose_cls_e2e_inference/demo_3dbp.json \ 97 | --dataset_convert_config ${tao_triton_root}/tao_triton/python/dataset_convert_specs/dataset_convert_config_pose_classification.yaml \ 98 | -m pose_classification_tao \ 99 | -x 1 \ 100 | -b 1 \ 101 | --mode Pose_classification \ 102 | -i https \ 103 | -u localhost:8000 \ 104 | --async \ 105 | --output_path ${tao_triton_root}/scripts/pose_cls_e2e_inference 106 | 107 | # Plot inference results 108 | python3 ./scripts/pose_cls_e2e_inference/plot_e2e_inference.py \ 109 | ./scripts/pose_cls_e2e_inference/results.json \ 110 | ./scripts/pose_cls_e2e_inference/demo.mp4 \ 111 | ./scripts/pose_cls_e2e_inference/results.mp4 112 | 113 | # Clean repo 114 | rm -r ${tao_triton_root}/deepstream_reference_apps 115 | -------------------------------------------------------------------------------- /scripts/pose_cls_e2e_inference/start_server.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | function check_wget_installed { 4 | if ! command -v wget > /dev/null; then 5 | echo "Wget not found. Please run sudo apt-get install wget" 6 | return false 7 | fi 8 | return 0 9 | } 10 | 11 | function check_ngc_cli_installation { 12 | if ! command -v ngc > /dev/null; then 13 | echo "[ERROR] The NGC CLI tool not found on device in /usr/bin/ or PATH env var" 14 | echo "[ERROR] Please follow: https://ngc.nvidia.com/setup/installers/cli" 15 | exit 16 | fi 17 | } 18 | 19 | get_ngc_key_from_environment() { 20 | # first check the global NGC_API_KEY environment variable. 21 | local ngc_key=$NGC_API_KEY 22 | # if env variable was not set, and a ~/.ngc/config exists 23 | # try to get it from there. 24 | if [ -z "$ngc_key" ] && [[ -f "$HOME/.ngc/config" ]] 25 | then 26 | ngc_key=$(cat $HOME/.ngc/config | grep apikey -m1 | awk '{print $3}') 27 | fi 28 | echo $ngc_key 29 | } 30 | 31 | check_ngc_cli_installation 32 | NGC_API_KEY="$(get_ngc_key_from_environment)" 33 | if [ -z "$NGC_API_KEY" ]; then 34 | echo -e 'Did not find environment variable "$NGC_API_KEY"' 35 | read -sp 'Please enter API key for ngc.nvidia.com: ' NGC_API_KEY 36 | echo 37 | fi 38 | 39 | set -e 40 | 41 | # Docker login to Nvidia GPU Cloud (NGC). 42 | docker login nvcr.io -u \$oauthtoken -p ${NGC_API_KEY} 43 | 44 | # Load config file 45 | script_path="$( cd "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )" 46 | if [ -z "$1" ]; then 47 | config_path="${script_path}/../config.sh" 48 | else 49 | config_path=$(readlink -f $1) 50 | fi 51 | 52 | # Configure the required environment variables. 53 | if [[ ! -f $config_path ]]; then 54 | echo 'Unable to load configuration file. Override path to file with -c argument.' 55 | exit 1 56 | fi 57 | source $config_path 58 | 59 | # Building the Triton docker with the tao-converter. 60 | docker build -f "${tao_triton_root}/docker/Dockerfile" \ 61 | -t ${tao_triton_server_docker}:${tao_triton_server_tag} ${tao_triton_root} 62 | 63 | mkdir -p ${default_model_download_path} && cd ${default_model_download_path} 64 | wget --content-disposition ${ngc_pose_classification} -O ${default_model_download_path}/poseclassificationnet_v1.0.zip && \ 65 | unzip ${default_model_download_path}/poseclassificationnet_v1.0.zip -d ${default_model_download_path}/pose_cls_model/ 66 | 67 | # Run the server container. 68 | echo "Running the server on ${gpu_id}" 69 | find ${tao_triton_root}/model_repository -mindepth 1 ! -regex "^${tao_triton_root}/model_repository/pose_classification_tao\(/.*\)?" -delete 70 | docker run -it --rm -v ${tao_triton_root}/model_repository:/model_repository \ 71 | -v ${default_model_download_path}:/tao_models \ 72 | -v ${tao_triton_root}/scripts:/tao_triton \ 73 | --gpus all \ 74 | -p 8000:8000 \ 75 | -p 8001:8001 \ 76 | -p 8002:8002 \ 77 | -e CUDA_VISIBLE_DEVICES=$gpu_id \ 78 | ${tao_triton_server_docker}:${tao_triton_server_tag} \ 79 | /tao_triton/pose_cls_e2e_inference/download_and_convert.sh 80 | -------------------------------------------------------------------------------- /scripts/re_id_e2e_inference/download_and_convert.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Generate a re_identification model. 4 | echo "Converting the re_identification model" 5 | mkdir -p /model_repository/re_identification_tao/1 6 | tao-converter /tao_models/re_id_model/resnet50_market1501.etlt \ 7 | -k nvidia_tao \ 8 | -d 3,256,128 \ 9 | -p input,1x3x256x128,4x3x256x128,16x3x256x128 \ 10 | -o fc_pred \ 11 | -t fp16 \ 12 | -m 16 \ 13 | -e /model_repository/re_identification_tao/1/model.plan 14 | /opt/tritonserver/bin/tritonserver --model-store /model_repository -------------------------------------------------------------------------------- /scripts/re_id_e2e_inference/plot_e2e_inference.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import re 3 | import sys 4 | import os 5 | import json 6 | import numpy as np 7 | import matplotlib.pyplot as plt 8 | from re_ranking import R1_mAP_reranking 9 | 10 | 11 | def main(): 12 | if sys.argv[2]: 13 | json_metadata_path = sys.argv[1] 14 | output_dir = sys.argv[2] 15 | f = open(json_metadata_path) 16 | pattern = re.compile(r'([-\d]+)_c(\d)') 17 | data = json.load(f) 18 | 19 | pids = [] 20 | camids = [] 21 | img_paths = [] 22 | embeddings = [] 23 | num_query = 0 24 | 25 | for row in data: 26 | img_path = row["img_path"] 27 | if "query" in img_path: 28 | num_query += 1 29 | embedding = row["embedding"] 30 | pid, camid = map(int, pattern.search(img_path).groups()) 31 | if pid == -1: continue # junk images are ignored 32 | camid -= 1 # index starts from 0 33 | embeddings.append(embedding) 34 | pids.append(pid) 35 | camids.append(camid) 36 | img_paths.append(img_path) 37 | metrics = R1_mAP_reranking(num_query, output_dir, feat_norm=True) 38 | metrics.reset() 39 | metrics.update(torch.tensor(embeddings), pids, camids, img_paths) 40 | cmc, _ = metrics.compute() 41 | f.close() 42 | 43 | plt.figure() 44 | cmc_percentages = [value * 100 for value in cmc] 45 | plt.xticks(np.arange(len(cmc_percentages)), np.arange(1, len(cmc_percentages)+1)) 46 | plt.plot(cmc_percentages, marker="*") 47 | plt.title('Cumulative Matching Characteristics (CMC) curve') 48 | plt.grid() 49 | plt.ylabel('Matching Rate[%]') 50 | plt.xlabel('Rank') 51 | output_cmc_curve_plot_path = os.path.join(output_dir, 'cmc_curve.png') 52 | plt.savefig(output_cmc_curve_plot_path) 53 | 54 | print("Output CMC curve plot saved at %s" % output_cmc_curve_plot_path) 55 | 56 | else: 57 | print("Usage: %s json_metadata_path output_dir" % __file__) 58 | 59 | 60 | if __name__ == '__main__': 61 | main() 62 | -------------------------------------------------------------------------------- /scripts/re_id_e2e_inference/re_ranking.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 2 | 3 | # Original source taken from https://github.com/michuanhaohao/reid-strong-baseline 4 | 5 | # Permission is hereby granted, free of charge, to any person obtaining a copy 6 | # of this software and associated documentation files (the "Software"), to deal 7 | # in the Software without restriction, including without limitation the rights 8 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | # copies of the Software, and to permit persons to whom the Software is 10 | # furnished to do so, subject to the following conditions: 11 | 12 | # The above copyright notice and this permission notice shall be included in all 13 | # copies or substantial portions of the Software. 14 | 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | # SOFTWARE. 22 | """Re-Identification Metrics.""" 23 | 24 | import numpy as np 25 | import torch 26 | 27 | from utils import re_ranking, eval_func 28 | 29 | 30 | def euclidean_distance(qf, gf): 31 | """Return a similiarity matrix based on euclidian distance. 32 | 33 | Args: 34 | qf (Numpy): Matrix A. 35 | gf (Numpy): Matrix B 36 | 37 | Returns: 38 | dist_mat (Numpy): Distance Matrix. 39 | 40 | """ 41 | m = qf.shape[0] 42 | n = gf.shape[0] 43 | dist_mat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + \ 44 | torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t() 45 | dist_mat.addmm_(qf, gf.t(), beta = 1, alpha = -2) 46 | return dist_mat.cpu().numpy() 47 | 48 | 49 | def cosine_similarity(qf, gf): 50 | """Return a similiarity matrix. 51 | 52 | Args: 53 | qf (Numpy): Matrix A. 54 | gf (Numpy): Matrix B 55 | 56 | Returns: 57 | dist_mat (Numpy): Distance Matrix. 58 | 59 | """ 60 | epsilon = 0.00001 61 | dist_mat = qf.mm(gf.t()) 62 | qf_norm = torch.norm(qf, p=2, dim=1, keepdim=True) # mx1 63 | gf_norm = torch.norm(gf, p=2, dim=1, keepdim=True) # nx1 64 | qg_normdot = qf_norm.mm(gf_norm.t()) 65 | 66 | dist_mat = dist_mat.mul(1 / qg_normdot).cpu().numpy() 67 | dist_mat = np.clip(dist_mat, -1 + epsilon, 1 - epsilon) 68 | dist_mat = np.arccos(dist_mat) 69 | return dist_mat 70 | 71 | 72 | class R1_mAP_reranking(): 73 | """R1_mAP Class for metrics with reranking.""" 74 | 75 | def __init__(self, num_query, output_dir, feat_norm=True): 76 | """Initialize the R1_mAP class.""" 77 | super(R1_mAP_reranking, self).__init__() 78 | self.num_query = num_query 79 | self.feat_norm = feat_norm 80 | self.feats = [] 81 | self.pids = [] 82 | self.camids = [] 83 | self.img_paths = [] 84 | self.output_dir = output_dir 85 | 86 | def reset(self): 87 | """Reset the data members.""" 88 | self.feats = [] 89 | self.pids = [] 90 | self.camids = [] 91 | self.img_paths = [] 92 | 93 | def update(self, feat, pid, camid, img_path): 94 | """Append to the data members. 95 | 96 | Args: 97 | feat (Tensor): Feature embedding. 98 | pid (int): Person IDs. 99 | camid (int): Camera IDs. 100 | img_path (str): Image Paths. 101 | 102 | """ 103 | self.feats.append(feat) 104 | self.pids.extend(np.asarray(pid)) 105 | self.camids.extend(np.asarray(camid)) 106 | self.img_paths.extend(img_path) 107 | 108 | def compute(self): 109 | """Compute the metrics. 110 | 111 | Returns: 112 | cmc (list): CMC Rank List. 113 | mAP (float): Mean average precision. 114 | 115 | """ 116 | feats = torch.cat(self.feats, dim=0) 117 | if self.feat_norm: 118 | print("The test features are normalized.") 119 | feats = torch.nn.functional.normalize(feats, dim=1, p=2) 120 | 121 | # query 122 | qf = feats[:self.num_query] 123 | q_pids = np.asarray(self.pids[:self.num_query]) 124 | q_camids = np.asarray(self.camids[:self.num_query]) 125 | q_img_paths = self.img_paths[:self.num_query] 126 | # gallery 127 | gf = feats[self.num_query:] 128 | g_pids = np.asarray(self.pids[self.num_query:]) 129 | g_camids = np.asarray(self.camids[self.num_query:]) 130 | g_img_paths = self.img_paths[self.num_query:] 131 | 132 | print("The distance matrix is processed by re-ranking.") 133 | distmat = re_ranking(qf, gf, k1=20, k2=6, lambda_value=0.3) 134 | cmc, mAP = eval_func(distmat, q_pids, g_pids, q_camids, g_camids, q_img_paths, g_img_paths, self.output_dir) 135 | return cmc, mAP 136 | -------------------------------------------------------------------------------- /scripts/re_id_e2e_inference/sample_dataset.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | import glob 4 | import re 5 | import random 6 | 7 | 8 | def sample_dataset(input_dir, output_dir, n_samples, use_ids = None): 9 | """Select a subset of images fom input_dir and move them to output_dir. 10 | 11 | Args: 12 | input_dir (str): Input Folder Path of the train images. 13 | output_dir (str): Output Folder Path of the test images. 14 | n_samples (int): Number of samples to use. 15 | use_ids(list int): List of IDs to grab from test and query folder. 16 | 17 | Returns: 18 | IDs used for sampling 19 | """ 20 | img_paths = glob.glob(os.path.join(input_dir, '*.jpg')) 21 | pattern = re.compile(r'([-\d]+)_c(\d)') 22 | id_to_img = {} 23 | 24 | # Grab images with matching ids 25 | for img_path in img_paths: 26 | pid, _ = map(int, pattern.search(img_path).groups()) 27 | if pid not in id_to_img: 28 | id_to_img[pid] = [] 29 | id_to_img[pid].append(img_path) 30 | 31 | # Create directory 32 | if not os.path.exists(output_dir): 33 | os.makedirs(output_dir) 34 | else: 35 | command = "rm -r " + output_dir 36 | os.system(command) 37 | os.makedirs(output_dir) 38 | 39 | assert id_to_img, "Dataset size cannot be 0." 40 | 41 | sampled_id_to_img = dict(random.sample(id_to_img.items(), n_samples)) 42 | 43 | for key, img_paths in sampled_id_to_img.items(): 44 | for img_path in img_paths: 45 | command = "cp " + img_path + " " + output_dir 46 | os.system(command) 47 | 48 | # Use same ids for test and query 49 | if use_ids: 50 | 51 | # Create query dir 52 | if not os.path.exists(output_dir): 53 | os.makedirs(output_dir) 54 | else: 55 | command = "rm -r " + output_dir 56 | os.system(command) 57 | os.makedirs(output_dir) 58 | 59 | # Find images in test with same id 60 | img_paths = glob.glob(os.path.join(input_dir, '*.jpg')) 61 | for id in use_ids: 62 | pattern = re.compile(r'([-\d]+)_c(\d)') 63 | for img_path in img_paths: 64 | pid, _ = map(int, pattern.search(img_path).groups()) 65 | if id == pid: 66 | print(img_path) 67 | command = "cp " + img_path + " " + output_dir 68 | os.system(command) 69 | 70 | return sampled_id_to_img.keys() 71 | 72 | 73 | def main(): 74 | if sys.argv[1]: 75 | data_dir = sys.argv[1] 76 | 77 | # Number of samples 78 | n_samples = 100 79 | 80 | # Create train dataset 81 | train_input_dir = os.path.join(data_dir, "bounding_box_train") 82 | train_output_dir = os.path.join(data_dir, "sample_train") 83 | sample_dataset(train_input_dir, train_output_dir, n_samples) 84 | 85 | # Create test dataset 86 | test_input_dir = os.path.join(data_dir, "bounding_box_test") 87 | test_output_dir = os.path.join(data_dir, "sample_test") 88 | ids = sample_dataset(test_input_dir, test_output_dir, n_samples) 89 | 90 | # Create query dataset 91 | query_input_dir = os.path.join(data_dir, "query") 92 | query_output_dir = os.path.join(data_dir, "sample_query") 93 | sample_dataset(query_input_dir, query_output_dir, n_samples, ids) 94 | 95 | else: 96 | print("Usage: %s data_dir" % __file__) 97 | 98 | 99 | if __name__ == '__main__': 100 | main() 101 | -------------------------------------------------------------------------------- /scripts/re_id_e2e_inference/start_client.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Load config file 4 | script_path="$( cd "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )" 5 | if [ -z "$1" ]; then 6 | config_path="${script_path}/../config.sh" 7 | else 8 | config_path=$(readlink -f $1) 9 | fi 10 | 11 | # Configure the required environment variables. 12 | if [[ ! -f $config_path ]]; then 13 | echo 'Unable to load configuration file. Override path to file with -c argument.' 14 | exit 1 15 | fi 16 | source $config_path 17 | 18 | # Download the dataset 19 | rm -rf ${tao_triton_root}/data 20 | mkdir ${tao_triton_root}/data 21 | pip3 install gdown 22 | gdown https://drive.google.com/uc?id=0B8-rUzbwVRk0c054eEozWG9COHM -O ${tao_triton_root}/data/market1501.zip 23 | 24 | # Extract the files 25 | unzip ${tao_triton_root}/data/market1501.zip -d ${tao_triton_root}/data 26 | mv ${tao_triton_root}/data/Market-1501-v15.09.15 ${tao_triton_root}/data/market1501 27 | rm ${tao_triton_root}/data/market1501.zip 28 | 29 | # Verify 30 | ls -l ${tao_triton_root}/data/market1501 31 | 32 | # Sample the dataset 33 | cd ${tao_triton_root} 34 | python3 ./scripts/re_id_e2e_inference/sample_dataset.py \ 35 | ${tao_triton_root}/data/market1501 36 | 37 | # Run the Triton client 38 | rm -f ./scripts/re_id_e2e_inference/results.json 39 | python3 -m tao_triton.python.entrypoints.tao_client ${tao_triton_root}/data/market1501/sample_query \ 40 | --test_dir ${tao_triton_root}/data/market1501/sample_test \ 41 | -m re_identification_tao \ 42 | -x 1 \ 43 | -b 16 \ 44 | --mode Re_identification \ 45 | -i https \ 46 | -u localhost:8000 \ 47 | --async \ 48 | --output_path ${tao_triton_root}/scripts/re_id_e2e_inference 49 | 50 | # Plot inference results 51 | python3 ./scripts/re_id_e2e_inference/plot_e2e_inference.py \ 52 | ./scripts/re_id_e2e_inference/results.json \ 53 | ./scripts/re_id_e2e_inference 54 | -------------------------------------------------------------------------------- /scripts/re_id_e2e_inference/start_server.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | function check_wget_installed { 4 | if ! command -v wget > /dev/null; then 5 | echo "Wget not found. Please run sudo apt-get install wget" 6 | return false 7 | fi 8 | return 0 9 | } 10 | 11 | function check_ngc_cli_installation { 12 | if ! command -v ngc > /dev/null; then 13 | echo "[ERROR] The NGC CLI tool not found on device in /usr/bin/ or PATH env var" 14 | echo "[ERROR] Please follow: https://ngc.nvidia.com/setup/installers/cli" 15 | exit 16 | fi 17 | } 18 | 19 | get_ngc_key_from_environment() { 20 | # first check the global NGC_API_KEY environment variable. 21 | local ngc_key=$NGC_API_KEY 22 | # if env variable was not set, and a ~/.ngc/config exists 23 | # try to get it from there. 24 | if [ -z "$ngc_key" ] && [[ -f "$HOME/.ngc/config" ]] 25 | then 26 | ngc_key=$(cat $HOME/.ngc/config | grep apikey -m1 | awk '{print $3}') 27 | fi 28 | echo $ngc_key 29 | } 30 | 31 | check_ngc_cli_installation 32 | NGC_API_KEY="$(get_ngc_key_from_environment)" 33 | if [ -z "$NGC_API_KEY" ]; then 34 | echo -e 'Did not find environment variable "$NGC_API_KEY"' 35 | read -sp 'Please enter API key for ngc.nvidia.com: ' NGC_API_KEY 36 | echo 37 | fi 38 | 39 | set -e 40 | 41 | # Docker login to Nvidia GPU Cloud (NGC). 42 | docker login nvcr.io -u \$oauthtoken -p ${NGC_API_KEY} 43 | 44 | # Load config file 45 | script_path="$( cd "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )" 46 | if [ -z "$1" ]; then 47 | config_path="${script_path}/../config.sh" 48 | else 49 | config_path=$(readlink -f $1) 50 | fi 51 | 52 | # Configure the required environment variables. 53 | if [[ ! -f $config_path ]]; then 54 | echo 'Unable to load configuration file. Override path to file with -c argument.' 55 | exit 1 56 | fi 57 | source $config_path 58 | 59 | # Building the Triton docker with the tao-converter. 60 | docker build -f "${tao_triton_root}/docker/Dockerfile" \ 61 | -t ${tao_triton_server_docker}:${tao_triton_server_tag} ${tao_triton_root} 62 | 63 | mkdir -p ${default_model_download_path} && cd ${default_model_download_path} 64 | rm -rf ${default_model_download_path}/re_id_model 65 | mkdir ${default_model_download_path}/re_id_model 66 | wget --no-check-certificate ${ngc_re_identification} -O ${default_model_download_path}/re_id_model/resnet50_market1501.etlt 67 | 68 | # Run the server container. 69 | echo "Running the server on ${gpu_id}" 70 | find ${tao_triton_root}/model_repository -mindepth 1 ! -regex "^${tao_triton_root}/model_repository/re_identification_tao\(/.*\)?" -delete 71 | docker run -it --rm -v ${tao_triton_root}/model_repository:/model_repository \ 72 | -v ${default_model_download_path}:/tao_models \ 73 | -v ${tao_triton_root}/scripts:/tao_triton \ 74 | --gpus all \ 75 | -p 8000:8000 \ 76 | -p 8001:8001 \ 77 | -p 8002:8002 \ 78 | -e CUDA_VISIBLE_DEVICES=$gpu_id \ 79 | ${tao_triton_server_docker}:${tao_triton_server_tag} \ 80 | /tao_triton/re_id_e2e_inference/download_and_convert.sh 81 | -------------------------------------------------------------------------------- /scripts/start_server.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | function check_wget_installed { 4 | if ! command -v wget > /dev/null; then 5 | echo "Wget not found. Please run sudo apt-get install wget" 6 | return false 7 | fi 8 | return 0 9 | } 10 | 11 | function check_ngc_cli_installation { 12 | if ! command -v ngc > /dev/null; then 13 | echo "[ERROR] The NGC CLI tool not found on device in /usr/bin/ or PATH env var" 14 | echo "[ERROR] Please follow: https://ngc.nvidia.com/setup/installers/cli" 15 | exit 16 | fi 17 | } 18 | 19 | get_ngc_key_from_environment() { 20 | # first check the global NGC_API_KEY environment variable. 21 | local ngc_key=$NGC_API_KEY 22 | # if env variable was not set, and a ~/.ngc/config exists 23 | # try to get it from there. 24 | if [ -z "$ngc_key" ] && [[ -f "$HOME/.ngc/config" ]] 25 | then 26 | ngc_key=$(cat $HOME/.ngc/config | grep apikey -m1 | awk '{print $3}') 27 | fi 28 | echo $ngc_key 29 | } 30 | 31 | check_ngc_cli_installation 32 | NGC_API_KEY="$(get_ngc_key_from_environment)" 33 | if [ -z "$NGC_API_KEY" ]; then 34 | echo -e 'Did not find environment variable "$NGC_API_KEY"' 35 | read -sp 'Please enter API key for ngc.nvidia.com: ' NGC_API_KEY 36 | echo 37 | fi 38 | 39 | set -e 40 | 41 | # Docker login to Nvidia GPU Cloud (NGC). 42 | docker login nvcr.io -u \$oauthtoken -p ${NGC_API_KEY} 43 | 44 | # load config file 45 | script_path="$( cd "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )" 46 | if [ -z "$1" ]; then 47 | config_path="${script_path}/config.sh" 48 | else 49 | config_path=$(readlink -f $1) 50 | fi 51 | 52 | # Configure the required environment variables. 53 | if [[ ! -f $config_path ]]; then 54 | echo 'Unable to load configuration file. Override path to file with -c argument.' 55 | exit 1 56 | fi 57 | source $config_path 58 | 59 | # Building the Triton docker with the tao-converter. 60 | docker build -f "${tao_triton_root}/docker/Dockerfile" \ 61 | -t ${tao_triton_server_docker}:${tao_triton_server_tag} ${tao_triton_root} 62 | 63 | mkdir -p ${default_model_download_path} && cd ${default_model_download_path} 64 | wget --content-disposition ${ngc_peoplenet} -O ${default_model_download_path}/peoplenet_${peoplenet_version}.zip && \ 65 | unzip ${default_model_download_path}/peoplenet_${peoplenet_version}.zip -d ${default_model_download_path}/peoplenet_model/ 66 | wget --content-disposition ${ngc_dashcamnet} -O ${default_model_download_path}/dashcamnet_${dashcamnet_version}.zip && \ 67 | unzip ${default_model_download_path}/dashcamnet_${dashcamnet_version}.zip -d ${default_model_download_path}/dashcamnet_model/ 68 | wget --content-disposition ${ngc_vehicletypenet} -O ${default_model_download_path}/vehicletypenet_${vehicletypenet_version}.zip && \ 69 | unzip ${default_model_download_path}/vehicletypenet_${vehicletypenet_version}.zip -d ${default_model_download_path}/vehicletypenet_model/ 70 | wget --content-disposition ${ngc_lprnet} -O ${default_model_download_path}/lprnet_pruned_v1.0.zip && \ 71 | unzip ${default_model_download_path}/lprnet_pruned_v1.0.zip -d ${default_model_download_path}/lprnet_model/ 72 | wget --content-disposition ${ngc_yolov3} -O ${default_model_download_path}/models.zip && \ 73 | unzip ${default_model_download_path}/models.zip -d ${default_model_download_path} && \ 74 | rm -rf ${default_model_download_path}/yolov3_model && \ 75 | mv ${default_model_download_path}/models/yolov3 ${default_model_download_path}/yolov3_model && \ 76 | rm -rf ${default_model_download_path}/retinanet_model && \ 77 | mv ${default_model_download_path}/models/retinanet ${default_model_download_path}/retinanet_model && \ 78 | rm -rf ${default_model_download_path}/models 79 | wget --content-disposition ${ngc_peoplesegnet} -O ${default_model_download_path}/peoplesegnet_deployable_v2.0.zip && \ 80 | unzip ${default_model_download_path}/peoplesegnet_deployable_v2.0.zip -d ${default_model_download_path}/peoplesegnet_model/ 81 | rm -rf ${default_model_download_path}/multitask_cls_model 82 | mkdir ${default_model_download_path}/multitask_cls_model 83 | wget --no-check-certificate ${ngc_mcls_classification} -O ${default_model_download_path}/multitask_cls_model/multitask_cls_resnet18.etlt 84 | wget --content-disposition ${ngc_pose_classification} -O ${default_model_download_path}/poseclassificationnet_v1.0.zip && \ 85 | unzip ${default_model_download_path}/poseclassificationnet_v1.0.zip -d ${default_model_download_path}/pose_cls_model/ 86 | rm -rf ${default_model_download_path}/re_id_model 87 | mkdir ${default_model_download_path}/re_id_model 88 | wget --no-check-certificate ${ngc_re_identification} -O ${default_model_download_path}/re_id_model/resnet50_market1501.etlt 89 | ngc registry model download-version $ngc_visual_changenet --dest ${default_model_download_path} 90 | mv ${default_model_download_path}/visual_changenet_segmentation_levircd_v${visual_changenet_version} ${default_model_download_path}/visual_changenet_segmentation_tao 91 | wget --content-disposition ${ngc_centerpose} -O ${default_model_download_path}/centerpose_ros_deployable_bottle_dla34_v1.0.zip && \ 92 | unzip ${default_model_download_path}/centerpose_ros_deployable_bottle_dla34_v1.0.zip -d ${default_model_download_path}/centerpose_model/ 93 | rm -rf ${default_model_download_path}/*.zip 94 | ngc registry model download-version ${ngc_foundationpose} --dest ${default_model_download_path} 95 | 96 | # Run the server container. 97 | echo "Running the server on ${gpu_id}" 98 | docker run -it --rm -v ${tao_triton_root}/model_repository:/model_repository \ 99 | -v ${default_model_download_path}:/tao_models \ 100 | -v ${tao_triton_root}/scripts:/tao_triton \ 101 | --gpus all \ 102 | -p 8000:8000 \ 103 | -p 8001:8001 \ 104 | -p 8002:8002 \ 105 | -e CUDA_VISIBLE_DEVICES=$gpu_id \ 106 | ${tao_triton_server_docker}:${tao_triton_server_tag} \ 107 | /tao_triton/download_and_convert.sh 108 | -------------------------------------------------------------------------------- /tao_triton/python/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | -------------------------------------------------------------------------------- /tao_triton/python/clustering_specs/clustering_config_centerpose.prototxt: -------------------------------------------------------------------------------- 1 | visualization_threshold: 0.3 2 | principle_point_x: 300.3225504557292 3 | principle_point_y: 400.1635182698568 4 | focal_length_x: 654.2994384765625 5 | focal_length_y: 654.2994384765625 6 | skew: 0.0 7 | axis_size: 0.5 8 | square_size: 10 9 | line_weight: 2 10 | scale_text: 0.5 -------------------------------------------------------------------------------- /tao_triton/python/clustering_specs/clustering_config_dashcamnet.prototxt: -------------------------------------------------------------------------------- 1 | linewidth: 4 2 | stride: 16 3 | classwise_clustering_config{ 4 | key: "car" 5 | value: { 6 | coverage_threshold: 0.005 7 | minimum_bounding_box_height: 4 8 | dbscan_config{ 9 | dbscan_eps: 0.3 10 | dbscan_min_samples: 0.05 11 | dbscan_confidence_threshold: 0.9 12 | } 13 | bbox_color{ 14 | R: 0 15 | G: 255 16 | B: 0 17 | } 18 | } 19 | } 20 | classwise_clustering_config{ 21 | key: "bicycle" 22 | value: { 23 | coverage_threshold: 0.005 24 | minimum_bounding_box_height: 4 25 | dbscan_config{ 26 | dbscan_eps: 0.3 27 | dbscan_min_samples: 0.05 28 | dbscan_confidence_threshold: 0.9 29 | } 30 | bbox_color{ 31 | R: 0 32 | G: 255 33 | B: 255 34 | } 35 | } 36 | } 37 | classwise_clustering_config{ 38 | key: "person" 39 | value: { 40 | coverage_threshold: 0.005 41 | minimum_bounding_box_height: 4 42 | dbscan_config{ 43 | dbscan_eps: 0.3 44 | dbscan_min_samples: 0.05 45 | dbscan_confidence_threshold: 0.2 46 | } 47 | bbox_color{ 48 | R: 255 49 | G: 0 50 | B: 0 51 | } 52 | } 53 | } 54 | classwise_clustering_config{ 55 | key: "road_sign" 56 | value: { 57 | coverage_threshold: 0.005 58 | minimum_bounding_box_height: 4 59 | dbscan_config{ 60 | dbscan_eps: 0.3 61 | dbscan_min_samples: 0.05 62 | dbscan_confidence_threshold: 0.2 63 | } 64 | bbox_color{ 65 | R: 255 66 | G: 0 67 | B: 0 68 | } 69 | } 70 | } -------------------------------------------------------------------------------- /tao_triton/python/clustering_specs/clustering_config_peoplenet.prototxt: -------------------------------------------------------------------------------- 1 | linewidth: 4 2 | stride: 16 3 | classwise_clustering_config{ 4 | key: "person" 5 | value: { 6 | coverage_threshold: 0.005 7 | minimum_bounding_box_height: 4 8 | dbscan_config{ 9 | dbscan_eps: 0.3 10 | dbscan_min_samples: 0.05 11 | dbscan_confidence_threshold: 0.9 12 | } 13 | bbox_color{ 14 | R: 0 15 | G: 255 16 | B: 0 17 | } 18 | } 19 | } 20 | classwise_clustering_config{ 21 | key: "bag" 22 | value: { 23 | coverage_threshold: 0.005 24 | minimum_bounding_box_height: 4 25 | dbscan_config{ 26 | dbscan_eps: 0.3 27 | dbscan_min_samples: 0.05 28 | dbscan_confidence_threshold: 0.9 29 | } 30 | bbox_color{ 31 | R: 0 32 | G: 0 33 | B: 255 34 | } 35 | } 36 | } 37 | classwise_clustering_config{ 38 | key: "face" 39 | value: { 40 | coverage_threshold: 0.005 41 | minimum_bounding_box_height: 4 42 | dbscan_config{ 43 | dbscan_eps: 0.3 44 | dbscan_min_samples: 0.05 45 | dbscan_confidence_threshold: 0.7 46 | } 47 | bbox_color{ 48 | R: 0 49 | G: 255 50 | B: 255 51 | } 52 | } 53 | } -------------------------------------------------------------------------------- /tao_triton/python/dataset_convert_specs/dataset_convert_config_pose_classification.yaml: -------------------------------------------------------------------------------- 1 | pose_type: "3dbp" 2 | num_joints: 34 3 | frame_width: 1920 4 | frame_height: 1080 5 | focal_length: 1200.0 6 | sequence_length_max: 300 7 | sequence_length_min: 10 8 | sequence_length: 100 9 | sequence_overlap: 0.5 10 | -------------------------------------------------------------------------------- /tao_triton/python/entrypoints/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | -------------------------------------------------------------------------------- /tao_triton/python/model/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Module for model handling in triton.""" -------------------------------------------------------------------------------- /tao_triton/python/model/centerpose_model.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Triton inference client for TAO Toolkit model.""" 23 | 24 | import os 25 | import numpy as np 26 | 27 | import tritonclient.grpc as grpcclient 28 | import tritonclient.grpc.model_config_pb2 as mc 29 | import tritonclient.http as httpclient 30 | from tritonclient.utils import InferenceServerException 31 | from tritonclient.utils import triton_to_np_dtype 32 | 33 | from tao_triton.python.model.triton_model import TritonModel 34 | 35 | 36 | class CenterPoseModel(TritonModel): 37 | """Simple class to run model inference using Triton client.""" 38 | 39 | def __init__(self, max_batch_size, input_names, output_names, 40 | channels, height, width, data_format, 41 | triton_dtype): 42 | """Set up a centerpose triton model instance. 43 | 44 | Args: 45 | max_batch_size(int): The maximum batch size of the TensorRT engine. 46 | input_names (str): List of the input node names 47 | output_names (str): List of the output node names 48 | channels (int): Number of chanels in the input dimensions 49 | height (int): Height of the input 50 | width (int): Width of the input 51 | data_format (str): The input dimension order. This can be "channels_first" 52 | or "channels_last". "channels_first" is in the CHW order, 53 | and "channels_last" is in HWC order. 54 | triton_dtype (proto): Triton input data type 55 | 56 | Returns: 57 | An instance of the ReIdentificationModel. 58 | """ 59 | super().__init__(max_batch_size, input_names, output_names, 60 | channels, height, width, data_format, 61 | triton_dtype) 62 | 63 | @staticmethod 64 | def parse_model(model_metadata, model_config): 65 | """Parse model metadata and model config from the triton server.""" 66 | if len(model_metadata.inputs) != 1: 67 | raise Exception("expecting 1 input, got {}".format( 68 | len(model_metadata.inputs))) 69 | if len(model_metadata.outputs) != 7: 70 | raise Exception("expecting 7 output, got {}".format( 71 | len(model_metadata.outputs))) 72 | 73 | if len(model_config.input) != 1: 74 | raise Exception( 75 | "expecting 1 input in model configuration, got {}".format( 76 | len(model_config.input))) 77 | 78 | input_metadata = model_metadata.inputs[0] 79 | input_config = model_config.input[0] 80 | output_metadata = model_metadata.outputs 81 | 82 | # Model input must have 3 dims, either CHW or HWC (not counting 83 | # the batch dimension), either CHW or HWC 84 | input_batch_dim = (model_config.max_batch_size > 0) 85 | expected_input_dims = 3 + (1 if input_batch_dim else 0) 86 | if len(input_metadata.shape) != expected_input_dims: 87 | raise Exception( 88 | "expecting input to have {} dimensions, model '{}' input has {}". 89 | format(expected_input_dims, model_metadata.name, 90 | len(input_metadata.shape))) 91 | 92 | if type(input_config.format) == str: 93 | FORMAT_ENUM_TO_INT = dict(mc.ModelInput.Format.items()) 94 | input_config.format = FORMAT_ENUM_TO_INT[input_config.format] 95 | 96 | if ((input_config.format != mc.ModelInput.FORMAT_NCHW) and 97 | (input_config.format != mc.ModelInput.FORMAT_NHWC)): 98 | raise Exception("unexpected input format " + 99 | mc.ModelInput.Format.Name(input_config.format) + 100 | ", expecting " + 101 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NCHW) + 102 | " or " + 103 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NHWC)) 104 | 105 | if input_config.format == mc.ModelInput.FORMAT_NHWC: 106 | h = input_metadata.shape[1 if input_batch_dim else 0] 107 | w = input_metadata.shape[2 if input_batch_dim else 1] 108 | c = input_metadata.shape[3 if input_batch_dim else 2] 109 | else: 110 | c = input_metadata.shape[1 if input_batch_dim else 0] 111 | h = input_metadata.shape[2 if input_batch_dim else 1] 112 | w = input_metadata.shape[3 if input_batch_dim else 2] 113 | 114 | print(model_config.max_batch_size, input_metadata.name, 115 | [data.name for data in output_metadata], c, h, w, input_config.format, 116 | input_metadata.datatype) 117 | 118 | return (model_config.max_batch_size, input_metadata.name, 119 | [data.name for data in output_metadata], c, h, w, input_config.format, 120 | input_metadata.datatype) 121 | -------------------------------------------------------------------------------- /tao_triton/python/model/detectnet_model.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Model class for the DetectNet_v2 Triton Model.""" 23 | 24 | import os 25 | 26 | import tritonclient.grpc as grpcclient 27 | import tritonclient.grpc.model_config_pb2 as mc 28 | import tritonclient.http as httpclient 29 | from tritonclient.utils import InferenceServerException 30 | from tritonclient.utils import triton_to_np_dtype 31 | 32 | from tao_triton.python.model.triton_model import TritonModel 33 | 34 | 35 | class DetectnetModel(TritonModel): 36 | """Simple class to run model inference using Triton client.""" 37 | 38 | def __init__(self, max_batch_size, input_names, output_names, 39 | channels, height, width, data_format, triton_dtype): 40 | """Set up a detectnet_v2 triton model instance. 41 | 42 | Args: 43 | max_batch_size(int): The maximum batch size of the TensorRT engine. 44 | input_names (str): List of the input node names 45 | output_names (str): List of the output node names 46 | channels (int): Number of chanels in the input dimensions 47 | height (int): Height of the input 48 | width (int): Width of the input 49 | data_format (str): The input dimension order. This can be "channels_first" 50 | or "channels_last". "channels_first" is in the CHW order, 51 | and "channels_last" is in HWC order. 52 | triton_dtype (proto): Triton input data type. 53 | channel_mode (str): String order of the C dimension of the input. 54 | "RGB" or "BGR" 55 | 56 | Returns: 57 | An instance of the DetectnetModel. 58 | """ 59 | super().__init__(max_batch_size, input_names, output_names, 60 | channels, height, width, data_format, 61 | triton_dtype) 62 | self.scale = 1. / 255.0 63 | 64 | @staticmethod 65 | def parse_model(model_metadata, model_config): 66 | """Simple class to parse model metadata and model config.""" 67 | if len(model_metadata.inputs) != 1: 68 | raise Exception("expecting 1 input, got {}".format( 69 | len(model_metadata.inputs))) 70 | 71 | if len(model_metadata.outputs) != 2: 72 | raise Exception("expecting 2 output, got {}".format( 73 | len(model_metadata.outputs))) 74 | 75 | if len(model_config.input) != 1: 76 | raise Exception( 77 | "expecting 1 input in model configuration, got {}".format( 78 | len(model_config.input))) 79 | 80 | if len(model_config.output) != 2: 81 | raise Exception( 82 | "expecting 2 outputs in model configuration, got {}".format( 83 | len(model_config.output))) 84 | 85 | input_metadata = model_metadata.inputs[0] 86 | input_config = model_config.input[0] 87 | output_metadata = model_metadata.outputs 88 | 89 | for data in output_metadata: 90 | if data.datatype != "FP32": 91 | raise Exception("expecting output datatype to be FP32, model '" + 92 | data.name + "' output type is " + 93 | data.datatype) 94 | 95 | # Model input must have 3 dims, either CHW or HWC (not counting 96 | # the batch dimension), either CHW or HWC 97 | input_batch_dim = (model_config.max_batch_size > 0) 98 | expected_input_dims = 3 + (1 if input_batch_dim else 0) 99 | if len(input_metadata.shape) != expected_input_dims: 100 | raise Exception( 101 | "expecting input to have {} dimensions, model '{}' input has {}". 102 | format(expected_input_dims, model_metadata.name, 103 | len(input_metadata.shape))) 104 | 105 | if type(input_config.format) == str: 106 | FORMAT_ENUM_TO_INT = dict(mc.ModelInput.Format.items()) 107 | input_config.format = FORMAT_ENUM_TO_INT[input_config.format] 108 | 109 | if ((input_config.format != mc.ModelInput.FORMAT_NCHW) and 110 | (input_config.format != mc.ModelInput.FORMAT_NHWC)): 111 | raise Exception("unexpected input format " + 112 | mc.ModelInput.Format.Name(input_config.format) + 113 | ", expecting " + 114 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NCHW) + 115 | " or " + 116 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NHWC)) 117 | 118 | if input_config.format == mc.ModelInput.FORMAT_NHWC: 119 | h = input_metadata.shape[1 if input_batch_dim else 0] 120 | w = input_metadata.shape[2 if input_batch_dim else 1] 121 | c = input_metadata.shape[3 if input_batch_dim else 2] 122 | else: 123 | c = input_metadata.shape[1 if input_batch_dim else 0] 124 | h = input_metadata.shape[2 if input_batch_dim else 1] 125 | w = input_metadata.shape[3 if input_batch_dim else 2] 126 | 127 | # This part should be be where the input and output names are returned. 128 | return (model_config.max_batch_size, input_metadata.name, 129 | [data.name for data in output_metadata], c, h, w, input_config.format, 130 | input_metadata.datatype) 131 | -------------------------------------------------------------------------------- /tao_triton/python/model/lprnet_model.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Triton inference client for TAO Toolkit model.""" 23 | 24 | import os 25 | import numpy as np 26 | 27 | import tritonclient.grpc as grpcclient 28 | import tritonclient.grpc.model_config_pb2 as mc 29 | import tritonclient.http as httpclient 30 | from tritonclient.utils import InferenceServerException 31 | from tritonclient.utils import triton_to_np_dtype 32 | 33 | from tao_triton.python.model.triton_model import TritonModel 34 | 35 | CHANNEL_MODES = ["rgb", "bgr", "l"] 36 | 37 | 38 | class LPRModel(TritonModel): 39 | """Simple class to run model inference using Triton client.""" 40 | 41 | def __init__(self, max_batch_size, input_names, output_names, 42 | channels, height, width, data_format, 43 | triton_dtype, channel_mode="RGB"): 44 | """Set up a lpr triton model instance. 45 | 46 | Args: 47 | max_batch_size(int): The maximum batch size of the TensorRT engine. 48 | input_names (str): List of the input node names 49 | output_names (str): List of the output node names 50 | channels (int): Number of chanels in the input dimensions 51 | height (int): Height of the input 52 | width (int): Width of the input 53 | data_format (str): The input dimension order. This can be "channels_first" 54 | or "channels_last". "channels_first" is in the CHW order, 55 | and "channels_last" is in HWC order. 56 | triton_dtype (proto): Triton input data type. 57 | channel_mode (str): String order of the C dimension of the input. 58 | "RGB" or "BGR" 59 | 60 | Returns: 61 | An instance of the LPRNetModel. 62 | """ 63 | super().__init__(max_batch_size, input_names, output_names, 64 | channels, height, width, data_format, 65 | triton_dtype) 66 | self.scale = 1. / 255.0 67 | 68 | @staticmethod 69 | def parse_model(model_metadata, model_config): 70 | """Parse model metadata and model config from the triton server.""" 71 | if len(model_metadata.inputs) != 1: 72 | raise Exception("expecting 1 input, got {}".format( 73 | len(model_metadata.inputs))) 74 | 75 | if len(model_metadata.outputs) != 2: 76 | raise Exception("expecting 2 output, got {}".format( 77 | len(model_metadata.outputs))) 78 | 79 | if len(model_config.input) != 1: 80 | raise Exception( 81 | "expecting 1 input in model configuration, got {}".format( 82 | len(model_config.input))) 83 | 84 | if len(model_config.output) != 2: 85 | raise Exception( 86 | "expecting 2 input in model configuration, got {}".format( 87 | len(model_config.input))) 88 | 89 | input_metadata = model_metadata.inputs[0] 90 | input_config = model_config.input[0] 91 | output_metadata = model_metadata.outputs 92 | 93 | for _, data in enumerate(output_metadata): 94 | if _ == 0 : 95 | if data.datatype != "INT32": 96 | raise Exception("expecting output datatype to be INT32, model '" + 97 | data.name + "' output type is " + 98 | data.datatype) 99 | if _ == 1 : 100 | if data.datatype != "FP32": 101 | raise Exception("expecting output datatype to be FP32, model '" + 102 | data.name + "' output type is " + 103 | data.datatype) 104 | 105 | # Model input must have 3 dims, either CHW or HWC (not counting 106 | # the batch dimension), either CHW or HWC 107 | input_batch_dim = (model_config.max_batch_size > 0) 108 | expected_input_dims = 3 + (1 if input_batch_dim else 0) 109 | if len(input_metadata.shape) != expected_input_dims: 110 | raise Exception( 111 | "expecting input to have {} dimensions, model '{}' input has {}". 112 | format(expected_input_dims, model_metadata.name, 113 | len(input_metadata.shape))) 114 | 115 | if type(input_config.format) == str: 116 | FORMAT_ENUM_TO_INT = dict(mc.ModelInput.Format.items()) 117 | input_config.format = FORMAT_ENUM_TO_INT[input_config.format] 118 | 119 | if ((input_config.format != mc.ModelInput.FORMAT_NCHW) and 120 | (input_config.format != mc.ModelInput.FORMAT_NHWC)): 121 | raise Exception("unexpected input format " + 122 | mc.ModelInput.Format.Name(input_config.format) + 123 | ", expecting " + 124 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NCHW) + 125 | " or " + 126 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NHWC)) 127 | 128 | if input_config.format == mc.ModelInput.FORMAT_NHWC: 129 | h = input_metadata.shape[1 if input_batch_dim else 0] 130 | w = input_metadata.shape[2 if input_batch_dim else 1] 131 | c = input_metadata.shape[3 if input_batch_dim else 2] 132 | else: 133 | c = input_metadata.shape[1 if input_batch_dim else 0] 134 | h = input_metadata.shape[2 if input_batch_dim else 1] 135 | w = input_metadata.shape[3 if input_batch_dim else 2] 136 | 137 | print(model_config.max_batch_size, input_metadata.name, 138 | [data.name for data in output_metadata], c, h, w, input_config.format, 139 | input_metadata.datatype) 140 | 141 | return (model_config.max_batch_size, input_metadata.name, 142 | [data.name for data in output_metadata], c, h, w, input_config.format, 143 | input_metadata.datatype) 144 | -------------------------------------------------------------------------------- /tao_triton/python/model/multitask_classification_model.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Triton inference client for TAO Toolkit model.""" 23 | 24 | import os 25 | import numpy as np 26 | 27 | import tritonclient.grpc as grpcclient 28 | import tritonclient.grpc.model_config_pb2 as mc 29 | import tritonclient.http as httpclient 30 | from tritonclient.utils import InferenceServerException 31 | from tritonclient.utils import triton_to_np_dtype 32 | 33 | from tao_triton.python.model.triton_model import TritonModel 34 | 35 | CHANNEL_MODES = ["rgb", "bgr", "l"] 36 | 37 | 38 | class MultitaskClassificationModel(TritonModel): 39 | """Simple class to run model inference using Triton client.""" 40 | 41 | def __init__(self, max_batch_size, input_names, output_names, 42 | channels, height, width, data_format, 43 | triton_dtype, channel_mode="RGB"): 44 | """Set up a multitask_classification triton model instance. 45 | 46 | Args: 47 | max_batch_size(int): The maximum batch size of the TensorRT engine. 48 | input_names (str): List of the input node names 49 | output_names (str): List of the output node names 50 | channels (int): Number of chanels in the input dimensions 51 | height (int): Height of the input 52 | width (int): Width of the input 53 | data_format (str): The input dimension order. This can be "channels_first" 54 | or "channels_last". "channels_first" is in the CHW order, 55 | and "channels_last" is in HWC order. 56 | triton_dtype (proto): Triton input data type. 57 | channel_mode (str): String order of the C dimension of the input. 58 | "RGB" or "BGR" 59 | 60 | Returns: 61 | An instance of the MultitaskClassificationModel. 62 | """ 63 | super().__init__(max_batch_size, input_names, output_names, 64 | channels, height, width, data_format, 65 | triton_dtype) 66 | self.scale = 1.0 67 | if channels == 1: 68 | self.mean = [117.3786] 69 | elif channels == 3: 70 | self.mean = [103.939, 116.779, 123.68] 71 | 72 | if channel_mode.lower() == "rgb": 73 | # Swap channels to make sure that they are in the 74 | # correct channel order. 75 | self.mean.reverse() 76 | self.mean = np.asarray(self.mean).astype(np.float32) 77 | if self.data_format == mc.ModelInput.FORMAT_NCHW: 78 | self.mean = self.mean[:, np.newaxis, np.newaxis] 79 | 80 | @staticmethod 81 | def parse_model(model_metadata, model_config): 82 | """Parse model metadata and model config from the triton server.""" 83 | if len(model_metadata.inputs) != 1: 84 | raise Exception("expecting 1 input, got {}".format( 85 | len(model_metadata.inputs))) 86 | if len(model_metadata.outputs) != 3: 87 | raise Exception("expecting 3 output, got {}".format( 88 | len(model_metadata.outputs))) 89 | 90 | if len(model_config.input) != 1: 91 | raise Exception( 92 | "expecting 1 input in model configuration, got {}".format( 93 | len(model_config.input))) 94 | 95 | input_metadata = model_metadata.inputs[0] 96 | input_config = model_config.input[0] 97 | output_metadata = model_metadata.outputs 98 | 99 | for _, data in enumerate(output_metadata): 100 | if data.datatype != "FP32": 101 | raise Exception("expecting output datatype to be FP32, model '" + 102 | data.name + "' output type is " + 103 | data.datatype) 104 | 105 | # Model input must have 3 dims, either CHW or HWC (not counting 106 | # the batch dimension), either CHW or HWC 107 | input_batch_dim = (model_config.max_batch_size > 0) 108 | expected_input_dims = 3 + (1 if input_batch_dim else 0) 109 | if len(input_metadata.shape) != expected_input_dims: 110 | raise Exception( 111 | "expecting input to have {} dimensions, model '{}' input has {}". 112 | format(expected_input_dims, model_metadata.name, 113 | len(input_metadata.shape))) 114 | 115 | if type(input_config.format) == str: 116 | FORMAT_ENUM_TO_INT = dict(mc.ModelInput.Format.items()) 117 | input_config.format = FORMAT_ENUM_TO_INT[input_config.format] 118 | 119 | if ((input_config.format != mc.ModelInput.FORMAT_NCHW) and 120 | (input_config.format != mc.ModelInput.FORMAT_NHWC)): 121 | raise Exception("unexpected input format " + 122 | mc.ModelInput.Format.Name(input_config.format) + 123 | ", expecting " + 124 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NCHW) + 125 | " or " + 126 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NHWC)) 127 | 128 | if input_config.format == mc.ModelInput.FORMAT_NHWC: 129 | h = input_metadata.shape[1 if input_batch_dim else 0] 130 | w = input_metadata.shape[2 if input_batch_dim else 1] 131 | c = input_metadata.shape[3 if input_batch_dim else 2] 132 | else: 133 | c = input_metadata.shape[1 if input_batch_dim else 0] 134 | h = input_metadata.shape[2 if input_batch_dim else 1] 135 | w = input_metadata.shape[3 if input_batch_dim else 2] 136 | 137 | print(model_config.max_batch_size, input_metadata.name, 138 | output_metadata[0].name, output_metadata[1].name, output_metadata[2].name, 139 | c, h, w, input_config.format, 140 | input_metadata.datatype) 141 | 142 | return (model_config.max_batch_size, input_metadata.name, 143 | [output_metadata[0].name, output_metadata[1].name, output_metadata[2].name], c, h, w, input_config.format, 144 | input_metadata.datatype) 145 | -------------------------------------------------------------------------------- /tao_triton/python/model/peoplesegnet_model.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Triton inference client for TAO Toolkit model.""" 23 | 24 | import os 25 | import numpy as np 26 | 27 | import tritonclient.grpc as grpcclient 28 | import tritonclient.grpc.model_config_pb2 as mc 29 | import tritonclient.http as httpclient 30 | from tritonclient.utils import InferenceServerException 31 | from tritonclient.utils import triton_to_np_dtype 32 | 33 | from tao_triton.python.model.triton_model import TritonModel 34 | 35 | CHANNEL_MODES = ["rgb", "bgr", "l"] 36 | 37 | 38 | class PeoplesegnetModel(TritonModel): 39 | """Simple class to run model inference using Triton client.""" 40 | 41 | def __init__(self, max_batch_size, input_names, output_names, 42 | channels, height, width, data_format, 43 | triton_dtype, channel_mode="RGB"): 44 | """Set up a peoplesegnet triton model instance. 45 | 46 | Args: 47 | max_batch_size(int): The maximum batch size of the TensorRT engine. 48 | input_names (str): List of the input node names 49 | output_names (str): List of the output node names 50 | channels (int): Number of chanels in the input dimensions 51 | height (int): Height of the input 52 | width (int): Width of the input 53 | data_format (str): The input dimension order. This can be "channels_first" 54 | or "channels_last". "channels_first" is in the CHW order, 55 | and "channels_last" is in HWC order. 56 | triton_dtype (proto): Triton input data type. 57 | channel_mode (str): String order of the C dimension of the input. 58 | "RGB" or "BGR" 59 | 60 | Returns: 61 | An instance of the PeoplesegnetModel. 62 | """ 63 | super().__init__(max_batch_size, input_names, output_names, 64 | channels, height, width, data_format, 65 | triton_dtype) 66 | self.scale = 0.017507 67 | 68 | @staticmethod 69 | def parse_model(model_metadata, model_config): 70 | """Parse model metadata and model config from the triton server.""" 71 | if len(model_metadata.inputs) != 1: 72 | raise Exception("expecting 1 input, got {}".format( 73 | len(model_metadata.inputs))) 74 | 75 | if len(model_metadata.outputs) != 2: 76 | raise Exception("expecting 2 output, got {}".format( 77 | len(model_metadata.outputs))) 78 | 79 | if len(model_config.input) != 1: 80 | raise Exception( 81 | "expecting 1 input in model configuration, got {}".format( 82 | len(model_config.input))) 83 | 84 | if len(model_config.output) != 2: 85 | raise Exception( 86 | "expecting 2 input in model configuration, got {}".format( 87 | len(model_config.input))) 88 | 89 | input_metadata = model_metadata.inputs[0] 90 | input_config = model_config.input[0] 91 | output_metadata = model_metadata.outputs 92 | 93 | 94 | for _, data in enumerate(output_metadata): 95 | if _ == 0 : 96 | if data.datatype != "FP32": 97 | raise Exception("expecting output datatype to be FP32, model '" + 98 | data.name + "' output type is " + 99 | data.datatype) 100 | if _ != 0 : 101 | if data.datatype != "FP32": 102 | raise Exception("expecting output datatype to be FP32, model '" + 103 | data.name + "' output type is " + 104 | data.datatype) 105 | 106 | # Model input must have 3 dims, either CHW or HWC (not counting 107 | # the batch dimension), either CHW or HWC 108 | input_batch_dim = (model_config.max_batch_size > 0) 109 | expected_input_dims = 3 + (1 if input_batch_dim else 0) 110 | if len(input_metadata.shape) != expected_input_dims: 111 | raise Exception( 112 | "expecting input to have {} dimensions, model '{}' input has {}". 113 | format(expected_input_dims, model_metadata.name, 114 | len(input_metadata.shape))) 115 | 116 | if type(input_config.format) == str: 117 | FORMAT_ENUM_TO_INT = dict(mc.ModelInput.Format.items()) 118 | input_config.format = FORMAT_ENUM_TO_INT[input_config.format] 119 | 120 | if ((input_config.format != mc.ModelInput.FORMAT_NCHW) and 121 | (input_config.format != mc.ModelInput.FORMAT_NHWC)): 122 | raise Exception("unexpected input format " + 123 | mc.ModelInput.Format.Name(input_config.format) + 124 | ", expecting " + 125 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NCHW) + 126 | " or " + 127 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NHWC)) 128 | 129 | if input_config.format == mc.ModelInput.FORMAT_NHWC: 130 | h = input_metadata.shape[1 if input_batch_dim else 0] 131 | w = input_metadata.shape[2 if input_batch_dim else 1] 132 | c = input_metadata.shape[3 if input_batch_dim else 2] 133 | else: 134 | c = input_metadata.shape[1 if input_batch_dim else 0] 135 | h = input_metadata.shape[2 if input_batch_dim else 1] 136 | w = input_metadata.shape[3 if input_batch_dim else 2] 137 | 138 | print(model_config.max_batch_size, input_metadata.name, 139 | [data.name for data in output_metadata], c, h, w, input_config.format, 140 | input_metadata.datatype) 141 | 142 | return (model_config.max_batch_size, input_metadata.name, 143 | [data.name for data in output_metadata], c, h, w, input_config.format, 144 | input_metadata.datatype) 145 | -------------------------------------------------------------------------------- /tao_triton/python/model/pose_classification_model.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Triton inference client for TAO Toolkit model.""" 23 | 24 | import os 25 | import numpy as np 26 | 27 | import tritonclient.grpc as grpcclient 28 | import tritonclient.grpc.model_config_pb2 as mc 29 | import tritonclient.http as httpclient 30 | from tritonclient.utils import InferenceServerException 31 | from tritonclient.utils import triton_to_np_dtype 32 | 33 | from tao_triton.python.model.triton_model import TritonModel 34 | 35 | 36 | class PoseClassificationModel(TritonModel): 37 | """Simple class to run model inference using Triton client.""" 38 | 39 | def __init__(self, max_batch_size, input_names, output_names, 40 | channels, seq_length, num_joint, num_person, 41 | triton_dtype): 42 | """Set up a pose_classification triton model instance. 43 | 44 | Args: 45 | max_batch_size(int): The maximum batch size of the TensorRT engine. 46 | input_names (str): List of the input node names 47 | output_names (str): List of the output node names 48 | channels (int): Number of chanels in the input dimensions 49 | seq_length (int): Length of the sequences 50 | num_joint (int): Number of joint points 51 | num_person (int): Number of persons 52 | triton_dtype (proto): Triton input data type 53 | 54 | Returns: 55 | An instance of the PoseClassificationModel. 56 | """ 57 | self.max_batch_size = max_batch_size 58 | self.input_names = input_names 59 | self.output_names = output_names 60 | self.channels = channels 61 | self.seq_length = seq_length 62 | self.num_joint = num_joint 63 | self.num_person = num_person 64 | self.triton_dtype = triton_dtype 65 | 66 | @staticmethod 67 | def parse_model(model_metadata, model_config): 68 | """Parse model metadata and model config from the triton server.""" 69 | if len(model_metadata.inputs) != 1: 70 | raise Exception("expecting 1 input, got {}".format( 71 | len(model_metadata.inputs))) 72 | if len(model_metadata.outputs) != 1: 73 | raise Exception("expecting 1 output, got {}".format( 74 | len(model_metadata.outputs))) 75 | 76 | if len(model_config.input) != 1: 77 | raise Exception( 78 | "expecting 1 input in model configuration, got {}".format( 79 | len(model_config.input))) 80 | 81 | input_metadata = model_metadata.inputs[0] 82 | input_config = model_config.input[0] 83 | output_metadata = model_metadata.outputs[0] 84 | 85 | if output_metadata.datatype != "FP32": 86 | raise Exception("expecting output datatype to be FP32, model '" + 87 | model_metadata.name + "' output type is " + 88 | output_metadata.datatype) 89 | 90 | # Model input must have 4 dims CTVM (not counting the batch dimension) 91 | input_batch_dim = (model_config.max_batch_size > 0) 92 | expected_input_dims = 4 + (1 if input_batch_dim else 0) 93 | if len(input_metadata.shape) != expected_input_dims: 94 | raise Exception( 95 | "expecting input to have {} dimensions, model '{}' input has {}". 96 | format(expected_input_dims, model_metadata.name, 97 | len(input_metadata.shape))) 98 | 99 | c = input_metadata.shape[1 if input_batch_dim else 0] 100 | t = input_metadata.shape[2 if input_batch_dim else 1] 101 | v = input_metadata.shape[3 if input_batch_dim else 2] 102 | m = input_metadata.shape[4 if input_batch_dim else 3] 103 | 104 | print(model_config.max_batch_size, input_metadata.name, 105 | output_metadata.name, c, t, v, m, 106 | input_metadata.datatype) 107 | 108 | return (model_config.max_batch_size, input_metadata.name, 109 | [output_metadata.name], c, t, v, m, 110 | input_metadata.datatype) 111 | -------------------------------------------------------------------------------- /tao_triton/python/model/re_identification_model.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Triton inference client for TAO Toolkit model.""" 23 | 24 | import os 25 | import numpy as np 26 | 27 | import tritonclient.grpc as grpcclient 28 | import tritonclient.grpc.model_config_pb2 as mc 29 | import tritonclient.http as httpclient 30 | from tritonclient.utils import InferenceServerException 31 | from tritonclient.utils import triton_to_np_dtype 32 | 33 | from tao_triton.python.model.triton_model import TritonModel 34 | 35 | 36 | class ReIdentificationModel(TritonModel): 37 | """Simple class to run model inference using Triton client.""" 38 | 39 | def __init__(self, max_batch_size, input_names, output_names, 40 | channels, height, width, data_format, 41 | triton_dtype): 42 | """Set up a re_identification triton model instance. 43 | 44 | Args: 45 | max_batch_size(int): The maximum batch size of the TensorRT engine. 46 | input_names (str): List of the input node names 47 | output_names (str): List of the output node names 48 | channels (int): Number of chanels in the input dimensions 49 | height (int): Height of the input 50 | width (int): Width of the input 51 | data_format (str): The input dimension order. This can be "channels_first" 52 | or "channels_last". "channels_first" is in the CHW order, 53 | and "channels_last" is in HWC order. 54 | triton_dtype (proto): Triton input data type 55 | 56 | Returns: 57 | An instance of the ReIdentificationModel. 58 | """ 59 | super().__init__(max_batch_size, input_names, output_names, 60 | channels, height, width, data_format, 61 | triton_dtype) 62 | 63 | @staticmethod 64 | def parse_model(model_metadata, model_config): 65 | """Parse model metadata and model config from the triton server.""" 66 | if len(model_metadata.inputs) != 1: 67 | raise Exception("expecting 1 input, got {}".format( 68 | len(model_metadata.inputs))) 69 | if len(model_metadata.outputs) != 1: 70 | raise Exception("expecting 1 output, got {}".format( 71 | len(model_metadata.outputs))) 72 | 73 | if len(model_config.input) != 1: 74 | raise Exception( 75 | "expecting 1 input in model configuration, got {}".format( 76 | len(model_config.input))) 77 | 78 | input_metadata = model_metadata.inputs[0] 79 | input_config = model_config.input[0] 80 | output_metadata = model_metadata.outputs[0] 81 | 82 | if output_metadata.datatype != "FP32": 83 | raise Exception("expecting output datatype to be FP32, model '" + 84 | model_metadata.name + "' output type is " + 85 | output_metadata.datatype) 86 | 87 | # Model input must have 3 dims, either CHW or HWC (not counting 88 | # the batch dimension), either CHW or HWC 89 | input_batch_dim = (model_config.max_batch_size > 0) 90 | expected_input_dims = 3 + (1 if input_batch_dim else 0) 91 | if len(input_metadata.shape) != expected_input_dims: 92 | raise Exception( 93 | "expecting input to have {} dimensions, model '{}' input has {}". 94 | format(expected_input_dims, model_metadata.name, 95 | len(input_metadata.shape))) 96 | 97 | if type(input_config.format) == str: 98 | FORMAT_ENUM_TO_INT = dict(mc.ModelInput.Format.items()) 99 | input_config.format = FORMAT_ENUM_TO_INT[input_config.format] 100 | 101 | if ((input_config.format != mc.ModelInput.FORMAT_NCHW) and 102 | (input_config.format != mc.ModelInput.FORMAT_NHWC)): 103 | raise Exception("unexpected input format " + 104 | mc.ModelInput.Format.Name(input_config.format) + 105 | ", expecting " + 106 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NCHW) + 107 | " or " + 108 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NHWC)) 109 | 110 | if input_config.format == mc.ModelInput.FORMAT_NHWC: 111 | h = input_metadata.shape[1 if input_batch_dim else 0] 112 | w = input_metadata.shape[2 if input_batch_dim else 1] 113 | c = input_metadata.shape[3 if input_batch_dim else 2] 114 | else: 115 | c = input_metadata.shape[1 if input_batch_dim else 0] 116 | h = input_metadata.shape[2 if input_batch_dim else 1] 117 | w = input_metadata.shape[3 if input_batch_dim else 2] 118 | 119 | print(model_config.max_batch_size, input_metadata.name, 120 | output_metadata.name, c, h, w, input_config.format, 121 | input_metadata.datatype) 122 | 123 | return (model_config.max_batch_size, input_metadata.name, 124 | [output_metadata.name], c, h, w, input_config.format, 125 | input_metadata.datatype) 126 | -------------------------------------------------------------------------------- /tao_triton/python/model/retinanet_model.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Triton inference client for TAO Toolkit model.""" 23 | 24 | import os 25 | import numpy as np 26 | 27 | import tritonclient.grpc as grpcclient 28 | import tritonclient.grpc.model_config_pb2 as mc 29 | import tritonclient.http as httpclient 30 | from tritonclient.utils import InferenceServerException 31 | from tritonclient.utils import triton_to_np_dtype 32 | 33 | from tao_triton.python.model.triton_model import TritonModel 34 | 35 | CHANNEL_MODES = ["rgb", "bgr", "l"] 36 | 37 | 38 | class RetinanetModel(TritonModel): 39 | """Simple class to run model inference using Triton client.""" 40 | 41 | def __init__(self, max_batch_size, input_names, output_names, 42 | channels, height, width, data_format, 43 | triton_dtype, channel_mode="RGB"): 44 | """Set up a retinanet triton model instance. 45 | 46 | Args: 47 | max_batch_size(int): The maximum batch size of the TensorRT engine. 48 | input_names (str): List of the input node names 49 | output_names (str): List of the output node names 50 | channels (int): Number of chanels in the input dimensions 51 | height (int): Height of the input 52 | width (int): Width of the input 53 | data_format (str): The input dimension order. This can be "channels_first" 54 | or "channels_last". "channels_first" is in the CHW order, 55 | and "channels_last" is in HWC order. 56 | triton_dtype (proto): Triton input data type. 57 | channel_mode (str): String order of the C dimension of the input. 58 | "RGB" or "BGR" 59 | 60 | Returns: 61 | An instance of the RetinanetModel. 62 | """ 63 | super().__init__(max_batch_size, input_names, output_names, 64 | channels, height, width, data_format, 65 | triton_dtype) 66 | self.scale = 1.0 67 | 68 | @staticmethod 69 | def parse_model(model_metadata, model_config): 70 | """Parse model metadata and model config from the triton server.""" 71 | if len(model_metadata.inputs) != 1: 72 | raise Exception("expecting 1 input, got {}".format( 73 | len(model_metadata.inputs))) 74 | 75 | if len(model_metadata.outputs) != 2: 76 | raise Exception("expecting 2 output, got {}".format( 77 | len(model_metadata.outputs))) 78 | 79 | if len(model_config.input) != 1: 80 | raise Exception( 81 | "expecting 1 input in model configuration, got {}".format( 82 | len(model_config.input))) 83 | 84 | if len(model_config.output) != 2: 85 | raise Exception( 86 | "expecting 2 output in model configuration, got {}".format( 87 | len(model_config.output))) 88 | 89 | input_metadata = model_metadata.inputs[0] 90 | input_config = model_config.input[0] 91 | output_metadata = model_metadata.outputs 92 | 93 | 94 | for _, data in enumerate(output_metadata): 95 | print("output_metadata is {}".format(output_metadata)) 96 | if _ == 0 : 97 | if data.datatype != "FP32": 98 | raise Exception("expecting output datatype to be FP32, model '" + 99 | data.name + "' output type is " + 100 | data.datatype) 101 | 102 | if _ != 0 : 103 | if data.datatype != "FP32": 104 | raise Exception("expecting output datatype to be FP32, model '" + 105 | data.name + "' output type is " + 106 | data.datatype) 107 | 108 | # Model input must have 3 dims, either CHW or HWC (not counting 109 | # the batch dimension), either CHW or HWC 110 | input_batch_dim = (model_config.max_batch_size > 0) 111 | expected_input_dims = 3 + (1 if input_batch_dim else 0) 112 | if len(input_metadata.shape) != expected_input_dims: 113 | raise Exception( 114 | "expecting input to have {} dimensions, model '{}' input has {}". 115 | format(expected_input_dims, model_metadata.name, 116 | len(input_metadata.shape))) 117 | 118 | if type(input_config.format) == str: 119 | FORMAT_ENUM_TO_INT = dict(mc.ModelInput.Format.items()) 120 | input_config.format = FORMAT_ENUM_TO_INT[input_config.format] 121 | 122 | if ((input_config.format != mc.ModelInput.FORMAT_NCHW) and 123 | (input_config.format != mc.ModelInput.FORMAT_NHWC)): 124 | raise Exception("unexpected input format " + 125 | mc.ModelInput.Format.Name(input_config.format) + 126 | ", expecting " + 127 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NCHW) + 128 | " or " + 129 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NHWC)) 130 | 131 | if input_config.format == mc.ModelInput.FORMAT_NHWC: 132 | h = input_metadata.shape[1 if input_batch_dim else 0] 133 | w = input_metadata.shape[2 if input_batch_dim else 1] 134 | c = input_metadata.shape[3 if input_batch_dim else 2] 135 | else: 136 | c = input_metadata.shape[1 if input_batch_dim else 0] 137 | h = input_metadata.shape[2 if input_batch_dim else 1] 138 | w = input_metadata.shape[3 if input_batch_dim else 2] 139 | 140 | print(model_config.max_batch_size, input_metadata.name, 141 | [data.name for data in output_metadata], c, h, w, input_config.format, 142 | input_metadata.datatype) 143 | 144 | return (model_config.max_batch_size, input_metadata.name, 145 | [data.name for data in output_metadata], c, h, w, input_config.format, 146 | input_metadata.datatype) 147 | -------------------------------------------------------------------------------- /tao_triton/python/model/triton_model.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Triton inference client for TAO Toolkit model.""" 23 | 24 | from abc import abstractmethod 25 | import os 26 | 27 | import tritonclient.grpc as grpcclient 28 | import tritonclient.grpc.model_config_pb2 as mc 29 | import tritonclient.http as httpclient 30 | from tritonclient.utils import InferenceServerException 31 | from tritonclient.utils import triton_to_np_dtype 32 | 33 | import numpy as np 34 | 35 | 36 | class TritonModel(object): 37 | """Simple class to run model inference using Triton client.""" 38 | 39 | def __init__(self, max_batch_size, input_names, output_names, 40 | channels, height, width, data_format, triton_dtype): 41 | """Set up a detectnet_v2 triton model instance. 42 | 43 | Args: 44 | max_batch_size(int): The maximum batch size of the TensorRT engine. 45 | input_names (str): List of the input node names 46 | output_names (str): List of the output node names 47 | channels (int): Number of chanels in the input dimensions 48 | height (int): Height of the input 49 | width (int): Width of the input 50 | data_format (str): The input dimension order. This can be "channels_first" 51 | or "channels_last". "channels_first" is in the CHW order, 52 | and "channels_last" is in HWC order. 53 | triton_dtype (proto): Triton input data type. 54 | channel_mode (str): String order of the C dimension of the input. 55 | "RGB" or "BGR" 56 | 57 | Returns: 58 | An instance of the DetectnetModel. 59 | """ 60 | self.max_batch_size = max_batch_size 61 | self.input_names = input_names 62 | self.output_names = output_names 63 | self.c = channels 64 | assert channels in [1, 3], ( 65 | "TAO Toolkit models only support 1 or 3 channel inputs." 66 | ) 67 | self.h = height 68 | self.w = width 69 | self.data_format = data_format 70 | self.triton_dtype = triton_dtype 71 | self.scale = 1 72 | if channels == 3: 73 | self.mean = [0., 0., 0.] 74 | else: 75 | self.mean = [0] 76 | self.mean = np.asarray(self.mean).astype(np.float32) 77 | if self.data_format == mc.ModelInput.FORMAT_NCHW: 78 | self.mean = self.mean[:, np.newaxis, np.newaxis] 79 | 80 | @staticmethod 81 | def parse_model(model_metadata, model_config): 82 | """Simple class to parse model metadata and model config.""" 83 | raise NotImplementedError("Base class doesn't implement this method.") 84 | 85 | @classmethod 86 | def from_metadata(cls, model_metadata, model_config): 87 | """Parse a model from the metadata config.""" 88 | parsed_outputs = cls.parse_model(model_metadata, model_config) 89 | max_batch_size, input_names, output_names, channels, height, width, \ 90 | data_format, triton_dtype = parsed_outputs 91 | return cls( 92 | max_batch_size, input_names, output_names, 93 | channels, height, width, data_format, 94 | triton_dtype 95 | ) 96 | 97 | def get_config(self): 98 | """Get dictionary config.""" 99 | config_dict = { 100 | "data_format": self.data_format, 101 | "max_batch_size": self.max_batch_size, 102 | "channels": self.c, 103 | "width": self.w, 104 | "height": self.h, 105 | "input_names": self.input_names, 106 | "output_names": self.output_names, 107 | "triton_dtype": self.triton_dtype 108 | } 109 | return config_dict 110 | 111 | def preprocess(self, image): 112 | """Function to preprocess image 113 | 114 | Performs mean subtraction and then normalization. 115 | 116 | Args: 117 | image (np.ndarray): Numpy ndarray of an input batch. 118 | 119 | Returns: 120 | image (np.ndarray): Preprocessed input image. 121 | """ 122 | image = (image - self.mean) * self.scale 123 | return image 124 | -------------------------------------------------------------------------------- /tao_triton/python/model/visual_changenet_model.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Triton inference client for TAO Toolkit model.""" 23 | 24 | import os 25 | import numpy as np 26 | 27 | import tritonclient.grpc as grpcclient 28 | import tritonclient.grpc.model_config_pb2 as mc 29 | import tritonclient.http as httpclient 30 | from tritonclient.utils import InferenceServerException 31 | from tritonclient.utils import triton_to_np_dtype 32 | 33 | from tao_triton.python.model.triton_model import TritonModel 34 | 35 | CHANNEL_MODES = ["rgb", "bgr", "l"] 36 | 37 | 38 | class VisualChangeNetModel(TritonModel): 39 | """Simple class to run model inference using Triton client.""" 40 | 41 | def __init__(self, max_batch_size, input_names, output_names, 42 | channels, height, width, data_format, 43 | triton_dtype, channel_mode="RGB"): 44 | """Set up a visual_changenet triton model instance. 45 | 46 | Args: 47 | max_batch_size(int): The maximum batch size of the TensorRT engine. 48 | input_names (str): List of the input node names 49 | output_names (str): List of the output node names 50 | channels (int): Number of chanels in the input dimensions 51 | height (int): Height of the input 52 | width (int): Width of the input 53 | data_format (str): The input dimension order. This can be "channels_first" 54 | or "channels_last". "channels_first" is in the CHW order, 55 | and "channels_last" is in HWC order. 56 | triton_dtype (proto): Triton input data type. 57 | channel_mode (str): String order of the C dimension of the input. 58 | "RGB" or "BGR" 59 | 60 | Returns: 61 | An instance of the VisualChangeNetModel. 62 | """ 63 | super().__init__(max_batch_size, input_names, output_names, 64 | channels, height, width, data_format, 65 | triton_dtype) 66 | self.scale = 1.0 67 | 68 | @staticmethod 69 | def parse_model(model_metadata, model_config): 70 | """Parse model metadata and model config from the triton server.""" 71 | print(len(model_metadata.inputs)) 72 | if len(model_metadata.inputs) != 2: 73 | raise Exception("expecting 2 input, got {}".format( 74 | len(model_metadata.inputs))) 75 | 76 | if len(model_metadata.outputs) != 5: 77 | raise Exception("expecting 5 output, got {}".format( 78 | len(model_metadata.outputs))) 79 | 80 | if len(model_config.input) != 2: 81 | raise Exception( 82 | "expecting 2 input in model configuration, got {}".format( 83 | len(model_config.input))) 84 | if len(model_config.output) != 5: 85 | raise Exception( 86 | "expecting 5 input in model configuration, got {}".format( 87 | len(model_config.output))) 88 | 89 | input_metadata = model_metadata.inputs 90 | input_config = model_config.input 91 | output_metadata = model_metadata.outputs 92 | 93 | 94 | for _, data in enumerate(output_metadata): 95 | if data.datatype != "FP32": 96 | raise Exception("expecting output datatype to be FP32, model '" + 97 | data.name + "' output type is " + 98 | data.datatype) 99 | 100 | # Model input must have 3 dims, either CHW or HWC (not counting 101 | # the batch dimension), either CHW or HWC 102 | 103 | input_batch_dim = (model_config.max_batch_size > 0) 104 | expected_input_dims = 3 + (1 if input_batch_dim else 0) 105 | if type(input_config)==tuple: 106 | for i in range(len(input_config)): 107 | if len(input_metadata[i].shape) != expected_input_dims: 108 | raise Exception( 109 | "expecting input to have {} dimensions, model '{}' input has {}". 110 | format(expected_input_dims, model_metadata.name, 111 | len(input_metadata[i].shape))) 112 | 113 | if type(input_config)==tuple: 114 | for i in range(len(input_config)): 115 | if type(input_config[i].format) == str: 116 | FORMAT_ENUM_TO_INT = dict(mc.ModelInput.Format.items()) 117 | input_config[i].format = FORMAT_ENUM_TO_INT[input_config[i].format] 118 | 119 | if type(input_config)==tuple: 120 | for i in range(len(input_config)): 121 | if ((input_config[i].format != mc.ModelInput.FORMAT_NCHW) and 122 | (input_config[i].format != mc.ModelInput.FORMAT_NHWC)): 123 | raise Exception("unexpected input format " + 124 | mc.ModelInput.Format.Name(input_config[i].format) + 125 | ", expecting " + 126 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NCHW) + 127 | " or " + 128 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NHWC)) 129 | 130 | if type(input_config)==tuple: #Can take for only 1 of the inputs 131 | for i in range(len(input_config)): 132 | if input_config[i].format == mc.ModelInput.FORMAT_NHWC: 133 | h = input_metadata[i].shape[1 if input_batch_dim else 0] 134 | w = input_metadata[i].shape[2 if input_batch_dim else 1] 135 | c = input_metadata[i].shape[3 if input_batch_dim else 2] 136 | else: 137 | c = input_metadata[i].shape[1 if input_batch_dim else 0] 138 | h = input_metadata[i].shape[2 if input_batch_dim else 1] 139 | w = input_metadata[i].shape[3 if input_batch_dim else 2] 140 | 141 | 142 | print(model_config.max_batch_size, [input_meta.name for input_meta in input_metadata], 143 | [data.name for data in output_metadata], c, h, w, input_config[0].format, 144 | input_metadata[0].datatype) 145 | return (model_config.max_batch_size, [input_meta.name for input_meta in input_metadata], 146 | [data.name for data in output_metadata], c, h, w, input_config[0].format, 147 | input_metadata[0].datatype) 148 | -------------------------------------------------------------------------------- /tao_triton/python/model/yolov3_model.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Triton inference client for TAO Toolkit model.""" 23 | 24 | import os 25 | import numpy as np 26 | 27 | import tritonclient.grpc as grpcclient 28 | import tritonclient.grpc.model_config_pb2 as mc 29 | import tritonclient.http as httpclient 30 | from tritonclient.utils import InferenceServerException 31 | from tritonclient.utils import triton_to_np_dtype 32 | 33 | from tao_triton.python.model.triton_model import TritonModel 34 | 35 | CHANNEL_MODES = ["rgb", "bgr", "l"] 36 | 37 | 38 | class YOLOv3Model(TritonModel): 39 | """Simple class to run model inference using Triton client.""" 40 | 41 | def __init__(self, max_batch_size, input_names, output_names, 42 | channels, height, width, data_format, 43 | triton_dtype, channel_mode="RGB"): 44 | """Set up a yolov3 triton model instance. 45 | 46 | Args: 47 | max_batch_size(int): The maximum batch size of the TensorRT engine. 48 | input_names (str): List of the input node names 49 | output_names (str): List of the output node names 50 | channels (int): Number of chanels in the input dimensions 51 | height (int): Height of the input 52 | width (int): Width of the input 53 | data_format (str): The input dimension order. This can be "channels_first" 54 | or "channels_last". "channels_first" is in the CHW order, 55 | and "channels_last" is in HWC order. 56 | triton_dtype (proto): Triton input data type. 57 | channel_mode (str): String order of the C dimension of the input. 58 | "RGB" or "BGR" 59 | 60 | Returns: 61 | An instance of the YOLOv3Model. 62 | """ 63 | super().__init__(max_batch_size, input_names, output_names, 64 | channels, height, width, data_format, 65 | triton_dtype) 66 | self.scale = 1.0 67 | 68 | @staticmethod 69 | def parse_model(model_metadata, model_config): 70 | """Parse model metadata and model config from the triton server.""" 71 | if len(model_metadata.inputs) != 1: 72 | raise Exception("expecting 1 input, got {}".format( 73 | len(model_metadata.inputs))) 74 | 75 | if len(model_metadata.outputs) != 4: 76 | raise Exception("expecting 4 output, got {}".format( 77 | len(model_metadata.outputs))) 78 | 79 | if len(model_config.input) != 1: 80 | raise Exception( 81 | "expecting 1 input in model configuration, got {}".format( 82 | len(model_config.input))) 83 | 84 | if len(model_config.output) != 4: 85 | raise Exception( 86 | "expecting 2 input in model configuration, got {}".format( 87 | len(model_config.input))) 88 | 89 | input_metadata = model_metadata.inputs[0] 90 | input_config = model_config.input[0] 91 | output_metadata = model_metadata.outputs 92 | 93 | 94 | for _, data in enumerate(output_metadata): 95 | if _ == 0 : 96 | if data.datatype != "INT32": 97 | raise Exception("expecting output datatype to be INT32, model '" + 98 | data.name + "' output type is " + 99 | data.datatype) 100 | if _ != 0 : 101 | if data.datatype != "FP32": 102 | raise Exception("expecting output datatype to be FP32, model '" + 103 | data.name + "' output type is " + 104 | data.datatype) 105 | 106 | # Model input must have 3 dims, either CHW or HWC (not counting 107 | # the batch dimension), either CHW or HWC 108 | input_batch_dim = (model_config.max_batch_size > 0) 109 | expected_input_dims = 3 + (1 if input_batch_dim else 0) 110 | if len(input_metadata.shape) != expected_input_dims: 111 | raise Exception( 112 | "expecting input to have {} dimensions, model '{}' input has {}". 113 | format(expected_input_dims, model_metadata.name, 114 | len(input_metadata.shape))) 115 | 116 | if type(input_config.format) == str: 117 | FORMAT_ENUM_TO_INT = dict(mc.ModelInput.Format.items()) 118 | input_config.format = FORMAT_ENUM_TO_INT[input_config.format] 119 | 120 | if ((input_config.format != mc.ModelInput.FORMAT_NCHW) and 121 | (input_config.format != mc.ModelInput.FORMAT_NHWC)): 122 | raise Exception("unexpected input format " + 123 | mc.ModelInput.Format.Name(input_config.format) + 124 | ", expecting " + 125 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NCHW) + 126 | " or " + 127 | mc.ModelInput.Format.Name(mc.ModelInput.FORMAT_NHWC)) 128 | 129 | if input_config.format == mc.ModelInput.FORMAT_NHWC: 130 | h = input_metadata.shape[1 if input_batch_dim else 0] 131 | w = input_metadata.shape[2 if input_batch_dim else 1] 132 | c = input_metadata.shape[3 if input_batch_dim else 2] 133 | else: 134 | c = input_metadata.shape[1 if input_batch_dim else 0] 135 | h = input_metadata.shape[2 if input_batch_dim else 1] 136 | w = input_metadata.shape[3 if input_batch_dim else 2] 137 | 138 | print(model_config.max_batch_size, input_metadata.name, 139 | [data.name for data in output_metadata], c, h, w, input_config.format, 140 | input_metadata.datatype) 141 | 142 | return (model_config.max_batch_size, input_metadata.name, 143 | [data.name for data in output_metadata], c, h, w, input_config.format, 144 | input_metadata.datatype) 145 | -------------------------------------------------------------------------------- /tao_triton/python/postprocessing/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | -------------------------------------------------------------------------------- /tao_triton/python/postprocessing/classification_postprocessor.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Simple class to run post processing of Triton Inference outputs.""" 23 | 24 | import os 25 | import numpy as np 26 | 27 | from tao_triton.python.postprocessing.postprocessor import Postprocessor 28 | 29 | 30 | class ClassificationPostprocessor(Postprocessor): 31 | """Class to run post processing of Triton Tensors.""" 32 | 33 | def __init__(self, batch_size, frames, output_path, data_format): 34 | """Initialize a post processor class for a classification model. 35 | 36 | Args: 37 | batch_size (int): Number of images in the batch. 38 | frames (list): List of images. 39 | output_path (str): Unix path to the output rendered images and labels. 40 | data_format (str): Order of the input model dimensions. 41 | "channels_first": CHW order. 42 | "channels_last": HWC order. 43 | """ 44 | super().__init__(batch_size, frames, output_path, data_format) 45 | self.output_name = "predictions/Softmax" 46 | 47 | def apply(self, output_tensors, this_id, render=True, batching=True): 48 | """Apply the post processor to the outputs to the classification outputs.""" 49 | output_array = output_tensors.as_numpy(self.output_name) 50 | if len(output_array) != self.batch_size: 51 | raise Exception("expected {} results, got {}".format( 52 | batch_size, len(output_array))) 53 | 54 | # Include special handling for non-batching models 55 | if not os.path.exists(self.output_path): 56 | os.makedirs(self.output_path) 57 | output_file = os.path.join(self.output_path, "results.txt") 58 | with open(output_file, "a") as wfile: 59 | for idx in range(self.batch_size): 60 | results = output_array[idx] 61 | current_idx = (int(this_id) - 1) * self.batch_size + idx 62 | if current_idx < len(self.frames): 63 | wfile.write("{}".format(self.frames[current_idx]._image_path)) 64 | if not batching: 65 | results = [results] 66 | for result in results: 67 | if output_array.dtype.type == np.object_: 68 | cls = "".join(chr(x) for x in result).split(':') 69 | else: 70 | cls = result.split(':') 71 | wfile.write( 72 | ", {:0.4f}({})={}".format( 73 | float(cls[0]), cls[1], cls[2] 74 | ) 75 | ) 76 | wfile.write("\n") 77 | else: 78 | break 79 | -------------------------------------------------------------------------------- /tao_triton/python/postprocessing/lprnet_postprocessor.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Simple class to run post processing of Triton Inference outputs.""" 23 | 24 | import os 25 | import numpy as np 26 | 27 | from tao_triton.python.postprocessing.postprocessor import Postprocessor 28 | 29 | repo_root = os.getenv("TAO_TRITON_REPO_ROOT", "") 30 | characters_list_file = os.path.join(repo_root, "model_repository/lprnet_tao/characters_list.txt") 31 | 32 | def get_classes_id(): 33 | with open(characters_list_file, "r") as f: 34 | temp_list = f.readlines() 35 | classes = [i.strip() for i in temp_list] 36 | blank_id = len(classes) 37 | 38 | return classes, blank_id 39 | 40 | 41 | def decode_ctc_conf(pred, 42 | classes, 43 | blank_id): 44 | ''' 45 | Decode ctc trained model's output. 46 | 47 | return decoded license plate and confidence. 48 | ''' 49 | pred_id = pred['tf_op_layer_ArgMax'] 50 | pred_conf = pred['tf_op_layer_Max'] 51 | decoded_lp = [] 52 | decoded_conf = [] 53 | 54 | for idx_in_batch, seq in enumerate(pred_id): 55 | seq_conf = pred_conf[idx_in_batch] 56 | prev = seq[0] 57 | tmp_seq = [prev] 58 | tmp_conf = [seq_conf[0]] 59 | for idx in range(1, len(seq)): 60 | if seq[idx] != prev: 61 | tmp_seq.append(seq[idx]) 62 | tmp_conf.append(seq_conf[idx]) 63 | prev = seq[idx] 64 | lp = "" 65 | output_conf = [] 66 | for index, i in enumerate(tmp_seq): 67 | if i != blank_id: 68 | lp += classes[i] 69 | output_conf.append(tmp_conf[index]) 70 | decoded_lp.append(lp) 71 | decoded_conf.append(output_conf) 72 | 73 | return decoded_lp, decoded_conf 74 | 75 | 76 | class LPRPostprocessor(Postprocessor): 77 | """Class to run post processing of Triton Tensors.""" 78 | 79 | def __init__(self, batch_size, frames, output_path, data_format): 80 | """Initialize a post processor class for a lprnet model. 81 | 82 | Args: 83 | batch_size (int): Number of images in the batch. 84 | frames (list): List of images. 85 | output_path (str): Unix path to the output rendered images and labels. 86 | data_format (str): Order of the input model dimensions. 87 | "channels_first": CHW order. 88 | "channels_last": HWC order. 89 | """ 90 | super().__init__(batch_size, frames, output_path, data_format) 91 | self.output_names = ["tf_op_layer_ArgMax", 92 | "tf_op_layer_Max"] 93 | 94 | 95 | def apply(self, results, this_id, render=True, batching=True): 96 | """Apply the post processor to the outputs to the lprnet outputs.""" 97 | 98 | output_array = {} 99 | 100 | for output_name in self.output_names: 101 | output_array[output_name] = results.as_numpy(output_name) 102 | 103 | classes, blank_id = get_classes_id() 104 | 105 | decoded_lp, _ = decode_ctc_conf(output_array, 106 | classes=classes, 107 | blank_id=blank_id) 108 | 109 | for idx in range(self.batch_size): 110 | current_idx = (int(this_id) - 1) * self.batch_size + idx 111 | image_id = self.frames[current_idx]._image_path 112 | print(image_id) 113 | 114 | print("inference result: {}\n".format(decoded_lp)) 115 | 116 | 117 | if not os.path.exists(self.output_path): 118 | os.makedirs(self.output_path) 119 | output_file = os.path.join(self.output_path, "results.txt") 120 | with open(output_file, "a") as wfile: 121 | wfile.write("{} : {}\n".format(image_id,decoded_lp)) 122 | -------------------------------------------------------------------------------- /tao_triton/python/postprocessing/multitask_classification_postprocessor.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Simple class to run post processing of Triton Inference outputs.""" 23 | 24 | import os 25 | import numpy as np 26 | 27 | from tao_triton.python.postprocessing.postprocessor import Postprocessor 28 | 29 | 30 | class MultitaskClassificationPostprocessor(Postprocessor): 31 | """Class to run post processing of Triton Tensors.""" 32 | 33 | def __init__(self, batch_size, frames, output_path, data_format): 34 | """Initialize a post processor class for a multitaskclassification model. 35 | 36 | Args: 37 | batch_size (int): Number of images in the batch. 38 | frames (list): List of images. 39 | output_path (str): Unix path to the output rendered images and labels. 40 | data_format (str): Order of the input model dimensions. 41 | "channels_first": CHW order. 42 | "channels_last": HWC order. 43 | """ 44 | super().__init__(batch_size, frames, output_path, data_format) 45 | self.output_name_0 = "base_color/Softmax" 46 | self.output_name_1 = "category/Softmax" 47 | self.output_name_2 = "season/Softmax" 48 | self.task_name = ["base_color", "category", "season"] 49 | self.class_mapping = {"base_color": {"0": "Black", "1": "Blue", "2": "Brown", "3": "Green", \ 50 | "4": "Grey", "5": "Navy Blue", "6": "Pink", "7": "Purple", "8": "Red", \ 51 | "9": "Silver", "10": "White"}, 52 | "category": {"0": "Bags", "1": "Bottomwear", "2": "Eyewear", "3": "Fragrance", \ 53 | "4": "Innerwear", "5": "Jewellery", "6": "Sandal", "7": "Shoes", "8": "Topwear", \ 54 | "9": "Watches"}, 55 | "season": {"0": "Fall", "1": "Spring", "2": "Summer", "3": "Winter"}} 56 | 57 | def apply(self, output_tensors, this_id, render=True, batching=True): 58 | """Apply the post processor to the outputs to the classification outputs.""" 59 | output_array_0 = output_tensors.as_numpy(self.output_name_0) 60 | output_array_1 = output_tensors.as_numpy(self.output_name_1) 61 | output_array_2 = output_tensors.as_numpy(self.output_name_2) 62 | output_array = [output_array_0 , output_array_1 , output_array_2] 63 | 64 | if not os.path.exists(self.output_path): 65 | os.makedirs(self.output_path) 66 | output_file = os.path.join(self.output_path, "results.txt") 67 | 68 | for image_idx in range(self.batch_size): 69 | current_idx = (int(this_id) - 1) * self.batch_size + image_idx 70 | if current_idx >= len(self.frames): 71 | break 72 | current_frame = self.frames[current_idx] 73 | img_in_path = current_frame._image_path 74 | filename = os.path.basename(current_frame._image_path) 75 | print("filename is {}".format(filename)) 76 | 77 | with open(output_file, "a") as j: 78 | j.write("\n{}:\n".format(filename)) 79 | 80 | for idx, task in enumerate(self.task_name): 81 | pred = output_array[idx].reshape(-1) 82 | print("Task {}:".format(task)) 83 | #print("Predictions: {}".format(pred)) 84 | class_name = self.class_mapping[task][str(np.argmax(pred))] 85 | print("Class name = {}".format(class_name)) 86 | print('********') 87 | 88 | with open(output_file, "a") as j: 89 | j.write("{}: {}\n".format(task,class_name)) 90 | -------------------------------------------------------------------------------- /tao_triton/python/postprocessing/pose_classification_postprocessor.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Simple class to run post processing of Triton Inference outputs.""" 23 | 24 | import os 25 | import numpy as np 26 | 27 | from tao_triton.python.postprocessing.postprocessor import Postprocessor 28 | 29 | 30 | class PoseClassificationPostprocessor(Postprocessor): 31 | """Class to run post processing of Triton Tensors.""" 32 | 33 | def __init__(self, batch_size, sequences, output_path): 34 | """Initialize a post processor class for a poseclassification model. 35 | 36 | Args: 37 | batch_size (int): Number of sequences in the batch. 38 | sequences (list): List of sequences. 39 | output_path (str): Unix path to the output rendered sequences and labels. 40 | """ 41 | self.batch_size = batch_size 42 | self.sequences = sequences 43 | self.output_path = output_path 44 | self.output_name = "fc_pred" 45 | 46 | def apply(self, output_tensors, this_id, render=True, batching=True, action_data=[]): 47 | """Apply the post processor to the outputs to the poseclassification outputs.""" 48 | output_array = output_tensors.as_numpy(self.output_name) 49 | if len(output_array) != self.batch_size: 50 | raise Exception("expected {} results, got {}".format( 51 | batch_size, len(output_array))) 52 | 53 | # Include special handling for non-batching models 54 | if not os.path.exists(self.output_path): 55 | os.makedirs(self.output_path) 56 | 57 | # Output a JSON file 58 | if len(action_data) > 0: 59 | for idx in range(self.batch_size): 60 | results = output_array[idx] 61 | current_idx = (int(this_id) - 1) * self.batch_size + idx 62 | if current_idx < len(self.sequences): 63 | for b in range(len(action_data)): 64 | for f in range(len(action_data[b]["batches"])): 65 | frame_num = action_data[b]["batches"][f]["frame_num"] 66 | for p in range(len(action_data[b]["batches"][f]["objects"])): 67 | object_id = action_data[b]["batches"][f]["objects"][p]["object_id"] 68 | segment_id = action_data[b]["batches"][f]["objects"][p]["segment_id"] 69 | if segment_id == current_idx: 70 | if not batching: 71 | results = [results] 72 | result = results[0] 73 | if output_array.dtype.type == np.object_: 74 | cls = "".join(chr(x) for x in result).split(':') 75 | else: 76 | cls = result.split(':') 77 | action_data[b]["batches"][f]["objects"][p]["action"] = cls[2] 78 | 79 | # Output a text file of results 80 | else: 81 | output_file = os.path.join(self.output_path, "results.txt") 82 | with open(output_file, "a") as wfile: 83 | for idx in range(self.batch_size): 84 | results = output_array[idx] 85 | current_idx = (int(this_id) - 1) * self.batch_size + idx 86 | if current_idx < len(self.sequences): 87 | wfile.write("{}".format(current_idx)) 88 | if not batching: 89 | results = [results] 90 | for result in results: 91 | if output_array.dtype.type == np.object_: 92 | cls = "".join(chr(x) for x in result).split(':') 93 | else: 94 | cls = result.split(':') 95 | wfile.write( 96 | ", {:0.4f}({})={}".format( 97 | float(cls[0]), cls[1], cls[2] 98 | ) 99 | ) 100 | wfile.write("\n") 101 | else: 102 | break 103 | -------------------------------------------------------------------------------- /tao_triton/python/postprocessing/postprocessor.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Simple class to run post processing of Triton Inference outputs.""" 23 | 24 | import os 25 | 26 | class Postprocessor(object): 27 | """Class to run post processing of Triton Tensors.""" 28 | 29 | def __init__(self, batch_size, frames, output_path, data_format): 30 | """Initialize a post processor class. 31 | 32 | Args: 33 | batch_size (int): Number of images in the batch. 34 | frames (list): List of images. 35 | output_path (str): Unix path to the output rendered images and labels. 36 | data_format (str): Order of the input model dimensions. 37 | "channels_first": CHW order. 38 | "channels_last": HWC order. 39 | """ 40 | self.batch_size = batch_size 41 | self.frames = frames 42 | self.output_path = output_path 43 | self.data_format = data_format 44 | if not os.path.exists(self.output_path): 45 | os.makedirs(self.output_path) 46 | self.initialized = True 47 | 48 | def apply(self, output_tensors, this_id, render=True): 49 | """Apply the post processor to the outputs.""" 50 | raise NotImplementedError("Base class doesn't implement any post-processing") 51 | -------------------------------------------------------------------------------- /tao_triton/python/postprocessing/re_identification_postprocessor.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Simple class to run post processing of Triton Inference outputs.""" 23 | 24 | import os 25 | import json 26 | 27 | from tao_triton.python.postprocessing.postprocessor import Postprocessor 28 | 29 | 30 | class ReIdentificationPostprocessor(Postprocessor): 31 | """Class to run post processing of Triton Tensors.""" 32 | 33 | def __init__(self, batch_size, query_frames, test_frames, output_path, data_format): 34 | """Initialize a post processor class for a reidentification model. 35 | 36 | Args: 37 | batch_size (int): Number of sequences in the batch. 38 | query_frames (list): List of query images. 39 | test_frames (list): List of test images. 40 | output_path (str): Unix path to the output embeddings. 41 | data_format (str): The input dimension order. This can be "channels_first" 42 | or "channels_last". "channels_first" is in the CHW order, 43 | and "channels_last" is in HWC order. 44 | """ 45 | super().__init__(batch_size, query_frames, output_path, data_format) 46 | self.batch_size = batch_size 47 | self.output_path = output_path 48 | self.output_name = "fc_pred" 49 | self.query_frames = query_frames 50 | self.test_frames = test_frames 51 | 52 | def apply(self, output_tensors, this_id, render=True, batching=True): 53 | """Apply the post processor to the outputs to the reidentification outputs.""" 54 | output_array = output_tensors.as_numpy(self.output_name) 55 | if len(output_array) != self.batch_size: 56 | raise Exception("expected {} results, got {}".format( 57 | self.batch_size, len(output_array))) 58 | if not os.path.exists(self.output_path): 59 | os.makedirs(self.output_path) 60 | output_file = os.path.join(self.output_path, "results.json") 61 | 62 | results = [] 63 | for image_idx in range(self.batch_size): 64 | current_idx = (int(this_id) - 1) * self.batch_size + image_idx 65 | embedding = output_array[image_idx] 66 | 67 | if current_idx < len(self.query_frames): 68 | current_frame = self.query_frames[current_idx] 69 | elif len(self.query_frames) <= current_idx < len(self.query_frames) + len(self.test_frames): 70 | current_frame = self.test_frames[current_idx-len(self.query_frames)] 71 | else: 72 | break 73 | 74 | result = {"img_path": current_frame._image_path, "embedding": embedding.tolist()} 75 | results.append(result) 76 | 77 | if os.path.exists(output_file): 78 | with open(output_file) as fp: 79 | results_list = json.load(fp) 80 | results_list.extend(results) 81 | results = results_list 82 | 83 | with open(output_file, "w") as json_file: 84 | json.dump(results, json_file) 85 | -------------------------------------------------------------------------------- /tao_triton/python/postprocessing/retinanet_postprocessor.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Simple class to run post processing of Triton Inference outputs.""" 23 | 24 | import os 25 | import struct 26 | import numpy as np 27 | from PIL import Image, ImageDraw 28 | 29 | from tao_triton.python.postprocessing.postprocessor import Postprocessor 30 | from tao_triton.python.postprocessing.utils import pool_context 31 | 32 | 33 | 34 | def trt_output_process_fn(self, y_encoded): 35 | "function to process TRT model output." 36 | det_out, keep_k = y_encoded 37 | result = [] 38 | for idx, k in enumerate(keep_k.reshape(-1)): 39 | det = det_out[idx].reshape(-1, 7)[:k] 40 | xmin = det[:, 3] * self.model_input_width 41 | ymin = det[:, 4] * self.model_input_height 42 | xmax = det[:, 5] * self.model_input_width 43 | ymax = det[:, 6] * self.model_input_height 44 | cls_id = det[:, 1] 45 | conf = det[:, 2] 46 | result.append(np.stack((cls_id, conf, xmin, ymin, xmax, ymax), axis=-1)) 47 | 48 | return result 49 | 50 | 51 | 52 | class RetinanetPostprocessor(Postprocessor): 53 | """Class to run post processing of Triton Tensors.""" 54 | 55 | def __init__(self, batch_size, frames, output_path, data_format): 56 | """Initialize a post processor class for a retinanet model. 57 | 58 | Args: 59 | batch_size (int): Number of images in the batch. 60 | frames (list): List of images. 61 | output_path (str): Unix path to the output rendered images and labels. 62 | data_format (str): Order of the input model dimensions. 63 | "channels_first": CHW order. 64 | "channels_last": HWC order. 65 | """ 66 | super().__init__(batch_size, frames, output_path, data_format) 67 | self.output_names = ["NMS", 68 | "NMS_1" ] 69 | self.threshold = 0.6 70 | self.keep_aspect_ratio = True 71 | self.class_mapping = {0: 'background', 1: 'bicycle', 2: 'car', 3: 'person', 4: 'road_sign'} 72 | self.model_input_width = 960 73 | self.model_input_height = 544 74 | 75 | def _get_bbox_and_kitti_label_single_img( 76 | self, img, img_ratio, y_decoded, 77 | is_draw_img, is_kitti_export 78 | ): 79 | """helper function to draw bbox on original img and get kitti label on single image. 80 | 81 | Note: img will be modified in-place. 82 | """ 83 | kitti_txt = "" 84 | draw = ImageDraw.Draw(img) 85 | color_list = ['Black', 'Red', 'Blue', 'Gold', 'Purple'] 86 | for i in y_decoded: 87 | if float(i[1]) < self.threshold: 88 | continue 89 | 90 | if self.keep_aspect_ratio: 91 | i[2:6] *= img_ratio 92 | else: 93 | orig_w, orig_h = img.size 94 | ratio_w = float(orig_w) / self.model_input_width 95 | ratio_h = float(orig_h) / self.model_input_height 96 | i[2] *= ratio_w 97 | i[3] *= ratio_h 98 | i[4] *= ratio_w 99 | i[5] *= ratio_h 100 | 101 | if is_kitti_export: 102 | kitti_txt += self.class_mapping[int(i[0])] + ' 0 0 0 ' + ' '.join([str(x) for x in i[2:6]])+' 0 0 0 0 0 0 0 ' + str(i[1])+'\n' 103 | 104 | if is_draw_img: 105 | draw.rectangle( 106 | ((i[2], i[3]), (i[4], i[5])), 107 | outline=color_list[int(i[0]) % len(color_list)] 108 | ) 109 | # txt pad 110 | draw.rectangle(((i[2], i[3]), (i[2] + 100, i[3]+10)), 111 | fill=color_list[int(i[0]) % len(color_list)]) 112 | 113 | draw.text((i[2], i[3]), "{0}: {1:.2f}".format(self.class_mapping[int(i[0])], i[1])) 114 | 115 | 116 | return img, kitti_txt 117 | 118 | 119 | def apply(self, results, this_id, render=True, batching=True): 120 | """Apply the post processor to the outputs to the retinanet outputs.""" 121 | 122 | output_array = [] 123 | 124 | for output_name in self.output_names: 125 | if output_name == "NMS_1": 126 | nms_1_bytes = results.as_numpy(output_name)[0].tobytes() 127 | nms_1 = np.frombuffer(nms_1_bytes, dtype=np.int32) 128 | output_array.append(nms_1) 129 | else: 130 | output_array.append(results.as_numpy(output_name)) 131 | 132 | #..and return results up to the actual batch size. 133 | #y_pred = [i.reshape(max_batch_size, -1)[:actual_batch_size] for i in output_array] 134 | y_pred = [i.reshape(1, -1)[:1] for i in output_array] 135 | print("y_pred is {}".format(y_pred)) 136 | 137 | y_pred_decoded = trt_output_process_fn(self, y_pred) 138 | 139 | 140 | for image_idx in range(self.batch_size): 141 | current_idx = (int(this_id) - 1) * self.batch_size + image_idx 142 | if current_idx >= len(self.frames): 143 | break 144 | current_frame = self.frames[current_idx] 145 | filename = os.path.basename(current_frame._image_path) 146 | 147 | img = Image.open(current_frame._image_path) 148 | orig_w, orig_h = img.size 149 | ratio = min(current_frame.w/float(orig_w), current_frame.h/float(orig_h)) 150 | new_w = int(round(orig_w*ratio)) 151 | ratio = float(orig_w)/new_w 152 | 153 | output_label_file = os.path.join( 154 | self.output_path, "infer_labels", 155 | "{}.txt".format(os.path.splitext(filename)[0]) 156 | ) 157 | output_image_file = os.path.join( 158 | self.output_path, "infer_images", 159 | "{}.jpg".format(os.path.splitext(filename)[0]) 160 | ) 161 | if not os.path.exists(os.path.dirname(output_label_file)): 162 | os.makedirs(os.path.dirname(output_label_file)) 163 | if not os.path.exists(os.path.dirname(output_image_file)): 164 | os.makedirs(os.path.dirname(output_image_file)) 165 | 166 | img, kitti_txt = self._get_bbox_and_kitti_label_single_img(img, ratio, y_pred_decoded[0], output_image_file, output_label_file) 167 | 168 | img.save(output_image_file) 169 | 170 | open(output_label_file, 'w').write(kitti_txt) 171 | 172 | -------------------------------------------------------------------------------- /tao_triton/python/postprocessing/utils.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | from contextlib import contextmanager 23 | from copy import deepcopy 24 | import logging 25 | from multiprocessing import Pool 26 | import os 27 | 28 | import numpy as np 29 | from sklearn.cluster import DBSCAN as dbscan 30 | from PIL import ImageDraw 31 | 32 | from tao_triton.python.types import KittiBbox 33 | from tao_triton.python.utils.kitti import write_kitti_annotation 34 | import tritonclient.grpc.model_config_pb2 as mc 35 | 36 | logger = logging.getLogger(__name__) 37 | 38 | 39 | @contextmanager 40 | def pool_context(*args, **kwargs): 41 | """Simple wrapper to get pool context with close function.""" 42 | pool = Pool(*args, **kwargs) 43 | try: 44 | yield pool 45 | finally: 46 | pool.terminate() 47 | 48 | 49 | def denormalize_bounding_bboxes( 50 | bbox_array, stride, offset, 51 | bbox_norm, num_classes, 52 | scale_w, scale_h, 53 | data_format, model_shape, frames, 54 | this_id 55 | ): 56 | """Convert bbox from relative coordinates to absolute coordinates.""" 57 | boxes = deepcopy(bbox_array) 58 | if data_format == mc.ModelInput.FORMAT_NCHW: 59 | _, model_height, model_width = model_shape 60 | else: 61 | model_height, model_width, _ = model_shape 62 | scales = np.zeros( 63 | (boxes.shape[0], 4, boxes.shape[2], boxes.shape[3]) 64 | ).astype(np.float32) 65 | for i in range(boxes.shape[0]): 66 | frame = frames[(this_id * boxes.shape[0] + i) % len(frames)] 67 | scales[i, 0, :, :].fill(float(frame.width/model_width)) 68 | scales[i, 1, :, :].fill(float(frame.height/model_height)) 69 | scales[i, 2, :, :].fill(float(frame.width/model_width)) 70 | scales[i, 3, :, :].fill(float(frame.height/model_height)) 71 | scales = np.asarray(scales).astype(np.float32) 72 | target_shape = boxes.shape[-2:] 73 | gc_centers = [ 74 | (np.arange(s) * stride + offset) 75 | for s in target_shape 76 | ] 77 | gc_centers = [s / n for s, n in zip(gc_centers, bbox_norm)] 78 | for n in range(num_classes): 79 | boxes[:, 4*n+0, :, :] -= gc_centers[0][:, np.newaxis] * scale_w 80 | boxes[:, 4*n+1, :, :] -= gc_centers[1] * scale_h 81 | boxes[:, 4*n+2, :, :] += gc_centers[0][:, np.newaxis] * scale_w 82 | boxes[:, 4*n+3, :, :] += gc_centers[1] * scale_h 83 | boxes[:, 4*n+0, :, :] *= -bbox_norm[0] 84 | boxes[:, 4*n+1, :, :] *= -bbox_norm[1] 85 | boxes[:, 4*n+2, :, :] *= bbox_norm[0] 86 | boxes[:, 4*n+3, :, :] *= bbox_norm[1] 87 | # Scale back boxes. 88 | boxes[:, 4*n+0, :, :] = np.minimum( 89 | np.maximum(boxes[:, 4*n+0, :, :], 0), 90 | model_width 91 | ) * scales[:, 0, :, :] 92 | boxes[:, 4*n+1, :, :] = np.minimum( 93 | np.maximum(boxes[:, 4*n+1, :, :], 0), 94 | model_height 95 | ) * scales[:, 1, :, :] 96 | boxes[:, 4*n+2, :, :] = np.minimum( 97 | np.maximum(boxes[:, 4*n+2, :, :], 0), 98 | model_width 99 | ) * scales[:, 2, :, :] 100 | boxes[:, 4*n+3, :, :] = np.minimum( 101 | np.maximum(boxes[:, 4*n+3, :, :], 0), 102 | model_height 103 | ) * scales[:, 3, :, :] 104 | return boxes 105 | 106 | 107 | def thresholded_indices(cov_array, num_classes, classes, cov_threshold): 108 | """Threshold out valid bboxes and extract the indices per class.""" 109 | valid_indices = [] 110 | batch_size, num_classes, _, _ = cov_array.shape 111 | for image_idx in range(batch_size): 112 | indices_per_class = [] 113 | for class_idx in range(num_classes): 114 | covs = cov_array[image_idx, class_idx, :, :].flatten() 115 | class_indices = covs > cov_threshold[classes[class_idx]] 116 | indices_per_class.append(class_indices) 117 | valid_indices.append(indices_per_class) 118 | return valid_indices 119 | 120 | 121 | def render_image(frame, image_wise_bboxes, output_image_file, box_color, linewidth=3): 122 | """Render images with overlain outputs.""" 123 | image = frame.load_image() 124 | draw = ImageDraw.Draw(image) 125 | for annotations in image_wise_bboxes: 126 | class_name = annotations.category 127 | box = annotations.box 128 | outline_color = (box_color[class_name].R, 129 | box_color[class_name].G, 130 | box_color[class_name].B) 131 | if (box[2] - box[0]) >= 0 and (box[3] - box[1]) >= 0: 132 | draw.rectangle(box, outline=outline_color) 133 | for i in range(linewidth): 134 | x1 = max(0, box[0] - i) 135 | y1 = max(0, box[1] - i) 136 | x2 = min(frame.width, box[2] + i) 137 | y2 = min(frame.height, box[3] + i) 138 | draw.rectangle(box, outline=outline_color) 139 | image.save(output_image_file) 140 | 141 | 142 | def iou_vectorized(rects): 143 | """ 144 | Intersection over union among a list of rectangles in LTRB format. 145 | 146 | Args: 147 | rects (np.array) : numpy array of shape (N, 4), LTRB format, assumes L 2: 42 | # multi-class visualisation 43 | # FIXME: Supports 10 classes - make it more flexible and not hard coded 44 | vis_multi = vis[:,:,0] 45 | 46 | #Code for visualising FN/FP (gt!=pred) 47 | if isinstance(gt, torch.Tensor): 48 | gt = gt.detach() 49 | gt = utils.make_grid( 50 | gt, 51 | pad_value=pad_value, 52 | padding=padding 53 | ) 54 | gt = np.array(gt.cpu()).transpose((1,2,0))[:,:,0] 55 | 56 | if vis.shape[2] == 1: 57 | vis = np.stack([vis, vis, vis], axis=-1) 58 | 59 | if num_class>2: 60 | color_coded = np.ones(np.shape(vis)) 61 | for i in range(10): 62 | color_coded[vis_multi==i] = colour_mappings[str(i)] 63 | color_coded = color_coded / 255 64 | color_coded = color_coded.astype(float) 65 | #Code for visualising FN/FP (gt!=pred) 66 | if isinstance(gt, np.ndarray): 67 | color_coded_mismatch = np.copy(color_coded) 68 | color_coded_mismatch[vis_multi!=gt] = (0,0,0) 69 | color_coded_mismatch = color_coded_mismatch.astype(float) 70 | return color_coded, color_coded_mismatch 71 | return color_coded 72 | else: 73 | return vis 74 | 75 | 76 | def _visualize_pred(pred): 77 | """Return visualized predictions.""" 78 | pred = torch.Tensor(pred) 79 | pred = torch.argmax(pred, dim=1, keepdims=True) 80 | pred_vis = pred * 255 81 | return pred_vis 82 | 83 | 84 | def de_norm(tensor_data): 85 | """Denormalize a torch tensor.""" 86 | tensor_data = tensor_data * 0.5 + 0.5 87 | tensor_data = torch.Tensor(tensor_data) 88 | return tensor_data 89 | 90 | 91 | class VisualChangeNetPostprocessor(Postprocessor): 92 | """Class to run post processing of Triton Tensors.""" 93 | 94 | def __init__(self, batch_size, frames, output_path, data_format): 95 | """Initialize a post processor class for a visual_changenet model. 96 | 97 | Args: 98 | batch_size (int): Number of images in the batch. 99 | frames (list): List of images. 100 | output_path (str): Unix path to the output rendered images and labels. 101 | data_format (str): Order of the input model dimensions. 102 | "channels_first": CHW order. 103 | "channels_last": HWC order. 104 | """ 105 | super().__init__(batch_size, frames, output_path, data_format) 106 | self.output_names = [ 107 | "output0", 108 | "output1", 109 | 'output2', 110 | 'output3', 111 | 'output_final' 112 | ] 113 | self.final_output = "output_final" 114 | 115 | def apply(self, output_tensors, this_id, render=True, batching=True, img_name=None): 116 | """Apply the post processor to the outputs to the visual_changenet outputs.""" 117 | output_cf = output_tensors.as_numpy(self.final_output) 118 | vis = make_numpy_grid(_visualize_pred(output_cf)) 119 | vis = np.clip(vis, a_min=0.0, a_max=1.0) 120 | file_name = os.path.join( 121 | self.output_path, f'eval_{str(img_name)}.jpg' 122 | ) 123 | plt.imsave(file_name, vis) 124 | -------------------------------------------------------------------------------- /tao_triton/python/postprocessing/yolov3_postprocessor.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Simple class to run post processing of Triton Inference outputs.""" 23 | 24 | import os 25 | import numpy as np 26 | from PIL import Image, ImageDraw 27 | 28 | from tao_triton.python.postprocessing.postprocessor import Postprocessor 29 | from tao_triton.python.postprocessing.utils import pool_context 30 | 31 | 32 | 33 | def trt_output_process_fn(y_encoded): 34 | "function to process TRT model output." 35 | keep_k, boxes, scores, cls_id = y_encoded 36 | result = [] 37 | for idx, k in enumerate(keep_k.reshape(-1)): 38 | mul = np.array([960, 39 | 544, 40 | 960, 41 | 544]) 42 | loc = boxes[idx].reshape(-1, 4)[:k] * mul 43 | cid = cls_id[idx].reshape(-1, 1)[:k] 44 | conf = scores[idx].reshape(-1, 1)[:k] 45 | result.append(np.concatenate((cid, conf, loc), axis=-1)) 46 | return result 47 | 48 | 49 | class YOLOv3Postprocessor(Postprocessor): 50 | """Class to run post processing of Triton Tensors.""" 51 | 52 | def __init__(self, batch_size, frames, output_path, data_format): 53 | """Initialize a post processor class for a yolov3 model. 54 | 55 | Args: 56 | batch_size (int): Number of images in the batch. 57 | frames (list): List of images. 58 | output_path (str): Unix path to the output rendered images and labels. 59 | data_format (str): Order of the input model dimensions. 60 | "channels_first": CHW order. 61 | "channels_last": HWC order. 62 | """ 63 | super().__init__(batch_size, frames, output_path, data_format) 64 | self.output_names = ["BatchedNMS", 65 | "BatchedNMS_1", 66 | "BatchedNMS_2", 67 | "BatchedNMS_3"] 68 | self.threshold = 0.8 69 | self.keep_aspect_ratio = True 70 | self.class_mapping = {0: 'bicycle', 1: 'car', 2: 'person', 3: 'road_sign'} 71 | 72 | def _get_bbox_and_kitti_label_single_img( 73 | self, img, img_ratio, y_decoded, 74 | is_draw_img, is_kitti_export 75 | ): 76 | """helper function to draw bbox on original img and get kitti label on single image. 77 | 78 | Note: img will be modified in-place. 79 | """ 80 | kitti_txt = "" 81 | draw = ImageDraw.Draw(img) 82 | color_list = ['Black', 'Red', 'Blue', 'Gold', 'Purple'] 83 | for i in y_decoded: 84 | if float(i[1]) < self.threshold: 85 | continue 86 | 87 | if self.keep_aspect_ratio: 88 | i[2:6] *= img_ratio 89 | else: 90 | orig_w, orig_h = img.size 91 | ratio_w = float(orig_w) / self.model_input_width 92 | ratio_h = float(orig_h) / self.model_input_height 93 | i[2] *= ratio_w 94 | i[3] *= ratio_h 95 | i[4] *= ratio_w 96 | i[5] *= ratio_h 97 | 98 | if is_kitti_export: 99 | kitti_txt += self.class_mapping[int(i[0])] + ' 0 0 0 ' + \ 100 | ' '.join([str(x) for x in i[2:6]])+' 0 0 0 0 0 0 0 ' + str(i[1])+'\n' 101 | 102 | if is_draw_img: 103 | draw.rectangle( 104 | ((i[2], i[3]), (i[4], i[5])), 105 | outline=color_list[int(i[0]) % len(color_list)] 106 | ) 107 | # txt pad 108 | draw.rectangle(((i[2], i[3]), (i[2] + 100, i[3]+10)), 109 | fill=color_list[int(i[0]) % len(color_list)]) 110 | 111 | draw.text((i[2], i[3]), "{0}: {1:.2f}".format(self.class_mapping[int(i[0])], i[1])) 112 | 113 | 114 | return img, kitti_txt 115 | 116 | 117 | def apply(self, results, this_id, render=True, batching=True): 118 | """Apply the post processor to the outputs to the yolov3 outputs.""" 119 | 120 | #output_array = {} 121 | output_array = [] 122 | 123 | for output_name in self.output_names: 124 | #print(results.as_numpy(output_name)) 125 | output_array.append(results.as_numpy(output_name)) 126 | 127 | #y_pred = [i.reshape(max_batch_size, -1)[:actual_batch_size] for i in output_array] 128 | y_pred = [i.reshape(1, -1)[:1] for i in output_array] 129 | 130 | y_pred_decoded = trt_output_process_fn(y_pred) 131 | 132 | 133 | for image_idx in range(self.batch_size): 134 | current_idx = (int(this_id) - 1) * self.batch_size + image_idx 135 | if current_idx >= len(self.frames): 136 | break 137 | current_frame = self.frames[current_idx] 138 | filename = os.path.basename(current_frame._image_path) 139 | 140 | img = Image.open(current_frame._image_path) 141 | orig_w, orig_h = img.size 142 | ratio = min(current_frame.w/float(orig_w), current_frame.h/float(orig_h)) 143 | new_w = int(round(orig_w*ratio)) 144 | ratio = float(orig_w)/new_w 145 | 146 | output_label_file = os.path.join( 147 | self.output_path, "infer_labels", 148 | "{}.txt".format(os.path.splitext(filename)[0]) 149 | ) 150 | output_image_file = os.path.join( 151 | self.output_path, "infer_images", 152 | "{}.jpg".format(os.path.splitext(filename)[0]) 153 | ) 154 | if not os.path.exists(os.path.dirname(output_label_file)): 155 | os.makedirs(os.path.dirname(output_label_file)) 156 | if not os.path.exists(os.path.dirname(output_image_file)): 157 | os.makedirs(os.path.dirname(output_image_file)) 158 | 159 | img, kitti_txt = self._get_bbox_and_kitti_label_single_img(img, ratio, y_pred_decoded[0], output_image_file, output_label_file) 160 | 161 | img.save(output_image_file) 162 | 163 | open(output_label_file, 'w').write(kitti_txt) 164 | 165 | -------------------------------------------------------------------------------- /tao_triton/python/proto/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps/9a30f9692bf29fb728520e9dba1c79be2bf65e74/tao_triton/python/proto/__init__.py -------------------------------------------------------------------------------- /tao_triton/python/proto/postprocessor_config.proto: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved. 2 | 3 | /** 4 | * postprocessing_config.proto: Protocol buffer definition for configuring the post-processing 5 | * block of the DetectNet-v2 model. 6 | */ 7 | 8 | syntax = "proto3"; 9 | 10 | message DBSCANConfig{ 11 | // DBSCAN parameters 12 | float dbscan_eps = 1; 13 | float dbscan_min_samples = 2; 14 | int32 neighborhood_size = 3; 15 | float dbscan_confidence_threshold = 4; 16 | } 17 | 18 | message ClusteringConfig { 19 | // Grid cells with coverage lower than this threshold will be ignored 20 | float coverage_threshold = 1; 21 | int32 minimum_bounding_box_height = 2; 22 | DBSCANConfig dbscan_config = 3; 23 | message BboxColor { 24 | int32 R = 1; 25 | int32 G = 2; 26 | int32 B = 3; 27 | } 28 | BboxColor bbox_color = 4; 29 | } 30 | 31 | // Create a dictionary of post processing config. 32 | message PostprocessingConfig { 33 | // Post processing config. 34 | map classwise_clustering_config = 1; 35 | int32 linewidth = 2; 36 | int32 stride = 3; 37 | } 38 | 39 | // Create a dictionary of post processing config for CenterPose. 40 | message CenterPoseConfig { 41 | // CenterPose post processing config, intrinsic matrix. 42 | float visualization_threshold = 1; 43 | float principle_point_x = 2; 44 | float principle_point_y = 3; 45 | float focal_length_x = 4; 46 | float focal_length_y = 5; 47 | float skew = 6; 48 | // CenterPose visualization parameters. 49 | float axis_size = 7; 50 | int32 square_size = 8; 51 | int32 line_weight = 9; 52 | float scale_text = 10; 53 | } -------------------------------------------------------------------------------- /tao_triton/python/types/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Data structures in Triton Samples.""" 23 | 24 | from tao_triton.python.types.annotation import * 25 | from tao_triton.python.types.frame import * 26 | from tao_triton.python.types.user_data import * 27 | 28 | __all__ = ('BaseAnnotation', "Frame", "KittiBbox", "UserData") -------------------------------------------------------------------------------- /tao_triton/python/types/annotation.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Definition for bbox annotations.""" 23 | 24 | import csv 25 | import logging 26 | import os 27 | 28 | logger = logging.getLogger(__name__) 29 | 30 | 31 | class BaseAnnotation(object): 32 | """Label annotation object.""" 33 | 34 | def __init__(self): 35 | """Initialze an annotation object.""" 36 | self.initialized = True 37 | 38 | def __str__(self): 39 | """String representation of the annotation object.""" 40 | raise NotImplementedError("This method is not implemented in the base class.") 41 | 42 | 43 | class KittiBbox(BaseAnnotation): 44 | """Label annotation for a kitti object.""" 45 | 46 | def __init__(self, category, truncation, 47 | occlusion, observation_angle, 48 | box, height, width, length, 49 | x, y, z, world_bbox_rot_y, 50 | confidence_score=None): 51 | """Initialize a kitti annotation object.""" 52 | self.category = category 53 | self.truncation = float(truncation) 54 | self.observation_angle = float(observation_angle) 55 | self.occlusion = int(occlusion) 56 | self.box = [float(x) for x in box] 57 | hwlxyz = [float(height), float(width), float(length), 58 | float(x), float(y), float(z)] 59 | self.world_bbox = hwlxyz[3:6] + hwlxyz[0:3] 60 | self.world_bbox_rot_y = world_bbox_rot_y 61 | self.confidence = confidence_score 62 | super(KittiBbox, self).__init__() 63 | 64 | def __str__(self): 65 | """String representation of the label file.""" 66 | assert self.initialized, ("Annotation should be initialized.") 67 | world_bbox_str = "{3:.2f} {4:.2f} {5:.2f} {0:.2f} {1:.2f} {2:.2f}".format( 68 | *self.world_bbox 69 | ) 70 | bbox_str = "{:0.3f} {:0.3f} {:0.3f} {:0.3f}".format(*self.box) 71 | if self.confidence is not None: 72 | return "{0} {1:.2f} {2} {3:0.2f} {4} {5} {6:.2f} {7:0.2f}".format( 73 | self.category, self.truncation, self.occlusion, self.observation_angle, 74 | bbox_str, world_bbox_str, self.world_bbox_rot_y, self.confidence 75 | ) 76 | return "{0} {1:.2f} {2} {3:0.2f} {4} {5} {6:.2f}".format( 77 | self.category, self.truncation, self.occlusion, self.observation_angle, 78 | bbox_str, world_bbox_str, self.world_bbox_rot_y 79 | ) 80 | 81 | -------------------------------------------------------------------------------- /tao_triton/python/types/user_data.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """User data requests.""" 23 | 24 | import sys 25 | 26 | if sys.version_info >= (3, 0): 27 | import queue 28 | else: 29 | import Queue as queue 30 | 31 | class UserData: 32 | """Data structure to gather queued requests.""" 33 | 34 | def __init__(self): 35 | self._completed_requests = queue.Queue() -------------------------------------------------------------------------------- /tao_triton/python/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Utitilies for the TLT-Triton.""" -------------------------------------------------------------------------------- /tao_triton/python/utils/kitti.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # Permission is hereby granted, free of charge, to any person obtaining 4 | # a copy of this software and associated documentation files (the 5 | # "Software"), to deal in the Software without restriction, including 6 | # without limitation the rights to use, copy, modify, merge, publish, 7 | # distribute, sublicense, and/or sell copies of the Software, and to 8 | # permit persons to whom the Software is furnished to do so, subject to 9 | # the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be 12 | # included in all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | """Utitilies to handle KITTI label file.""" 23 | 24 | import os 25 | 26 | from tao_triton.python.types import KittiBbox 27 | 28 | 29 | def write_kitti_annotation(label_file, objects): 30 | """Write a kitti annotation file.""" 31 | if not os.path.exists(os.path.dirname(label_file)): 32 | raise NotFoundError("Label file cannot be written to dir: {}".format( 33 | os.path.dirname(label_file)) 34 | ) 35 | assert isinstance(objects, list), ( 36 | "The annotation must be a list of objects.""" 37 | ) 38 | with open(label_file, "w") as lfile: 39 | for label in objects: 40 | if not isinstance(label, KittiBbox): 41 | raise NotImplementedError("Cannot serialize label object") 42 | lfile.write("{}\n".format(str(label))) 43 | return lfile.closed -------------------------------------------------------------------------------- /tao_triton/python/utils/preprocess_input.py: -------------------------------------------------------------------------------- 1 | """Utilities for ImageNet data preprocessing & prediction decoding.""" 2 | 3 | from __future__ import absolute_import 4 | from __future__ import division 5 | from __future__ import print_function 6 | 7 | import logging 8 | import numpy as np 9 | logger = logging.getLogger(__name__) 10 | 11 | 12 | def _preprocess_numpy_input(x, data_format, mode, color_mode, img_mean, **kwargs): 13 | """Preprocesses a Numpy array encoding a batch of images. 14 | 15 | # Arguments 16 | x: Input array, 3D or 4D. 17 | data_format: Data format of the image array. 18 | mode: One of "caffe", "tf" or "torch". 19 | - caffe: will convert the images from RGB to BGR, 20 | then will zero-center each color channel with 21 | respect to the ImageNet dataset, 22 | without scaling. 23 | - tf: will scale pixels between -1 and 1, 24 | sample-wise. 25 | - torch: will scale pixels between 0 and 1 and then 26 | will normalize each channel with respect to the 27 | ImageNet dataset. 28 | 29 | # Returns 30 | Preprocessed Numpy array. 31 | """ 32 | if not issubclass(x.dtype.type, np.floating): 33 | #x = x.astype(backend.floatx(), copy=False) 34 | x = x.astype(float32, copy=False) 35 | 36 | if mode == 'tf': 37 | if img_mean and len(img_mean) > 0: 38 | logger.debug("image_mean is ignored in tf mode.") 39 | x /= 127.5 40 | x -= 1. 41 | return x 42 | 43 | if mode == 'torch': 44 | if img_mean and len(img_mean) > 0: 45 | logger.debug("image_mean is ignored in torch mode.") 46 | x /= 255. 47 | if color_mode == "rgb": 48 | mean = [0.485, 0.456, 0.406] 49 | std = [0.224, 0.224, 0.224] 50 | elif color_mode == "grayscale": 51 | mean = [0.449] 52 | std = [0.224] 53 | else: 54 | raise NotImplementedError("Invalid color mode: {}".format(color_mode)) 55 | else: 56 | if color_mode == "rgb": 57 | if data_format == 'channels_first': 58 | # 'RGB'->'BGR' 59 | if x.ndim == 3: 60 | x = x[::-1, ...] 61 | else: 62 | x = x[:, ::-1, ...] 63 | else: 64 | # 'RGB'->'BGR' 65 | x = x[..., ::-1] 66 | if not img_mean: 67 | mean = [103.939, 116.779, 123.68] 68 | else: 69 | assert len(img_mean) == 3, "image_mean must be a list of 3 values \ 70 | for RGB input." 71 | mean = img_mean 72 | std = None 73 | else: 74 | if not img_mean: 75 | mean = [117.3786] 76 | else: 77 | assert len(img_mean) == 1, "image_mean must be a list of a single value \ 78 | for gray image input." 79 | mean = img_mean 80 | std = None 81 | 82 | # Zero-center by mean pixel 83 | if data_format == 'channels_first': 84 | for idx in range(len(mean)): 85 | if x.ndim == 3: 86 | x[idx, :, :] -= mean[idx] 87 | if std is not None: 88 | x[idx, :, :] /= std[idx] 89 | else: 90 | x[:, idx, :, :] -= mean[idx] 91 | if std is not None: 92 | x[:, idx, :, :] /= std[idx] 93 | else: 94 | for idx in range(len(mean)): 95 | x[..., idx] -= mean[idx] 96 | if std is not None: 97 | x[..., idx] /= std[idx] 98 | return x 99 | 100 | 101 | 102 | def preprocess_input(x, data_format=None, mode='caffe', color_mode="rgb", img_mean=None, **kwargs): 103 | """Preprocesses a tensor or Numpy array encoding a batch of images. 104 | 105 | # Arguments 106 | x: Input Numpy or symbolic tensor, 3D or 4D. 107 | The preprocessed data is written over the input data 108 | if the data types are compatible. To avoid this 109 | behaviour, `numpy.copy(x)` can be used. 110 | data_format: Data format of the image tensor/array. 111 | mode: One of "caffe", "tf" or "torch". 112 | - caffe: will convert the images from RGB to BGR, 113 | then will zero-center each color channel with 114 | respect to the ImageNet dataset, 115 | without scaling. 116 | - tf: will scale pixels between -1 and 1, 117 | sample-wise. 118 | - torch: will scale pixels between 0 and 1 and then 119 | will normalize each channel with respect to the 120 | ImageNet dataset. 121 | 122 | # Returns 123 | Preprocessed tensor or Numpy array. 124 | 125 | # Raises 126 | ValueError: In case of unknown `data_format` argument. 127 | """ 128 | data_format = "channels_first" 129 | 130 | return _preprocess_numpy_input(x, data_format=data_format, 131 | mode=mode, color_mode=color_mode, 132 | img_mean=img_mean, **kwargs) 133 | -------------------------------------------------------------------------------- /tao_triton/requirements-pip.txt: -------------------------------------------------------------------------------- 1 | argon2-cffi==20.1.0 2 | astroid==2.5.6 3 | async-generator==1.10 4 | attrdict==2.0.1 5 | attrs==20.3.0 6 | backcall==0.2.0 7 | bleach==3.3.0 8 | certifi==2020.12.5 9 | cffi==1.14.5 10 | chardet<4,>=3.0.2 11 | decorator==5.0.7 12 | defusedxml==0.7.1 13 | entrypoints==0.3 14 | fvcore==0.1.5.post20221221 15 | gevent==21.1.2 16 | geventhttpclient==1.4.4 17 | greenlet==1.0.0 18 | grpcio==1.37.0 19 | importlib-metadata==4.0.1 20 | ipykernel==5.5.3 21 | ipython==7.2.0 22 | ipython-genutils==0.2.0 23 | ipywidgets==7.6.3 24 | isort==5.8.0 25 | jedi==0.17.0 26 | Jinja2==2.11.3 27 | joblib==1.0.1 28 | jsonschema==3.2.0 29 | jupyter==1.0.0 30 | jupyter-client==6.1.12 31 | jupyter-console==6.4.0 32 | jupyter-core==4.7.1 33 | jupyterlab-pygments==0.1.2 34 | jupyterlab-widgets==1.0.0 35 | kornia==0.7.2 36 | lazy-object-proxy==1.6.0 37 | MarkupSafe==1.1.1 38 | matplotlib==3.5.1 39 | mccabe==0.6.1 40 | mistune==0.8.4 41 | ninja==1.11.1.1 42 | nbclient==0.5.3 43 | nbconvert==6.0.7 44 | nbformat==5.1.3 45 | nest-asyncio==1.5.1 46 | notebook==6.4.1 47 | numpy==1.19.5 48 | open3d==0.15.2 49 | opencv-python==4.5.5.64 50 | packaging==20.9 51 | pandocfilters==1.4.3 52 | parso==0.8.2 53 | pexpect==4.8.0 54 | pickleshare==0.7.5 55 | Pillow==8.3.2 56 | prometheus-client==0.10.1 57 | prompt-toolkit==2.0.10 58 | protobuf==3.15.8 59 | ptyprocess==0.7.0 60 | pycparser==2.20 61 | Pygments==2.8.1 62 | pylint==2.8.2 63 | pyparsing==2.4.7 64 | pyrr==0.10.3 65 | pyrsistent==0.17.3 66 | python-dateutil==2.8.1 67 | python-rapidjson==1.0 68 | pyyaml==6.0 69 | pyzmq==22.0.3 70 | qtconsole==5.0.3 71 | QtPy==1.9.0 72 | scikit-image==0.19.2 73 | scikit-learn==0.24.2 74 | scipy==1.5.4 75 | Send2Trash==1.5.0 76 | six==1.15.0 77 | terminado==0.9.4 78 | testpath==0.4.4 79 | threadpoolctl==2.1.0 80 | toml==0.10.2 81 | torch==1.10.2 82 | torchvision==0.11.3 83 | tornado==6.1 84 | tqdm==4.60.0 85 | traitlets==4.3.3 86 | transformations==2022.9.26 87 | trimesh==4.0.4 88 | typed-ast==1.4.3 89 | typing-extensions==4.10.0 90 | warp-lang==0.10.1 91 | wcwidth==0.2.5 92 | webencodings==0.5.1 93 | widgetsnbextension==3.5.1 94 | wrapt==1.12.1 95 | zipp==3.4.1 96 | zope.event==4.5.0 97 | zope.interface==5.4.0 98 | git+https://github.com/NVlabs/nvdiffrast.git 99 | git+https://github.com/facebookresearch/pytorch3d.git@stable -------------------------------------------------------------------------------- /tao_triton/setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps/9a30f9692bf29fb728520e9dba1c79be2bf65e74/tao_triton/setup.py --------------------------------------------------------------------------------