[
{
"path": ".gitignore",
"content": "# Project particular ignores\nbackup/\nvideos/inputs/*\nvideos/outputs/*\nimages*/\n*.mp4\n.DS_Store\n.vscode/\n\n# Intermediate files to convert MobileNetV2 from TF->TF-TRT\nutils/*.png\nmobilenetv2/inference_graph_*/\nmobilenetv2/converted_trt_*.pb\n\n# Avoid weights, or engines\n*.weights\n*.trt\n*.engine\n*.onnx\n\n# Environment files (not templates)\nserver/*.env\n\n# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n*.o\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\npip-wheel-metadata/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n.python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n"
},
{
"path": "Dockerfile",
"content": "# Installs maskcam on a BalenaOS container (devkit or Photon)\nFROM balenalib/jetson-nano-ubuntu:20210201\n\n# Don't prompt with any configuration questions\nENV DEBIAN_FRONTEND noninteractive\n\n# Switch the nvidia apt source repos and\n# install some utilities\n\nRUN \\\n apt-get update && apt-get install -y \\\n lbzip2 wget tar python3 git\n\nENV UDEV=1\n\n# Download and install BSP binaries for L4T 32.4.2\n# This is mostly from Balena's Alan Boris at:\n# https://github.com/balena-io-playground/jetson-nano-sample-new/blob/master/CUDA/Dockerfile\n\nRUN apt-get update && apt-get install -y wget tar python3 libegl1 && \\\n wget https://developer.nvidia.com/embedded/L4T/r32_Release_v4.2/t210ref_release_aarch64/Tegra210_Linux_R32.4.2_aarch64.tbz2 && \\\n tar xf Tegra210_Linux_R32.4.2_aarch64.tbz2 && \\\n cd Linux_for_Tegra && \\\n sed -i 's/config.tbz2\\\"/config.tbz2\\\" --exclude=etc\\/hosts --exclude=etc\\/hostname/g' apply_binaries.sh && \\\n sed -i 's/install --owner=root --group=root \\\"${QEMU_BIN}\\\" \\\"${L4T_ROOTFS_DIR}\\/usr\\/bin\\/\\\"/#install --owner=root --group=root \\\"${QEMU_BIN}\\\" \\\"${L4T_ROOTFS_DIR}\\/usr\\/bin\\/\\\"/g' nv_tegra/nv-apply-debs.sh && \\\n sed -i 's/LC_ALL=C chroot . mount -t proc none \\/proc/ /g' nv_tegra/nv-apply-debs.sh && \\\n sed -i 's/umount ${L4T_ROOTFS_DIR}\\/proc/ /g' nv_tegra/nv-apply-debs.sh && \\\n sed -i 's/chroot . \\// /g' nv_tegra/nv-apply-debs.sh && \\\n ./apply_binaries.sh -r / --target-overlay && cd .. && \\\n rm -rf Tegra210_Linux_R32.4.2_aarch64.tbz2 && \\\n rm -rf Linux_for_Tegra && \\\n echo \"/usr/lib/aarch64-linux-gnu/tegra\" > /etc/ld.so.conf.d/nvidia-tegra.conf && \\\n echo \"/usr/lib/aarch64-linux-gnu/tegra-egl\" > /etc/ld.so.conf.d/nvidia-tegra-egl.conf && ldconfig\n\n# Install GStreamer and remove unnecessary files\nRUN apt-get install -y \\\n libssl1.0.0 \\\n libgstreamer1.0-0 \\\n gstreamer1.0-tools \\\n gstreamer1.0-plugins-good \\\n gstreamer1.0-plugins-bad \\\n gstreamer1.0-plugins-ugly \\\n gstreamer1.0-libav \\\n libgstrtspserver-1.0-0 \\\n libjansson4=2.11-1 \\\n cuda-toolkit-10-2 && \\\n ldconfig\nRUN \\\n rm -rf /usr/src/nvidia/graphics_demos \\\n /usr/local/cuda-10.2/samples \\\n /usr/local/cuda-10.2/doc \n\n# Install DeepStream\nRUN apt-get install -y deepstream-5.0 && \\\n rm -rf /opt/nvidia/deepstream/deepstream-5.0/samples \\\n /usr/lib/aarch64-linux-gnu/libcudnn_static_v8.a \\\n /usr/lib/aarch64-linux-gnu/libcudnn_cnn_infer_static_v8.a \\\n /usr/lib/aarch64-linux-gnu/libnvinfer_static.a \\\n /usr/lib/aarch64-linux-gnu/libcudnn_adv_infer_static_v8.a \\\n /usr/lib/aarch64-linux-gnu/libcublas_static.a \\\n /usr/lib/aarch64-linux-gnu/libcudnn_adv_train_static_v8.a \\\n /usr/lib/aarch64-linux-gnu/libcudnn_ops_infer_static_v8.a \\\n /usr/lib/aarch64-linux-gnu/libcublasLt_static.a \\\n /usr/lib/aarch64-linux-gnu/libcudnn_cnn_train_static_v8.a \\\n /usr/lib/aarch64-linux-gnu/libcudnn_ops_train_static_v8.a \\\n /usr/lib/aarch64-linux-gnu/libmyelin_compiler_static.a \\\n /usr/lib/aarch64-linux-gnu/libmyelin_executor_static.a \\\n /usr/lib/aarch64-linux-gnu/libnvinfer_plugin_static.a && \\\n ldconfig\n\n# Install system-level python3 packages\nRUN apt-get update && apt-get install -y \\\n gir1.2-gst-rtsp-server-1.0 \\\n python3-pip \\\n python3-opencv \\\n python3-libnvinfer \\\n python3-scipy \\\n cython3 \\\n python3-sklearn \\\n python-gi-dev \\\n unzip && ldconfig\n\n# These system-level packages don't provide egg-info files, add them manually so that pip knows\nCOPY docker/opencv_python-3.2.0.egg-info /usr/lib/python3/dist-packages/\nCOPY docker/scikit-learn-0.19.1.egg-info /usr/lib/python3/dist-packages/\n\n# Install gst-python (python bindings for GStreamer)\nRUN \\\n export GST_CFLAGS=\"-pthread -I/usr/include/gstreamer-1.0 -I/usr/include/glib-2.0 -I/usr/lib/x86_64-linux-gnu/glib-2.0/include\" && \\\n export GST_LIBS=\"-lgstreamer-1.0 -lgobject-2.0 -lglib-2.0\" && \\\n git clone https://github.com/GStreamer/gst-python.git && \\\n cd gst-python && git checkout 1a8f48a && \\\n ./autogen.sh PYTHON=python3 && \\\n ./configure PYTHON=python3 && \\\n make && make install\n\n# Install pyds (python bindings for DeepStream)\nRUN cd /opt/nvidia/deepstream/deepstream-5.0/lib && python3 setup.py install\n\n# Upgrade here to avoid re-running on code changes\nRUN pip3 install --upgrade pip\n\n# ---- Below steps are run before copying full maskcam code to allow layer caching ----\n\n# Compile YOLOv4 plugin for DeepStream\nCOPY deepstream_plugin_yolov4 /deepstream_plugin_yolov4\nENV CUDA_VER=10.2\nRUN cd /deepstream_plugin_yolov4 && make\n\n# Get TensorRT engine (pretrained YOLOv4-tiny)\n# Model trained on smaller dataset\n# RUN wget -P / https://maskcam.s3.us-east-2.amazonaws.com/facemask_y4tiny_1024_608_fp16.trt\n\n# Model trained on bigger dataset, merged with MAFA, WiderFace, Kaggle Medical Masks and FDDB\nRUN wget -P / https://maskcam.s3.us-east-2.amazonaws.com/maskcam_y4t_1024_608_fp16.trt\n# RUN wget -P / https://maskcam.s3.us-east-2.amazonaws.com/maskcam_y4t_1120_640_fp16.trt\n\n# Install requirements with pinned versions\nCOPY requirements.txt /maskcam_requirements.txt\nRUN pip3 install -r /maskcam_requirements.txt\n\n# ---- Note: all layers below this will be re-generated each time code changes ----\n# Copy full maskcam code\nCOPY . /opt/maskcam_1.0/\nWORKDIR /opt/maskcam_1.0\n\n# Move pre-copied files to their maskcam location\n# NOTE: Ignoring errors with `exit 0` to avoid breaking on balena livepush\nRUN rm -r deepstream_plugin_yolov4 && mv /deepstream_plugin_yolov4 . ; exit 0\nRUN mv /*.trt yolo/ ; exit 0\n\n# Preload library to avoids Gst errors \"cannot allocate memory in static TLS block\"\nENV LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libgomp.so.1\n\n# Un-pinned versions of maskcam requirements (comment pip3 install above before this)\n# RUN pip3 install -r requirements.in -c docker/constraints.docker\n\nRUN chmod +x docker/start.sh\nRUN chmod +x maskcam_run.py\nCMD [\"docker/start.sh\"]\n"
},
{
"path": "LICENSE.md",
"content": "MIT License\nCopyright (c) 2020, 2021 Berkeley Design Technology, Inc.. All rights reserved.\nPermission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:\nThe above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.\nTHE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.\n"
},
{
"path": "README.md",
"content": "# MaskCam \n\n\n ![](\"/docs/imgs/MaskCam-Demo1.gif\")
\n
\n\nMaskCam is a prototype reference design for a Jetson Nano-based smart camera system that measures crowd face mask usage in real-time, with all AI computation performed at the edge. MaskCam detects and tracks people in its field of view and determines whether they are wearing a mask via an object detection, tracking, and voting algorithm. It uploads statistics (not videos) to the cloud, where a web GUI can be used to monitor face mask compliance in the field of view. It saves interesting video snippets to local disk (e.g., a sudden influx of lots of people not wearing masks) and can optionally stream video via RTSP.\n\nMaskCam can be run on a Jetson Nano Developer Kit, or on a Jetson Nano module (SOM) with the ConnectTech Photon carrier board. It was designed to use the Raspberry Pi High Quality Camera but will also work with pretty much any USB webcam that is supported on Linux.\n\nThe on-device software stack is mostly written in Python and runs under JetPack 4.4.1 or 4.5. Edge AI processing is handled by NVIDIA’s DeepStream video analytics framework, YOLOv4-tiny, and Tryolabs' [Norfair](https://github.com/tryolabs/norfair) tracker. MaskCam reports statistics to and receives commands from the cloud using MQTT and a web-based GUI. The software is containerized and for evaluation can be easily installed on a Jetson Nano DevKit using docker with just a couple of commands. For production, MaskCam can run under balenaOS, which makes it easy to manage and deploy multiple devices.\n\nWe urge you to try it out! It’s easy to install on a Jetson Nano Developer Kit and requires only a web cam. (The cloud-based statistics server and web GUI are optional, but are also dockerized and easy to install on any reasonable Linux system.) [See below for installation instructions.](#running-maskcam-from-a-container-on-a-jetson-nano-developer-kit)\n\nMaskCam was developed by Berkeley Design Technology, Inc. (BDTI) and Tryolabs S.A., with development funded by NVIDIA. MaskCam is offered under the MIT License. For more information about MaskCam, please see the [report from BDTI](https://www.bdti.com/maskcam). If you have questions, please email us at maskcam@bdti.com. Thanks!\n\n## Table of contents \n- [Start Here!](#start-here)\n - [Running MaskCam from a Container on a Jetson Nano Developer Kit](#running-maskcam-from-a-container-on-a-jetson-nano-developer-kit)\n - [Viewing the Live Video Stream](#viewing-the-live-video-stream)\n - [Setting Device Configuration Parameters](#setting-device-configuration-parameters)\n- [MQTT Server Setup](#mqtt-and-web-server-setup)\n - [Running the MQTT Broker and Web Server](#running-the-mqtt-broker-and-web-server)\n - [Setup a Device with Your Server](#setup-a-device-with-your-server)\n - [Checking MQTT Connection](#checking-mqtt-connection)\n- [Working With the MaskCam Container](#working-with-the-maskcam-container)\n - [Development Mode: Manually Running MaskCam](#development-mode-manually-running-maskcam)\n - [Debugging: Running MaskCam Modules as Standalone Processes](#debugging-running-maskcam-modules-as-standalone-processes)\n- [Additional Information](#additional-information)\n - [Running on Jetson Nano Developer Kit Using BalenaOS](#running-on-jetson-nano-developer-kit-using-balenaos)\n - [Custom Container Development](#custom-container-development)\n - [Building From Source on Jetson Nano Developer Kit](#building-from-source-on-jetson-nano-developer-kit)\n - [Using Your Own Detection Model](#using-your-own-detection-model)\n - [Installing MaskCam Manually (Without a Container)](#installing-maskcam-manually-without-a-container)\n - [Running on Jetson Nano with Photon Carrier Board](#running-on-jetson-nano-with-photon-carrier-board)\n - [Useful Development Scripts](#useful-development-scripts)\n- [Troubleshooting Common Errors](#troubleshooting-common-errors)\n\n\n## Start Here!\n### Running MaskCam from a Container on a Jetson Nano Developer Kit\nThe easiest and fastest way to get MaskCam running on your Jetson Nano Dev Kit is using our pre-built containers. You will need:\n\n1. A Jetson Nano Dev Kit running JetPack 4.4.1 or 4.5\n2. An external DC 5 volt, 4 amp power supply connected through the Dev Kit's barrel jack connector (J25). (See [these instructions](https://www.jetsonhacks.com/2019/04/10/jetson-nano-use-more-power/) on how to enable barrel jack power.) This software makes full use of the GPU, so it will not run with USB power.\n3. A USB webcam attached to your Nano\n4. Another computer with a program that can display RTSP streams -- we suggest [VLC](https://www.videolan.org/vlc/index.html) or [QuickTime](https://www.apple.com/quicktime/download/).\n\nFirst, the MaskCam container needs to be downloaded from Docker Hub. On your Nano, run:\n```\n# This will take 10 minutes or more to download\nsudo docker pull maskcam/maskcam-beta\n```\n\nFind your local Jetson Nano IP address using `ifconfig`. This address will be used later to view a live video stream from the camera and to interact with the Nano from a web server.\n\nMake sure a USB camera is connected to the Nano, and then start MaskCam by running the following command. Make sure to substitute `` with your Nano's IP address.\n```\n# Connect USB camera before running this!\nsudo docker run --runtime nvidia --privileged --rm -it --env MASKCAM_DEVICE_ADDRESS= -p 1883:1883 -p 8080:8080 -p 8554:8554 maskcam/maskcam-beta\n```\n\nThe MaskCam container should start running the `maskcam_run.py` script, using the USB camera as the default input device (`/dev/video0`). It will produce various status output messages (and error messages, if it encounters problems). If there are errors, the process will automatically end after several seconds. Check the [Troubleshooting](#troubleshooting-common-errors) section for tips on resolving errors.\n\nOtherwise, after 30 seconds or so, it should continually generate status messages (such as `Processed 100 frames...`). Leave it running (don't press `Ctrl+C`, but be aware that the device will start heating up) and continue to the next section to visualize the video!\n\n### Viewing the Live Video Stream\nIf you scroll through the logs and don't see any errors, you should find a message like:\n\n```Streaming at rtsp://aaa.bbb.ccc.ddd:8554/maskcam```\n\nwhere `aaa.bbb.ccc.ddd` is the address that you provided in `MASKCAM_DEVICE_ADDRESS` previously. If you didn't provide an address, you'll see some unknown address label there, but the streaming will still work.\n\nYou can copy-paste that URL into your RSTP streaming viewer ([see how](https://user-images.githubusercontent.com/12506292/111346333-e14d8800-865c-11eb-9242-0ffa4f50547f.mp4) to do it with VLC) on another computer. If all goes well, you should be rewarded with streaming video of your Nano, with green boxes around faces wearing masks and red boxes around faces not wearing masks. An example video of the live streaming in action is shown below.\n\n\n ![](\"/docs/imgs/MaskCam-Live1.gif\")
\n
\n\nThis video stream gives a general demonstration of how MaskCam works. However, MaskCam also has other features, such as the ability to send mask detection statistics to the cloud and view them through a web browser. If you'd like to see these features in action, you'll need to set up an MQTT server, which is covered in the [MQTT Server Setup section](#mqtt-and-web-server-setup).\n\nIf you encounter any errors running the live stream, check the [Troubleshooting](#troubleshooting-common-errors) section for tips on resolving errors.\n\n### Setting Device Configuration Parameters\nMaskCam uses environment variables to configure parameters without having to rebuild the container or manually change the configuration file each time the program is run. For example, in the previous section we set the `MASKCAM_DEVICE_ADDRESS` variable to indicate our Nano's IP address. A list of configurable parameters is shown in [maskcam_config.txt](maskcam_config.txt). The mapping between environment variable names and configuration parameters is defined in [maskcam/config.py](maskcam/config.py).\n\nThis section shows how to set environment variables to change configuration parameters. For example, if you want to use the `/dev/video1` camera device rather than `/dev/video0`, you can define `MASKCAM_INPUT` when running the container:\n\n```\n# Run with MASKCAM_INPUT and MASKCAM_DEVICE_ADDRESS\nsudo docker run --runtime nvidia --privileged --rm -it --env MASKCAM_INPUT=v4l2:///dev/video1 --env MASKCAM_DEVICE_ADDRESS= -p 1883:1883 -p 8080:8080 -p 8554:8554 maskcam/maskcam-beta\n```\n\nAnother useful input device that you might want to use is a CSI camera (like the Raspberry Pi camera), and in that case you need to set `MASKCAM_INPUT=argus://0` instead of the value shown above.\n\nAs another example, if you have an already set up our MQTT and web server (as shown in [MQTT Server Setup section](#mqtt-and-web-server-setup)), you need to define\ntwo addtional environment variables, `MQTT_BROKER_IP` and `MQTT_DEVICE_NAME`. This allows your device to find the MQTT server and identify itself:\n\n```\n# Run with MQTT_BROKER_IP, MQTT_DEVICE_NAME, and MASKCAM_DEVICE_ADDRESS\nsudo docker run --runtime nvidia --privileged --rm -it --env MQTT_BROKER_IP= --env MQTT_DEVICE_NAME= --env MASKCAM_DEVICE_ADDRESS= -p 1883:1883 -p 8080:8080 -p 8554:8554 maskcam/maskcam-beta\n```\n\n*If you have too many `--env` variables to add, it might be easier to create a [.env file](https://docs.docker.com/compose/env-file/) and point to it using the `--env-file` flag instead.*\n\n\n## MQTT and Web Server Setup\n### Running the MQTT Broker and Web Server\nMaskCam is intended to be set up with a web server that stores mask detection statistics and allows users to remotely interact with the device. We wrote code for instantiating a [server](server/) that receives statistics from the device, stores them in a database, and has a web-based GUI frontend to display them. A screenshot of the frontend for an example device is shown below.\n\n\n ![](\"/docs/imgs/maskcam-frontend.PNG\")
\n
\n\nYou can test out and explore this functionality by starting the server on a PC on your local network and pointing your Jetson Nano MaskCam device to it. This section gives instructions on how to do so. The MQTT broker and web server can be built and run on a Linux or OSX machine; we've tested it on Ubuntu 18.04LTS and OSX Big Sur. It can also be set up in an online AWS EC2 instance if you want to access it from outside of your local network.\n\nThe server consists of several docker containers that run together using [docker-compose](https://docs.docker.com/compose/install/). Install docker-compose on your machine by following the [installation instructions for your platform](https://docs.docker.com/compose/install/) before continuing. All other necessary packages and libraries will be automatically installed when you set up the containers in the next steps.\n\nAfter installing docker-compose, clone this repo:\n```\ngit clone https://github.com/bdtinc/maskcam.git\n```\n\nGo to the `server/` folder, which has all the needed components implemented on four containers: the Mosquitto broker, backend API, database, and Streamlit frontend.\n\nThese containers are configured using environment variables, so create the `.env` files by copying the default templates:\n```\ncd server\ncp database.env.template database.env\ncp frontend.env.template frontend.env\ncp backend.env.template backend.env\n```\n\nThe only file that needs to be changed is `database.env`. Open it with a text editor and replace the ``, ``, and `` fields with your own values. Here are some example values, but you better be more creative for security reasons:\n```\nPOSTGRES_USER=postgres\nPOSTGRES_PASSWORD=some_password\nPOSTGRES_DB=maskcam\n```\n\n*NOTE:* If you want to change any of the `database.env` values after building the containers, the easiest thing to do is to delete the `pgdata` volume by running `docker volume rm pgdata`. It will also delete all stored database information and statistics.\n\nAfter editing the database environment file, you're ready to build all the containers and run them with a single command:\n\n```\nsudo docker-compose up -d\n```\n\nWait a couple minutes after issuing the command to make sure that all containers are built and running. Then, check the local IP of your computer by running the `ifconfig` command. (It should be an address that starts with `192.168...`, `10...` or `172...`.) This is the server IP that will be used for connecting to the server (since the server is hosted on this computer).\n\nNext, open a web browser and enter the server IP to visit the frontend webpage:\n```\nhttp://:8501/\n```\nIf you see a `ConnectionError` in the frontend, wait a couple more seconds and reload the page. The backend container can take some time to finish the database setup.\n\n*NOTE:* If you're setting the server up on a remote instance like an AWS EC2, make sure you have ports `1883` (MQTT) and `8501` (web frontend) open for inbound and outbound traffic.\n\n\n### Setup a Device With Your Server\nOnce you've got the server set up on a local machine (or in a AWS EC2 instance with a public IP), switch back to the Jetson Nano device. Run the MaskCam container using the following command, where `MQTT_BROKER_IP` is set to the IP of your server. (If you're using an AWS EC2 server, make sure to configure port `1883` for inbound and outbound traffic before running this command.)\n\n```\n# Run with MQTT_BROKER_IP, MQTT_DEVICE_NAME, and MASKCAM_DEVICE_ADDRESS\nsudo docker run --runtime nvidia --privileged --rm -it --env MQTT_BROKER_IP= --env MQTT_DEVICE_NAME=my-jetson-1 --env MASKCAM_DEVICE_ADDRESS= -p 1883:1883 -p 8080:8080 -p 8554:8554 maskcam/maskcam-beta\n```\n\nAnd that's it. If the device has access to the server's IP, then you should see in the output logs some successful connection messages and then see your device listed in the drop-down menu of the frontend (reload the page if you don't see it). In the frontend, select `Group data by: Second` and hit `Refresh status` to see how the plot changes when new data arrives.\n\nCheck the next section if the MQTT connection is not established from the device to the server.\n\n### Checking MQTT Connection\nIf you're running the MQTT broker on a machine in your local network, make sure its IP is accessible from the Jetson device:\n```\nping \n```\n\n*NOTE:* Remember to use the network address of the computer you set up the server on, which you can check using the `ifconfig` command and looking for an address that should start with `192.168...`, `10...` or `172...`\n\nIf you're setting up a remote server and using its public IP to connect\nfrom your device, chances are you're not setting properly the port `1883` to be opened for inbound and outbound traffic.\nIf you want to check the port is correctly configured, use `nc` from a local machine or your jetson:\n```\nnc -vz 1883\n```\nRemember you also need to open port `8501` to access the web server frontend from a web browser, as explained in the [server configuration section](#running-the-mqtt-broker-and-web-server) (but that's not relevant for the MQTT communication with the device).\n\n\n\n\n## Working With the MaskCam Container\n### Development Mode: Manually Running MaskCam\nIf you want to play around with the code, you probably don't want the container to automatically start running the `maskcam_run.py` script.\nThe easiest way to achieve that, is by defining the environment variable `DEV_MODE=1`:\n```\ndocker run --runtime nvidia --privileged --rm -it --env DEV_MODE=1 -p 1883:1883 -p 8080:8080 -p 8554:8554 maskcam/maskcam-beta\n```\nThis will cause the container to start a `/bin/bash` prompt (see [docker/start.sh](docker/start.sh) for details), from which you could run the script manually, or any\nof its sub-modules as standalone processes:\n\n```\n# e.g: Run with a different input instead of default `/dev/video0`\n./maskcam_run.py v4l2:///dev/video1\n\n# e.g: Disable tracker to visualize raw detections and scores\nMASKCAM_DISABLE_TRACKER=1 ./maskcam_run.py\n```\n\n### Debugging: Running MaskCam Modules as Standalone Processes\nThe script `maskcam_run.py`, which is the main entrypoint for the MaskCam software, has two roles:\n - Handles all the MQTT communication (send stats and receive commands)\n - Orchestrates all other processes that live under `maskcam/maskcam_*.py`.\n\nBut you can actually run any of those modules as standalone processes, which can be easier for debugging.\n\nYou need to set `DEV_MODE=1` as explained in the previous section to access the container prompt, and then you can run the python modules:\n\n```\n# e.g: Run only the static file server process\npython3 -m maskcam.maskcam_fileserver\n# e.g: Serve another directory to test\npython3 -m maskcam.maskcam_fileserver /tmp\n\n# e.g: Run only the inference and streaming processes\npython3 -m maskcam.maskcam_streaming &\n# Hit enter until you get a prompt and then:\npython3 -m maskcam.maskcam_inference\n```\n\n**Note:** In the last example, `maskcam_streaming` is running on background,\nso it will not terminate if you press `Ctrl+C` (only `maskcam_inference` will,\nsince it's running on the foreground).\n\nTo check that the streaming is still running and then bring it to foreground to terminate it, run:\n```\njobs\nfg %1\n# Now you can hit Ctrl+C to terminate streaming\n```\n\n## Additional Information\nFurther information about working with and customizing MaskCam is provided on separate pages in the [docs](docs) folder. This section gives a brief description and link to each page.\n\n### Running on Jetson Nano Developer Kit Using BalenaOS\n[BalenaOS](https://www.balena.io/os/) is a lightweight operating system designed for running containers on embedded devices. It provides several advantages for fleet deployment and management, especially when combined with balena's balenaCloud mangament system. If you'd like to try running MaskCam with balenaOS instead of JetPack OS on your Jetson Nano, please follow the instructions at [BalenaOS-DevKit-Nano-Setup.md](docs/BalenaOS-DevKit-Nano-Setup.md).\n\n### Custom Container Development\nMaskCam is intended to be a reference design for any connected smart camera application. You can create your own application by starting from our pre-built container, modifying it to add the code files and packages needed for your program, and then re-building the container. The [Custom-Container-Development.md](docs/Custom-Container-Development.md) gives instructions on how to build your own container based off MaskCam.\n\n#### Building From Source on Jetson Nano Developer Kit\nPlease see [How to Build your Own Container from Source on the Jetson Nano](https://github.com/bdtinc/maskcam/blob/main/docs/Custom-Container-Development.md#how-to-build-your-own-container-from-source-on-the-jetson-nano) for instructions on how to build a custom MaskCam container on your Jetson Nano Developer Kit.\n\n#### Using Your Own Detection Model\nPlease see [How to Use Your Own Detection Model](https://github.com/bdtinc/maskcam/blob/main/docs/Custom-Container-Development.md#how-to-use-your-own-detection-model) for instructions on how to use your own detection model rather than our mask detection model.\n\n### Installing MaskCam Manually (Without a Container)\nMaskCam can also be installed manually, rather than by downloading our pre-built container. Using a manual installation of MaskCam can help with development if you'd prefer not to work with containers. If you'd like to install MaskCam without using containers, please see [docs/Manual-Dependencies-Installation.md](docs/Manual-Dependencies-Installation.md).\n\n### Running on Jetson Nano with Photon Carrier Board\nFor our hardware prototype of MaskCam, we used a Jetson Nano module and a [Connect Tech Photon carrier board](https://connecttech.com/product/photon-jetson-nano-ai-camera-platform/), rather than the Jetson Nano Developer Kit. We used the Photon because the Developer Kit is not sold or warrantied for production use. Using the Photon allowed us to quickly create a production-ready prototype using off-the-shelf hardware. If you have a Photon carrier board and Jetson Nano module, you can install MaskCam on them by using the setup instructions at [docs/Photon-Nano-Setup.md](docs/Photon-Nano-Setup.md).\n\n### Useful Development Scripts\nDuring development, some scripts were produced which might be useful for other developers to debug or update the software. These include an MQTT sniffer, a script to run the TensorRT model on images, and to convert a model trained with the original YOLO Darknet implementation to TensorRT format. Basic usage for all these tools is covered on [docs/Useful-Development-Scripts.md](docs/Useful-Development-Scripts.md).\n\n\n## Troubleshooting Common Errors\nIf you run into any errors or issues while working with MaskCam, this section gives common errors and their solutions. \n\nMaskCam consists of many different processes running in parallel. As a consequence, when there's an error on a particular process, all of them will be sent termination signals and finish gracefully. This means that you need to scroll up through the output to find out the original error that caused a failure. It should be very notorious, flagged as a red **ERROR** log entry, followed by the name of the process that failed and a message.\n\n#### Error: camera not connected/not recognized\nIf you see an error containing the message `Cannot identify device '/dev/video0'`, among other Gst and v4l messages, it means the program couldn't find the camera device. Make sure your camera is connected to the Nano and recognized by the host Ubuntu OS by issuing `ls /dev` and checking if `/dev/video0` is present in the output.\n\n#### Error: not running in privileged mode\nIn this case, you'll see a bunch of annoying messages like:\n```\nError: Can't initialize nvrm channel\nCouldn't create ddkvic Session: Cannot allocate memory\nnvbuf_utils: Could not create Default NvBufferSession\n```\nYou'll probably see multiple failures in other MaskCam processes as well. To resolve these errors, make sure you're running docker with the `--privileged` flag, as described in the [first section](#running-maskcam-from-a-container-on-a-jetson-nano-developer-kit).\n\n#### Error: reason not negotiated/camera capabilities\nIf you get an error that looks like: `v4l-camera-source / reason not-negotiated`\nThen the problem is that the USB camera you're using doesn't support the default `camera-framerate=30` (frames per second). If you don't have another camera, try running the script under utils/gst_capabilities.sh and find the lines with type `video/x-raw ...`\n\nFind any suitable `framerate=X/1` (with `X` being an integer like 24, 15, etc.) and set the corresponding configuration parameter with `--env MASKCAM_CAMERA_FRAMERATE=X` (see [previous section](#setting-device-configuration-parameters)).\n\n#### Error: Streaming or file server are not accessible (nothing else seems to fail)\nMake sure you're mapping the right ports from the container, with the `-p container_port:host_port` parameters indicated in the previous sections. The default port numbers, that should be exposed by the container, are configured in [maskcam_config.txt](maskcam_config.txt) as:\n```\nfileserver-port=8080\nstreaming-port=8554\nmqtt-broker-port=1883\n```\nThese port mappings are why we use `docker run ... -p 1883:1883 -p 8080:8080 -p 8554:8554 ...` with the run command. Remember that all the ports can be overriden using environment variables, as described in the [previous section](#setting-device-configuration-parameters). Other ports like `udp-port-*` are not intended to be accessible from outside the container, they are used for communication between the inference process and the streaming and file-saving processes.\n\n#### Other Errors\nSometimes after restarting the process or the whole docker container many times, some GPU resources can get stuck and cause unexpected errors. If that's the case, try rebooting the device and running the container again. If you find that the container fails systematically after running some sequence, please don't hesitate to [report an Issue](https://github.com/bdtinc/maskcam/issues) with the relevant context and we'll try to reproduce and fix it.\n\n## Questions? Need Help?\nEmail us at maskcam@bdti.com, and be sure to check out our [independent report on the development of MaskCam](https://bdti.com/maskcam)!\n\n"
},
{
"path": "deepstream_plugin_yolov4/Makefile",
"content": "################################################################################\n# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.\n#\n# Permission is hereby granted, free of charge, to any person obtaining a\n# copy of this software and associated documentation files (the \"Software\"),\n# to deal in the Software without restriction, including without limitation\n# the rights to use, copy, modify, merge, publish, distribute, sublicense,\n# and/or sell copies of the Software, and to permit persons to whom the\n# Software is furnished to do so, subject to the following conditions:\n#\n# The above copyright notice and this permission notice shall be included in\n# all copies or substantial portions of the Software.\n#\n# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL\n# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING\n# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER\n# DEALINGS IN THE SOFTWARE.\n################################################################################\n\nCUDA_VER?=\nifeq ($(CUDA_VER),)\n $(error \"CUDA_VER is not set\")\nendif\nCC:= g++\nNVCC:=/usr/local/cuda-$(CUDA_VER)/bin/nvcc\n\nCFLAGS:= -Wall -std=c++11 -shared -fPIC -Wno-error=deprecated-declarations\nCFLAGS+= -I../../includes -I/usr/local/cuda-$(CUDA_VER)/include -I/opt/nvidia/deepstream/deepstream-5.0/sources/includes\n\nLIBS:= -lnvinfer_plugin -lnvinfer -lnvparsers -L/usr/local/cuda-$(CUDA_VER)/lib64 -lcudart -lcublas -lstdc++fs -L/opt/nvidia/deepstream/deepstream-5.0/lib\nLFLAGS:= -shared -Wl,--start-group $(LIBS) -Wl,--end-group\n\nINCS:= $(wildcard *.h)\nSRCFILES:= nvdsinfer_yolo_engine.cpp \\\n nvdsparsebbox_Yolo.cpp \\\n yoloPlugins.cpp \\\n trt_utils.cpp \\\n yolo.cpp \\\n kernels.cu\nTARGET_LIB:= libnvdsinfer_custom_impl_Yolo.so\n\nTARGET_OBJS:= $(SRCFILES:.cpp=.o)\nTARGET_OBJS:= $(TARGET_OBJS:.cu=.o)\n\nall: $(TARGET_LIB)\n\n%.o: %.cpp $(INCS) Makefile\n\t$(CC) -c -o $@ $(CFLAGS) $<\n\n%.o: %.cu $(INCS) Makefile\n\t$(NVCC) -c -o $@ --compiler-options '-fPIC' $<\n\n$(TARGET_LIB) : $(TARGET_OBJS)\n\t$(CC) -o $@ $(TARGET_OBJS) $(LFLAGS)\n\nclean:\n\trm -rf $(TARGET_LIB)\n"
},
{
"path": "deepstream_plugin_yolov4/README.md",
"content": "## YOLOv4 plugin for DeepStream\nThis plugin was obtained from: https://github.com/Tianxiaomo/pytorch-YOLOv4/tree/master/DeepStream\n\nIt must be compiled locally on the jetson device:\n```\nexport CUDA_VER=10.2\nmake\n```\n"
},
{
"path": "deepstream_plugin_yolov4/kernels.cu",
"content": "/*\n * Copyright (c) 2018-2019 NVIDIA Corporation. All rights reserved.\n *\n * NVIDIA Corporation and its licensors retain all intellectual property\n * and proprietary rights in and to this software, related documentation\n * and any modifications thereto. Any use, reproduction, disclosure or\n * distribution of this software and related documentation without an express\n * license agreement from NVIDIA Corporation is strictly prohibited.\n *\n */\n\n#include \n#include \n#include \n#include \n#include \n\ninline __device__ float sigmoidGPU(const float& x) { return 1.0f / (1.0f + __expf(-x)); }\n\n__global__ void gpuYoloLayerV3(const float* input, float* output, const uint gridSize, const uint numOutputClasses,\n const uint numBBoxes)\n{\n uint x_id = blockIdx.x * blockDim.x + threadIdx.x;\n uint y_id = blockIdx.y * blockDim.y + threadIdx.y;\n uint z_id = blockIdx.z * blockDim.z + threadIdx.z;\n\n if ((x_id >= gridSize) || (y_id >= gridSize) || (z_id >= numBBoxes))\n {\n return;\n }\n\n const int numGridCells = gridSize * gridSize;\n const int bbindex = y_id * gridSize + x_id;\n\n output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)]\n = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)]);\n\n output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)]\n = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)]);\n\n output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)]\n = __expf(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)]);\n\n output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)]\n = __expf(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)]);\n\n output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)]\n = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)]);\n\n for (uint i = 0; i < numOutputClasses; ++i)\n {\n output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))]\n = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))]);\n }\n}\n\ncudaError_t cudaYoloLayerV3(const void* input, void* output, const uint& batchSize, const uint& gridSize,\n const uint& numOutputClasses, const uint& numBBoxes,\n uint64_t outputSize, cudaStream_t stream);\n\ncudaError_t cudaYoloLayerV3(const void* input, void* output, const uint& batchSize, const uint& gridSize,\n const uint& numOutputClasses, const uint& numBBoxes,\n uint64_t outputSize, cudaStream_t stream)\n{\n dim3 threads_per_block(16, 16, 4);\n dim3 number_of_blocks((gridSize / threads_per_block.x) + 1,\n (gridSize / threads_per_block.y) + 1,\n (numBBoxes / threads_per_block.z) + 1);\n for (unsigned int batch = 0; batch < batchSize; ++batch)\n {\n gpuYoloLayerV3<<>>(\n reinterpret_cast(input) + (batch * outputSize),\n reinterpret_cast(output) + (batch * outputSize), gridSize, numOutputClasses,\n numBBoxes);\n }\n return cudaGetLastError();\n}\n"
},
{
"path": "deepstream_plugin_yolov4/nvdsinfer_yolo_engine.cpp",
"content": "/*\n * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved.\n *\n * Permission is hereby granted, free of charge, to any person obtaining a\n * copy of this software and associated documentation files (the \"Software\"),\n * to deal in the Software without restriction, including without limitation\n * the rights to use, copy, modify, merge, publish, distribute, sublicense,\n * and/or sell copies of the Software, and to permit persons to whom the\n * Software is furnished to do so, subject to the following conditions:\n *\n * The above copyright notice and this permission notice shall be included in\n * all copies or substantial portions of the Software.\n *\n * THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL\n * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING\n * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER\n * DEALINGS IN THE SOFTWARE.\n */\n\n#include \"nvdsinfer_custom_impl.h\"\n#include \"nvdsinfer_context.h\"\n#include \"yoloPlugins.h\"\n#include \"yolo.h\"\n\n#include \n\n#define USE_CUDA_ENGINE_GET_API 1\n\nstatic bool getYoloNetworkInfo (NetworkInfo &networkInfo, const NvDsInferContextInitParams* initParams)\n{\n std::string yoloCfg = initParams->customNetworkConfigFilePath;\n std::string yoloType;\n\n std::transform (yoloCfg.begin(), yoloCfg.end(), yoloCfg.begin(), [] (uint8_t c) {\n return std::tolower (c);});\n\n if (yoloCfg.find(\"yolov2\") != std::string::npos) {\n if (yoloCfg.find(\"yolov2-tiny\") != std::string::npos)\n yoloType = \"yolov2-tiny\";\n else\n yoloType = \"yolov2\";\n } else if (yoloCfg.find(\"yolov3\") != std::string::npos) {\n if (yoloCfg.find(\"yolov3-tiny\") != std::string::npos)\n yoloType = \"yolov3-tiny\";\n else\n yoloType = \"yolov3\";\n } else {\n std::cerr << \"Yolo type is not defined from config file name:\"\n << yoloCfg << std::endl;\n return false;\n }\n\n networkInfo.networkType = yoloType;\n networkInfo.configFilePath = initParams->customNetworkConfigFilePath;\n networkInfo.wtsFilePath = initParams->modelFilePath;\n networkInfo.deviceType = (initParams->useDLA ? \"kDLA\" : \"kGPU\");\n networkInfo.inputBlobName = \"data\";\n\n if (networkInfo.configFilePath.empty() ||\n networkInfo.wtsFilePath.empty()) {\n std::cerr << \"Yolo config file or weights file is NOT specified.\"\n << std::endl;\n return false;\n }\n\n if (!fileExists(networkInfo.configFilePath) ||\n !fileExists(networkInfo.wtsFilePath)) {\n std::cerr << \"Yolo config file or weights file is NOT exist.\"\n << std::endl;\n return false;\n }\n\n return true;\n}\n\n#if !USE_CUDA_ENGINE_GET_API\nIModelParser* NvDsInferCreateModelParser(\n const NvDsInferContextInitParams* initParams) {\n NetworkInfo networkInfo;\n if (!getYoloNetworkInfo(networkInfo, initParams)) {\n return nullptr;\n }\n\n return new Yolo(networkInfo);\n}\n#else\nextern \"C\"\nbool NvDsInferYoloCudaEngineGet(nvinfer1::IBuilder * const builder,\n const NvDsInferContextInitParams * const initParams,\n nvinfer1::DataType dataType,\n nvinfer1::ICudaEngine *& cudaEngine);\n\nextern \"C\"\nbool NvDsInferYoloCudaEngineGet(nvinfer1::IBuilder * const builder,\n const NvDsInferContextInitParams * const initParams,\n nvinfer1::DataType dataType,\n nvinfer1::ICudaEngine *& cudaEngine)\n{\n NetworkInfo networkInfo;\n if (!getYoloNetworkInfo(networkInfo, initParams)) {\n return false;\n }\n\n Yolo yolo(networkInfo);\n cudaEngine = yolo.createEngine (builder);\n if (cudaEngine == nullptr)\n {\n std::cerr << \"Failed to build cuda engine on \"\n << networkInfo.configFilePath << std::endl;\n return false;\n }\n\n return true;\n}\n#endif\n"
},
{
"path": "deepstream_plugin_yolov4/nvdsparsebbox_Yolo.cpp",
"content": "/*\n * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.\n *\n * Permission is hereby granted, free of charge, to any person obtaining a\n * copy of this software and associated documentation files (the \"Software\"),\n * to deal in the Software without restriction, including without limitation\n * the rights to use, copy, modify, merge, publish, distribute, sublicense,\n * and/or sell copies of the Software, and to permit persons to whom the\n * Software is furnished to do so, subject to the following conditions:\n *\n * The above copyright notice and this permission notice shall be included in\n * all copies or substantial portions of the Software.\n *\n * THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL\n * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING\n * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER\n * DEALINGS IN THE SOFTWARE.\n */\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"nvdsinfer_custom_impl.h\"\n#include \"trt_utils.h\"\n\nstatic const int NUM_CLASSES_YOLO = 4; // This is just for checkings, keep updated by hand\n\nextern \"C\" bool NvDsInferParseCustomYoloV3(\n std::vector const &outputLayersInfo,\n NvDsInferNetworkInfo const &networkInfo,\n NvDsInferParseDetectionParams const &detectionParams,\n std::vector &objectList);\n\nextern \"C\" bool NvDsInferParseCustomYoloV3Tiny(\n std::vector const &outputLayersInfo,\n NvDsInferNetworkInfo const &networkInfo,\n NvDsInferParseDetectionParams const &detectionParams,\n std::vector &objectList);\n\nextern \"C\" bool NvDsInferParseCustomYoloV2(\n std::vector const &outputLayersInfo,\n NvDsInferNetworkInfo const &networkInfo,\n NvDsInferParseDetectionParams const &detectionParams,\n std::vector &objectList);\n\nextern \"C\" bool NvDsInferParseCustomYoloV2Tiny(\n std::vector const &outputLayersInfo,\n NvDsInferNetworkInfo const &networkInfo,\n NvDsInferParseDetectionParams const &detectionParams,\n std::vector &objectList);\n\nextern \"C\" bool NvDsInferParseCustomYoloTLT(\n std::vector const &outputLayersInfo,\n NvDsInferNetworkInfo const &networkInfo,\n NvDsInferParseDetectionParams const &detectionParams,\n std::vector &objectList);\n\nextern \"C\" bool NvDsInferParseCustomYoloV4(\n std::vector const &outputLayersInfo,\n NvDsInferNetworkInfo const &networkInfo,\n NvDsInferParseDetectionParams const &detectionParams,\n std::vector &objectList);\n\n/* This is a sample bounding box parsing function for the sample YoloV3 detector model */\nstatic NvDsInferParseObjectInfo convertBBox(const float &bx, const float &by, const float &bw,\n const float &bh, const int &stride, const uint &netW,\n const uint &netH)\n{\n NvDsInferParseObjectInfo b;\n // Restore coordinates to network input resolution\n float xCenter = bx * stride;\n float yCenter = by * stride;\n float x0 = xCenter - bw / 2;\n float y0 = yCenter - bh / 2;\n float x1 = x0 + bw;\n float y1 = y0 + bh;\n\n x0 = clamp(x0, 0, netW);\n y0 = clamp(y0, 0, netH);\n x1 = clamp(x1, 0, netW);\n y1 = clamp(y1, 0, netH);\n\n b.left = x0;\n b.width = clamp(x1 - x0, 0, netW);\n b.top = y0;\n b.height = clamp(y1 - y0, 0, netH);\n\n return b;\n}\n\nstatic void addBBoxProposal(const float bx, const float by, const float bw, const float bh,\n const uint stride, const uint &netW, const uint &netH, const int maxIndex,\n const float maxProb, std::vector &binfo)\n{\n NvDsInferParseObjectInfo bbi = convertBBox(bx, by, bw, bh, stride, netW, netH);\n if (bbi.width < 1 || bbi.height < 1)\n return;\n\n bbi.detectionConfidence = maxProb;\n bbi.classId = maxIndex;\n binfo.push_back(bbi);\n}\n\nstatic std::vector\ndecodeYoloV2Tensor(\n const float *detections, const std::vector &anchors,\n const uint gridSizeW, const uint gridSizeH, const uint stride, const uint numBBoxes,\n const uint numOutputClasses, const uint &netW,\n const uint &netH)\n{\n std::vector binfo;\n for (uint y = 0; y < gridSizeH; ++y)\n {\n for (uint x = 0; x < gridSizeW; ++x)\n {\n for (uint b = 0; b < numBBoxes; ++b)\n {\n const float pw = anchors[b * 2];\n const float ph = anchors[b * 2 + 1];\n\n const int numGridCells = gridSizeH * gridSizeW;\n const int bbindex = y * gridSizeW + x;\n const float bx = x + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 0)];\n const float by = y + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 1)];\n const float bw = pw * exp(detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 2)]);\n const float bh = ph * exp(detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 3)]);\n\n const float objectness = detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 4)];\n\n float maxProb = 0.0f;\n int maxIndex = -1;\n\n for (uint i = 0; i < numOutputClasses; ++i)\n {\n float prob = (detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + (5 + i))]);\n\n if (prob > maxProb)\n {\n maxProb = prob;\n maxIndex = i;\n }\n }\n maxProb = objectness * maxProb;\n\n addBBoxProposal(bx, by, bw, bh, stride, netW, netH, maxIndex, maxProb, binfo);\n }\n }\n }\n return binfo;\n}\n\nstatic std::vector\ndecodeYoloV3Tensor(\n const float *detections, const std::vector &mask, const std::vector &anchors,\n const uint gridSizeW, const uint gridSizeH, const uint stride, const uint numBBoxes,\n const uint numOutputClasses, const uint &netW,\n const uint &netH)\n{\n std::vector binfo;\n for (uint y = 0; y < gridSizeH; ++y)\n {\n for (uint x = 0; x < gridSizeW; ++x)\n {\n for (uint b = 0; b < numBBoxes; ++b)\n {\n const float pw = anchors[mask[b] * 2];\n const float ph = anchors[mask[b] * 2 + 1];\n\n const int numGridCells = gridSizeH * gridSizeW;\n const int bbindex = y * gridSizeW + x;\n const float bx = x + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 0)];\n const float by = y + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 1)];\n const float bw = pw * detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 2)];\n const float bh = ph * detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 3)];\n\n const float objectness = detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 4)];\n\n float maxProb = 0.0f;\n int maxIndex = -1;\n\n for (uint i = 0; i < numOutputClasses; ++i)\n {\n float prob = (detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + (5 + i))]);\n\n if (prob > maxProb)\n {\n maxProb = prob;\n maxIndex = i;\n }\n }\n maxProb = objectness * maxProb;\n\n addBBoxProposal(bx, by, bw, bh, stride, netW, netH, maxIndex, maxProb, binfo);\n }\n }\n }\n return binfo;\n}\n\nstatic inline std::vector\nSortLayers(const std::vector &outputLayersInfo)\n{\n std::vector outLayers;\n for (auto const &layer : outputLayersInfo)\n {\n outLayers.push_back(&layer);\n }\n std::sort(outLayers.begin(), outLayers.end(),\n [](const NvDsInferLayerInfo *a, const NvDsInferLayerInfo *b) {\n return a->inferDims.d[1] < b->inferDims.d[1];\n });\n return outLayers;\n}\n\nstatic bool NvDsInferParseYoloV3(\n std::vector const &outputLayersInfo,\n NvDsInferNetworkInfo const &networkInfo,\n NvDsInferParseDetectionParams const &detectionParams,\n std::vector &objectList,\n const std::vector &anchors,\n const std::vector> &masks)\n{\n const uint kNUM_BBOXES = 3;\n\n const std::vector sortedLayers =\n SortLayers(outputLayersInfo);\n\n if (sortedLayers.size() != masks.size())\n {\n std::cerr << \"ERROR: yoloV3 output layer.size: \" << sortedLayers.size()\n << \" does not match mask.size: \" << masks.size() << std::endl;\n return false;\n }\n\n if (NUM_CLASSES_YOLO != detectionParams.numClassesConfigured)\n {\n std::cerr << \"WARNING: Num classes mismatch. Configured:\"\n << detectionParams.numClassesConfigured\n << \", detected by network: \" << NUM_CLASSES_YOLO << std::endl;\n }\n\n std::vector objects;\n\n for (uint idx = 0; idx < masks.size(); ++idx)\n {\n const NvDsInferLayerInfo &layer = *sortedLayers[idx]; // 255 x Grid x Grid\n\n assert(layer.inferDims.numDims == 3);\n const uint gridSizeH = layer.inferDims.d[1];\n const uint gridSizeW = layer.inferDims.d[2];\n const uint stride = DIVUP(networkInfo.width, gridSizeW);\n assert(stride == DIVUP(networkInfo.height, gridSizeH));\n\n std::vector outObjs =\n decodeYoloV3Tensor((const float *)(layer.buffer), masks[idx], anchors, gridSizeW, gridSizeH, stride, kNUM_BBOXES,\n NUM_CLASSES_YOLO, networkInfo.width, networkInfo.height);\n objects.insert(objects.end(), outObjs.begin(), outObjs.end());\n }\n\n objectList = objects;\n\n return true;\n}\n\nstatic NvDsInferParseObjectInfo convertBBoxYoloV4(const float &bx1, const float &by1, const float &bx2,\n const float &by2, const uint &netW, const uint &netH)\n{\n NvDsInferParseObjectInfo b;\n // Restore coordinates to network input resolution\n\n float x1 = bx1 * netW;\n float y1 = by1 * netH;\n float x2 = bx2 * netW;\n float y2 = by2 * netH;\n\n x1 = clamp(x1, 0, netW);\n y1 = clamp(y1, 0, netH);\n x2 = clamp(x2, 0, netW);\n y2 = clamp(y2, 0, netH);\n\n b.left = x1;\n b.width = clamp(x2 - x1, 0, netW);\n b.top = y1;\n b.height = clamp(y2 - y1, 0, netH);\n\n return b;\n}\n\nstatic void addBBoxProposalYoloV4(const float bx, const float by, const float bw, const float bh,\n const uint &netW, const uint &netH, const int maxIndex,\n const float maxProb, std::vector &binfo)\n{\n NvDsInferParseObjectInfo bbi = convertBBoxYoloV4(bx, by, bw, bh, netW, netH);\n if (bbi.width < 1 || bbi.height < 1)\n return;\n\n bbi.detectionConfidence = maxProb;\n bbi.classId = maxIndex;\n binfo.push_back(bbi);\n}\n\nstatic std::vector\ndecodeYoloV4Tensor(\n const float *boxes, const float *scores,\n const uint num_bboxes, NvDsInferParseDetectionParams const &detectionParams,\n const uint &netW, const uint &netH)\n{\n std::vector binfo;\n\n uint bbox_location = 0;\n uint score_location = 0;\n for (uint b = 0; b < num_bboxes; ++b)\n {\n float bx1 = boxes[bbox_location];\n float by1 = boxes[bbox_location + 1];\n float bx2 = boxes[bbox_location + 2];\n float by2 = boxes[bbox_location + 3];\n\n float maxProb = 0.0f;\n int maxIndex = -1;\n\n for (uint c = 0; c < detectionParams.numClassesConfigured; ++c)\n {\n float prob = scores[score_location + c];\n if (prob > maxProb)\n {\n maxProb = prob;\n maxIndex = c;\n }\n }\n\n if (maxProb > detectionParams.perClassPreclusterThreshold[maxIndex])\n {\n addBBoxProposalYoloV4(bx1, by1, bx2, by2, netW, netH, maxIndex, maxProb, binfo);\n }\n\n bbox_location += 4;\n score_location += detectionParams.numClassesConfigured;\n }\n\n return binfo;\n}\n\n/* C-linkage to prevent name-mangling */\n\nstatic bool NvDsInferParseYoloV4(\n std::vector const &outputLayersInfo,\n NvDsInferNetworkInfo const &networkInfo,\n NvDsInferParseDetectionParams const &detectionParams,\n std::vector &objectList)\n{\n if (NUM_CLASSES_YOLO != detectionParams.numClassesConfigured)\n {\n std::cerr << \"WARNING: Num classes mismatch. Configured:\"\n << detectionParams.numClassesConfigured\n << \", detected by network: \" << NUM_CLASSES_YOLO << std::endl;\n }\n\n std::vector objects;\n\n const NvDsInferLayerInfo &boxes = outputLayersInfo[0]; // num_boxes x 4\n const NvDsInferLayerInfo &scores = outputLayersInfo[1]; // num_boxes x num_classes\n const NvDsInferLayerInfo &subbox = outputLayersInfo[2];\n //* printf(\"%d\\n\", subbox.inferDims.numDims);\n // 3 dimensional: [num_boxes, 1, 4]\n assert(boxes.inferDims.numDims == 3);\n // 2 dimensional: [num_boxes, num_classes]\n assert(scores.inferDims.numDims == 2);\n\n // The second dimension should be num_classes\n assert(detectionParams.numClassesConfigured == scores.inferDims.d[1]);\n\n uint num_bboxes = boxes.inferDims.d[0];\n\n // std::cout << \"Network Info: \" << networkInfo.height << \" \" << networkInfo.width << std::endl;\n\n std::vector outObjs =\n decodeYoloV4Tensor(\n (const float *)(boxes.buffer), (const float *)(scores.buffer), num_bboxes, detectionParams,\n networkInfo.width, networkInfo.height);\n\n objects.insert(objects.end(), outObjs.begin(), outObjs.end());\n\n objectList = objects;\n\n return true;\n}\n\nextern \"C\" bool NvDsInferParseCustomYoloV4(\n std::vector const &outputLayersInfo,\n NvDsInferNetworkInfo const &networkInfo,\n NvDsInferParseDetectionParams const &detectionParams,\n std::vector &objectList)\n{\n return NvDsInferParseYoloV4(\n outputLayersInfo, networkInfo, detectionParams, objectList);\n}\n\nextern \"C\" bool NvDsInferParseCustomYoloV3(\n std::vector const &outputLayersInfo,\n NvDsInferNetworkInfo const &networkInfo,\n NvDsInferParseDetectionParams const &detectionParams,\n std::vector &objectList)\n{\n static const std::vector kANCHORS = {\n 10.0, 13.0, 16.0, 30.0, 33.0, 23.0, 30.0, 61.0, 62.0,\n 45.0, 59.0, 119.0, 116.0, 90.0, 156.0, 198.0, 373.0, 326.0};\n static const std::vector> kMASKS = {\n {6, 7, 8},\n {3, 4, 5},\n {0, 1, 2}};\n return NvDsInferParseYoloV3(\n outputLayersInfo, networkInfo, detectionParams, objectList,\n kANCHORS, kMASKS);\n}\n\nextern \"C\" bool NvDsInferParseCustomYoloV3Tiny(\n std::vector const &outputLayersInfo,\n NvDsInferNetworkInfo const &networkInfo,\n NvDsInferParseDetectionParams const &detectionParams,\n std::vector &objectList)\n{\n static const std::vector kANCHORS = {\n 10, 14, 23, 27, 37, 58, 81, 82, 135, 169, 344, 319};\n static const std::vector> kMASKS = {\n {3, 4, 5},\n //{0, 1, 2}}; // as per output result, select {1,2,3}\n {1, 2, 3}};\n\n return NvDsInferParseYoloV3(\n outputLayersInfo, networkInfo, detectionParams, objectList,\n kANCHORS, kMASKS);\n}\n\nstatic bool NvDsInferParseYoloV2(\n std::vector const &outputLayersInfo,\n NvDsInferNetworkInfo const &networkInfo,\n NvDsInferParseDetectionParams const &detectionParams,\n std::vector &objectList)\n{\n // copy anchor data from yolov2.cfg file\n std::vector anchors = {0.57273, 0.677385, 1.87446, 2.06253, 3.33843,\n 5.47434, 7.88282, 3.52778, 9.77052, 9.16828};\n const uint kNUM_BBOXES = 5;\n\n if (outputLayersInfo.empty())\n {\n std::cerr << \"Could not find output layer in bbox parsing\" << std::endl;\n ;\n return false;\n }\n const NvDsInferLayerInfo &layer = outputLayersInfo[0];\n\n if (NUM_CLASSES_YOLO != detectionParams.numClassesConfigured)\n {\n std::cerr << \"WARNING: Num classes mismatch. Configured:\"\n << detectionParams.numClassesConfigured\n << \", detected by network: \" << NUM_CLASSES_YOLO << std::endl;\n }\n\n assert(layer.inferDims.numDims == 3);\n const uint gridSizeH = layer.inferDims.d[1];\n const uint gridSizeW = layer.inferDims.d[2];\n const uint stride = DIVUP(networkInfo.width, gridSizeW);\n assert(stride == DIVUP(networkInfo.height, gridSizeH));\n for (auto &anchor : anchors)\n {\n anchor *= stride;\n }\n std::vector objects =\n decodeYoloV2Tensor((const float *)(layer.buffer), anchors, gridSizeW, gridSizeH, stride, kNUM_BBOXES,\n NUM_CLASSES_YOLO, networkInfo.width, networkInfo.height);\n\n objectList = objects;\n\n return true;\n}\n\nextern \"C\" bool NvDsInferParseCustomYoloV2(\n std::vector const &outputLayersInfo,\n NvDsInferNetworkInfo const &networkInfo,\n NvDsInferParseDetectionParams const &detectionParams,\n std::vector &objectList)\n{\n return NvDsInferParseYoloV2(\n outputLayersInfo, networkInfo, detectionParams, objectList);\n}\n\nextern \"C\" bool NvDsInferParseCustomYoloV2Tiny(\n std::vector const &outputLayersInfo,\n NvDsInferNetworkInfo const &networkInfo,\n NvDsInferParseDetectionParams const &detectionParams,\n std::vector &objectList)\n{\n return NvDsInferParseYoloV2(\n outputLayersInfo, networkInfo, detectionParams, objectList);\n}\n\nextern \"C\" bool NvDsInferParseCustomYoloTLT(\n std::vector const &outputLayersInfo,\n NvDsInferNetworkInfo const &networkInfo,\n NvDsInferParseDetectionParams const &detectionParams,\n std::vector &objectList)\n{\n\n if (outputLayersInfo.size() != 4)\n {\n std::cerr << \"Mismatch in the number of output buffers.\"\n << \"Expected 4 output buffers, detected in the network :\"\n << outputLayersInfo.size() << std::endl;\n return false;\n }\n\n const int topK = 200;\n const int *keepCount = static_cast(outputLayersInfo.at(0).buffer);\n const float *boxes = static_cast(outputLayersInfo.at(1).buffer);\n const float *scores = static_cast(outputLayersInfo.at(2).buffer);\n const float *cls = static_cast(outputLayersInfo.at(3).buffer);\n\n for (int i = 0; (i < keepCount[0]) && (objectList.size() <= topK); ++i)\n {\n const float *loc = &boxes[0] + (i * 4);\n const float *conf = &scores[0] + i;\n const float *cls_id = &cls[0] + i;\n\n if (conf[0] > 1.001)\n continue;\n\n if ((loc[0] < 0) || (loc[1] < 0) || (loc[2] < 0) || (loc[3] < 0))\n continue;\n\n if ((loc[0] > networkInfo.width) || (loc[2] > networkInfo.width) || (loc[1] > networkInfo.height) || (loc[3] > networkInfo.width))\n continue;\n\n if ((loc[2] < loc[0]) || (loc[3] < loc[1]))\n continue;\n\n if (((loc[3] - loc[1]) > networkInfo.height) || ((loc[2] - loc[0]) > networkInfo.width))\n continue;\n\n NvDsInferParseObjectInfo curObj{static_cast(cls_id[0]),\n loc[0], loc[1], (loc[2] - loc[0]),\n (loc[3] - loc[1]), conf[0]};\n objectList.push_back(curObj);\n }\n\n return true;\n}\n\n/* Check that the custom function has been defined correctly */\nCHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV4);\nCHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV3);\nCHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV3Tiny);\nCHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV2);\nCHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV2Tiny);\nCHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloTLT);\n"
},
{
"path": "deepstream_plugin_yolov4/trt_utils.cpp",
"content": "/*\n * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved.\n *\n * Permission is hereby granted, free of charge, to any person obtaining a\n * copy of this software and associated documentation files (the \"Software\"),\n * to deal in the Software without restriction, including without limitation\n * the rights to use, copy, modify, merge, publish, distribute, sublicense,\n * and/or sell copies of the Software, and to permit persons to whom the\n * Software is furnished to do so, subject to the following conditions:\n *\n * The above copyright notice and this permission notice shall be included in\n * all copies or substantial portions of the Software.\n *\n * THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL\n * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING\n * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER\n * DEALINGS IN THE SOFTWARE.\n */\n\n#include \"trt_utils.h\"\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"NvInferPlugin.h\"\n\nstatic void leftTrim(std::string& s)\n{\n s.erase(s.begin(), std::find_if(s.begin(), s.end(), [](int ch) { return !isspace(ch); }));\n}\n\nstatic void rightTrim(std::string& s)\n{\n s.erase(std::find_if(s.rbegin(), s.rend(), [](int ch) { return !isspace(ch); }).base(), s.end());\n}\n\nstd::string trim(std::string s)\n{\n leftTrim(s);\n rightTrim(s);\n return s;\n}\n\nfloat clamp(const float val, const float minVal, const float maxVal)\n{\n assert(minVal <= maxVal);\n return std::min(maxVal, std::max(minVal, val));\n}\n\nbool fileExists(const std::string fileName, bool verbose)\n{\n if (!std::experimental::filesystem::exists(std::experimental::filesystem::path(fileName)))\n {\n if (verbose) std::cout << \"File does not exist : \" << fileName << std::endl;\n return false;\n }\n return true;\n}\n\nstd::vector loadWeights(const std::string weightsFilePath, const std::string& networkType)\n{\n assert(fileExists(weightsFilePath));\n std::cout << \"Loading pre-trained weights...\" << std::endl;\n std::ifstream file(weightsFilePath, std::ios_base::binary);\n assert(file.good());\n std::string line;\n\n if (networkType == \"yolov2\")\n {\n // Remove 4 int32 bytes of data from the stream belonging to the header\n file.ignore(4 * 4);\n }\n else if ((networkType == \"yolov3\") || (networkType == \"yolov3-tiny\")\n || (networkType == \"yolov2-tiny\"))\n {\n // Remove 5 int32 bytes of data from the stream belonging to the header\n file.ignore(4 * 5);\n }\n else\n {\n std::cout << \"Invalid network type\" << std::endl;\n assert(0);\n }\n\n std::vector weights;\n char floatWeight[4];\n while (!file.eof())\n {\n file.read(floatWeight, 4);\n assert(file.gcount() == 4);\n weights.push_back(*reinterpret_cast(floatWeight));\n if (file.peek() == std::istream::traits_type::eof()) break;\n }\n std::cout << \"Loading weights of \" << networkType << \" complete!\"\n << std::endl;\n std::cout << \"Total Number of weights read : \" << weights.size() << std::endl;\n return weights;\n}\n\nstd::string dimsToString(const nvinfer1::Dims d)\n{\n std::stringstream s;\n assert(d.nbDims >= 1);\n for (int i = 0; i < d.nbDims - 1; ++i)\n {\n s << std::setw(4) << d.d[i] << \" x\";\n }\n s << std::setw(4) << d.d[d.nbDims - 1];\n\n return s.str();\n}\n\nvoid displayDimType(const nvinfer1::Dims d)\n{\n std::cout << \"(\" << d.nbDims << \") \";\n for (int i = 0; i < d.nbDims; ++i)\n {\n switch (d.type[i])\n {\n case nvinfer1::DimensionType::kSPATIAL: std::cout << \"kSPATIAL \"; break;\n case nvinfer1::DimensionType::kCHANNEL: std::cout << \"kCHANNEL \"; break;\n case nvinfer1::DimensionType::kINDEX: std::cout << \"kINDEX \"; break;\n case nvinfer1::DimensionType::kSEQUENCE: std::cout << \"kSEQUENCE \"; break;\n }\n }\n std::cout << std::endl;\n}\n\nint getNumChannels(nvinfer1::ITensor* t)\n{\n nvinfer1::Dims d = t->getDimensions();\n assert(d.nbDims == 3);\n\n return d.d[0];\n}\n\nuint64_t get3DTensorVolume(nvinfer1::Dims inputDims)\n{\n assert(inputDims.nbDims == 3);\n return inputDims.d[0] * inputDims.d[1] * inputDims.d[2];\n}\n\nnvinfer1::ILayer* netAddMaxpool(int layerIdx, std::map& block,\n nvinfer1::ITensor* input, nvinfer1::INetworkDefinition* network)\n{\n assert(block.at(\"type\") == \"maxpool\");\n assert(block.find(\"size\") != block.end());\n assert(block.find(\"stride\") != block.end());\n\n int size = std::stoi(block.at(\"size\"));\n int stride = std::stoi(block.at(\"stride\"));\n\n nvinfer1::IPoolingLayer* pool\n = network->addPooling(*input, nvinfer1::PoolingType::kMAX, nvinfer1::DimsHW{size, size});\n assert(pool);\n std::string maxpoolLayerName = \"maxpool_\" + std::to_string(layerIdx);\n pool->setStride(nvinfer1::DimsHW{stride, stride});\n pool->setPaddingMode(nvinfer1::PaddingMode::kSAME_UPPER);\n pool->setName(maxpoolLayerName.c_str());\n\n return pool;\n}\n\nnvinfer1::ILayer* netAddConvLinear(int layerIdx, std::map& block,\n std::vector& weights,\n std::vector& trtWeights, int& weightPtr,\n int& inputChannels, nvinfer1::ITensor* input,\n nvinfer1::INetworkDefinition* network)\n{\n assert(block.at(\"type\") == \"convolutional\");\n assert(block.find(\"batch_normalize\") == block.end());\n assert(block.at(\"activation\") == \"linear\");\n assert(block.find(\"filters\") != block.end());\n assert(block.find(\"pad\") != block.end());\n assert(block.find(\"size\") != block.end());\n assert(block.find(\"stride\") != block.end());\n\n int filters = std::stoi(block.at(\"filters\"));\n int padding = std::stoi(block.at(\"pad\"));\n int kernelSize = std::stoi(block.at(\"size\"));\n int stride = std::stoi(block.at(\"stride\"));\n int pad;\n if (padding)\n pad = (kernelSize - 1) / 2;\n else\n pad = 0;\n // load the convolution layer bias\n nvinfer1::Weights convBias{nvinfer1::DataType::kFLOAT, nullptr, filters};\n float* val = new float[filters];\n for (int i = 0; i < filters; ++i)\n {\n val[i] = weights[weightPtr];\n weightPtr++;\n }\n convBias.values = val;\n trtWeights.push_back(convBias);\n // load the convolutional layer weights\n int size = filters * inputChannels * kernelSize * kernelSize;\n nvinfer1::Weights convWt{nvinfer1::DataType::kFLOAT, nullptr, size};\n val = new float[size];\n for (int i = 0; i < size; ++i)\n {\n val[i] = weights[weightPtr];\n weightPtr++;\n }\n convWt.values = val;\n trtWeights.push_back(convWt);\n nvinfer1::IConvolutionLayer* conv = network->addConvolution(\n *input, filters, nvinfer1::DimsHW{kernelSize, kernelSize}, convWt, convBias);\n assert(conv != nullptr);\n std::string convLayerName = \"conv_\" + std::to_string(layerIdx);\n conv->setName(convLayerName.c_str());\n conv->setStride(nvinfer1::DimsHW{stride, stride});\n conv->setPadding(nvinfer1::DimsHW{pad, pad});\n\n return conv;\n}\n\nnvinfer1::ILayer* netAddConvBNLeaky(int layerIdx, std::map& block,\n std::vector& weights,\n std::vector& trtWeights, int& weightPtr,\n int& inputChannels, nvinfer1::ITensor* input,\n nvinfer1::INetworkDefinition* network)\n{\n assert(block.at(\"type\") == \"convolutional\");\n assert(block.find(\"batch_normalize\") != block.end());\n assert(block.at(\"batch_normalize\") == \"1\");\n assert(block.at(\"activation\") == \"leaky\");\n assert(block.find(\"filters\") != block.end());\n assert(block.find(\"pad\") != block.end());\n assert(block.find(\"size\") != block.end());\n assert(block.find(\"stride\") != block.end());\n\n bool batchNormalize, bias;\n if (block.find(\"batch_normalize\") != block.end())\n {\n batchNormalize = (block.at(\"batch_normalize\") == \"1\");\n bias = false;\n }\n else\n {\n batchNormalize = false;\n bias = true;\n }\n // all conv_bn_leaky layers assume bias is false\n assert(batchNormalize == true && bias == false);\n UNUSED(batchNormalize);\n UNUSED(bias);\n\n int filters = std::stoi(block.at(\"filters\"));\n int padding = std::stoi(block.at(\"pad\"));\n int kernelSize = std::stoi(block.at(\"size\"));\n int stride = std::stoi(block.at(\"stride\"));\n int pad;\n if (padding)\n pad = (kernelSize - 1) / 2;\n else\n pad = 0;\n\n /***** CONVOLUTION LAYER *****/\n /*****************************/\n // batch norm weights are before the conv layer\n // load BN biases (bn_biases)\n std::vector bnBiases;\n for (int i = 0; i < filters; ++i)\n {\n bnBiases.push_back(weights[weightPtr]);\n weightPtr++;\n }\n // load BN weights\n std::vector bnWeights;\n for (int i = 0; i < filters; ++i)\n {\n bnWeights.push_back(weights[weightPtr]);\n weightPtr++;\n }\n // load BN running_mean\n std::vector bnRunningMean;\n for (int i = 0; i < filters; ++i)\n {\n bnRunningMean.push_back(weights[weightPtr]);\n weightPtr++;\n }\n // load BN running_var\n std::vector bnRunningVar;\n for (int i = 0; i < filters; ++i)\n {\n // 1e-05 for numerical stability\n bnRunningVar.push_back(sqrt(weights[weightPtr] + 1.0e-5));\n weightPtr++;\n }\n // load Conv layer weights (GKCRS)\n int size = filters * inputChannels * kernelSize * kernelSize;\n nvinfer1::Weights convWt{nvinfer1::DataType::kFLOAT, nullptr, size};\n float* val = new float[size];\n for (int i = 0; i < size; ++i)\n {\n val[i] = weights[weightPtr];\n weightPtr++;\n }\n convWt.values = val;\n trtWeights.push_back(convWt);\n nvinfer1::Weights convBias{nvinfer1::DataType::kFLOAT, nullptr, 0};\n trtWeights.push_back(convBias);\n nvinfer1::IConvolutionLayer* conv = network->addConvolution(\n *input, filters, nvinfer1::DimsHW{kernelSize, kernelSize}, convWt, convBias);\n assert(conv != nullptr);\n std::string convLayerName = \"conv_\" + std::to_string(layerIdx);\n conv->setName(convLayerName.c_str());\n conv->setStride(nvinfer1::DimsHW{stride, stride});\n conv->setPadding(nvinfer1::DimsHW{pad, pad});\n\n /***** BATCHNORM LAYER *****/\n /***************************/\n size = filters;\n // create the weights\n nvinfer1::Weights shift{nvinfer1::DataType::kFLOAT, nullptr, size};\n nvinfer1::Weights scale{nvinfer1::DataType::kFLOAT, nullptr, size};\n nvinfer1::Weights power{nvinfer1::DataType::kFLOAT, nullptr, size};\n float* shiftWt = new float[size];\n for (int i = 0; i < size; ++i)\n {\n shiftWt[i]\n = bnBiases.at(i) - ((bnRunningMean.at(i) * bnWeights.at(i)) / bnRunningVar.at(i));\n }\n shift.values = shiftWt;\n float* scaleWt = new float[size];\n for (int i = 0; i < size; ++i)\n {\n scaleWt[i] = bnWeights.at(i) / bnRunningVar[i];\n }\n scale.values = scaleWt;\n float* powerWt = new float[size];\n for (int i = 0; i < size; ++i)\n {\n powerWt[i] = 1.0;\n }\n power.values = powerWt;\n trtWeights.push_back(shift);\n trtWeights.push_back(scale);\n trtWeights.push_back(power);\n // Add the batch norm layers\n nvinfer1::IScaleLayer* bn = network->addScale(\n *conv->getOutput(0), nvinfer1::ScaleMode::kCHANNEL, shift, scale, power);\n assert(bn != nullptr);\n std::string bnLayerName = \"batch_norm_\" + std::to_string(layerIdx);\n bn->setName(bnLayerName.c_str());\n /***** ACTIVATION LAYER *****/\n /****************************/\n nvinfer1::ITensor* bnOutput = bn->getOutput(0);\n nvinfer1::IActivationLayer* leaky = network->addActivation(\n *bnOutput, nvinfer1::ActivationType::kLEAKY_RELU);\n leaky->setAlpha(0.1);\n assert(leaky != nullptr);\n std::string leakyLayerName = \"leaky_\" + std::to_string(layerIdx);\n leaky->setName(leakyLayerName.c_str());\n\n return leaky;\n}\n\nnvinfer1::ILayer* netAddUpsample(int layerIdx, std::map& block,\n std::vector& weights,\n std::vector& trtWeights, int& inputChannels,\n nvinfer1::ITensor* input, nvinfer1::INetworkDefinition* network)\n{\n assert(block.at(\"type\") == \"upsample\");\n nvinfer1::Dims inpDims = input->getDimensions();\n assert(inpDims.nbDims == 3);\n assert(inpDims.d[1] == inpDims.d[2]);\n int h = inpDims.d[1];\n int w = inpDims.d[2];\n int stride = std::stoi(block.at(\"stride\"));\n // add pre multiply matrix as a constant\n nvinfer1::Dims preDims{3,\n {1, stride * h, w},\n {nvinfer1::DimensionType::kCHANNEL, nvinfer1::DimensionType::kSPATIAL,\n nvinfer1::DimensionType::kSPATIAL}};\n int size = stride * h * w;\n nvinfer1::Weights preMul{nvinfer1::DataType::kFLOAT, nullptr, size};\n float* preWt = new float[size];\n /* (2*h * w)\n [ [1, 0, ..., 0],\n [1, 0, ..., 0],\n [0, 1, ..., 0],\n [0, 1, ..., 0],\n ...,\n ...,\n [0, 0, ..., 1],\n [0, 0, ..., 1] ]\n */\n for (int i = 0, idx = 0; i < h; ++i)\n {\n for (int s = 0; s < stride; ++s)\n {\n for (int j = 0; j < w; ++j, ++idx)\n {\n preWt[idx] = (i == j) ? 1.0 : 0.0;\n }\n }\n }\n preMul.values = preWt;\n trtWeights.push_back(preMul);\n nvinfer1::IConstantLayer* preM = network->addConstant(preDims, preMul);\n assert(preM != nullptr);\n std::string preLayerName = \"preMul_\" + std::to_string(layerIdx);\n preM->setName(preLayerName.c_str());\n // add post multiply matrix as a constant\n nvinfer1::Dims postDims{3,\n {1, h, stride * w},\n {nvinfer1::DimensionType::kCHANNEL, nvinfer1::DimensionType::kSPATIAL,\n nvinfer1::DimensionType::kSPATIAL}};\n size = stride * h * w;\n nvinfer1::Weights postMul{nvinfer1::DataType::kFLOAT, nullptr, size};\n float* postWt = new float[size];\n /* (h * 2*w)\n [ [1, 1, 0, 0, ..., 0, 0],\n [0, 0, 1, 1, ..., 0, 0],\n ...,\n ...,\n [0, 0, 0, 0, ..., 1, 1] ]\n */\n for (int i = 0, idx = 0; i < h; ++i)\n {\n for (int j = 0; j < stride * w; ++j, ++idx)\n {\n postWt[idx] = (j / stride == i) ? 1.0 : 0.0;\n }\n }\n postMul.values = postWt;\n trtWeights.push_back(postMul);\n nvinfer1::IConstantLayer* post_m = network->addConstant(postDims, postMul);\n assert(post_m != nullptr);\n std::string postLayerName = \"postMul_\" + std::to_string(layerIdx);\n post_m->setName(postLayerName.c_str());\n // add matrix multiply layers for upsampling\n nvinfer1::IMatrixMultiplyLayer* mm1\n = network->addMatrixMultiply(*preM->getOutput(0), nvinfer1::MatrixOperation::kNONE, *input,\n nvinfer1::MatrixOperation::kNONE);\n assert(mm1 != nullptr);\n std::string mm1LayerName = \"mm1_\" + std::to_string(layerIdx);\n mm1->setName(mm1LayerName.c_str());\n nvinfer1::IMatrixMultiplyLayer* mm2\n = network->addMatrixMultiply(*mm1->getOutput(0), nvinfer1::MatrixOperation::kNONE,\n *post_m->getOutput(0), nvinfer1::MatrixOperation::kNONE);\n assert(mm2 != nullptr);\n std::string mm2LayerName = \"mm2_\" + std::to_string(layerIdx);\n mm2->setName(mm2LayerName.c_str());\n return mm2;\n}\n\nvoid printLayerInfo(std::string layerIndex, std::string layerName, std::string layerInput,\n std::string layerOutput, std::string weightPtr)\n{\n std::cout << std::setw(6) << std::left << layerIndex << std::setw(15) << std::left << layerName;\n std::cout << std::setw(20) << std::left << layerInput << std::setw(20) << std::left\n << layerOutput;\n std::cout << std::setw(6) << std::left << weightPtr << std::endl;\n}\n"
},
{
"path": "deepstream_plugin_yolov4/trt_utils.h",
"content": "/*\n * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved.\n *\n * Permission is hereby granted, free of charge, to any person obtaining a\n * copy of this software and associated documentation files (the \"Software\"),\n * to deal in the Software without restriction, including without limitation\n * the rights to use, copy, modify, merge, publish, distribute, sublicense,\n * and/or sell copies of the Software, and to permit persons to whom the\n * Software is furnished to do so, subject to the following conditions:\n *\n * The above copyright notice and this permission notice shall be included in\n * all copies or substantial portions of the Software.\n *\n * THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL\n * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING\n * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER\n * DEALINGS IN THE SOFTWARE.\n */\n\n\n#ifndef __TRT_UTILS_H__\n#define __TRT_UTILS_H__\n\n#include \n#include