Repository: bdtinc/maskcam Branch: main Commit: 4841c2c49235 Files: 106 Total size: 388.6 KB Directory structure: gitextract_82ysmahp/ ├── .gitignore ├── Dockerfile ├── LICENSE.md ├── README.md ├── deepstream_plugin_yolov4/ │ ├── Makefile │ ├── README.md │ ├── kernels.cu │ ├── nvdsinfer_yolo_engine.cpp │ ├── nvdsparsebbox_Yolo.cpp │ ├── trt_utils.cpp │ ├── trt_utils.h │ ├── yolo.cpp │ ├── yolo.h │ ├── yoloPlugins.cpp │ └── yoloPlugins.h ├── docker/ │ ├── constraints.docker │ ├── opencv_python-3.2.0.egg-info │ ├── scikit-learn-0.19.1.egg-info │ └── start.sh ├── docker-compose.yml ├── docs/ │ ├── BalenaOS-DevKit-Nano-Setup.md │ ├── BalenaOS-Photon-Nano-Setup.md │ ├── Custom-Container-Development.md │ ├── Manual-Dependencies-Installation.md │ └── Useful-Development-Scripts.md ├── maskcam/ │ ├── common.py │ ├── config.py │ ├── maskcam_filesave.py │ ├── maskcam_fileserver.py │ ├── maskcam_inference.py │ ├── maskcam_streaming.py │ ├── mqtt_commander.py │ ├── mqtt_common.py │ ├── prints.py │ └── utils.py ├── maskcam_config.txt ├── maskcam_run.py ├── requirements-dev.in ├── requirements.in ├── requirements.txt ├── server/ │ ├── backend/ │ │ ├── Dockerfile │ │ ├── alembic.ini │ │ ├── app/ │ │ │ ├── api/ │ │ │ │ ├── __init__.py │ │ │ │ ├── exceptions.py │ │ │ │ └── routes/ │ │ │ │ ├── device_routes.py │ │ │ │ └── statistic_routes.py │ │ │ ├── core/ │ │ │ │ └── config.py │ │ │ ├── db/ │ │ │ │ ├── cruds/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── crud_device.py │ │ │ │ │ ├── crud_statistic.py │ │ │ │ │ └── crud_video_file.py │ │ │ │ ├── migrations/ │ │ │ │ │ ├── env.py │ │ │ │ │ ├── script.py.mako │ │ │ │ │ └── versions/ │ │ │ │ │ ├── 6a4d853aabce_added_database.py │ │ │ │ │ ├── 6d5c250f098c_added_device_file_server_address.py │ │ │ │ │ ├── 8f58cd776eda_added_database.py │ │ │ │ │ └── fb245977373f_added_video_file_table.py │ │ │ │ ├── schema/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── base.py │ │ │ │ │ ├── models.py │ │ │ │ │ └── schemas.py │ │ │ │ └── utils/ │ │ │ │ ├── __init__.py │ │ │ │ ├── enums.py │ │ │ │ └── utils.py │ │ │ ├── main.py │ │ │ └── mqtt/ │ │ │ ├── broker.py │ │ │ ├── publisher.py │ │ │ └── subscriber.py │ │ ├── prestart.sh │ │ ├── requirements.txt │ │ └── test_crud.py │ ├── backend.env.template │ ├── build_docker.sh │ ├── database.env.template │ ├── docker-compose.yml │ ├── frontend/ │ │ ├── Dockerfile │ │ ├── main.py │ │ ├── requirements.txt │ │ ├── session_manager.py │ │ └── utils/ │ │ ├── api_utils.py │ │ └── format_utils.py │ ├── frontend.env.template │ ├── mosquitto.conf │ ├── server_setup.sh │ └── stop_docker.sh ├── utils/ │ ├── combine_coco.py │ ├── gst_capabilities.sh │ ├── mqtt-test/ │ │ ├── broker.py │ │ ├── publisher.py │ │ └── suscriber.py │ ├── onnx_fix_mobilenet.py │ ├── remove_images_coco.py │ ├── tf1_trt_inference.py │ ├── tf2_trt_convert.py │ └── tf_trt_convert.py └── yolo/ ├── config.py ├── config_images.yml ├── data/ │ ├── obj.data │ └── obj.names ├── facemask-yolov4-tiny.cfg ├── facemask-yolov4.cfg ├── integrations/ │ └── yolo/ │ ├── detector_trt.py │ ├── utils_pytorch.py │ └── yolo_adaptor.py ├── run_yolo_images.py └── train_cu90.sh ================================================ FILE CONTENTS ================================================ ================================================ FILE: .gitignore ================================================ # Project particular ignores backup/ videos/inputs/* videos/outputs/* images*/ *.mp4 .DS_Store .vscode/ # Intermediate files to convert MobileNetV2 from TF->TF-TRT utils/*.png mobilenetv2/inference_graph_*/ mobilenetv2/converted_trt_*.pb # Avoid weights, or engines *.weights *.trt *.engine *.onnx # Environment files (not templates) server/*.env # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] *$py.class # C extensions *.so *.o # Distribution / packaging .Python build/ develop-eggs/ dist/ downloads/ eggs/ .eggs/ lib/ lib64/ parts/ sdist/ var/ wheels/ pip-wheel-metadata/ share/python-wheels/ *.egg-info/ .installed.cfg *.egg MANIFEST # PyInstaller # Usually these files are written by a python script from a template # before PyInstaller builds the exe, so as to inject date/other infos into it. *.manifest *.spec # Installer logs pip-log.txt pip-delete-this-directory.txt # Unit test / coverage reports htmlcov/ .tox/ .nox/ .coverage .coverage.* .cache nosetests.xml coverage.xml *.cover *.py,cover .hypothesis/ .pytest_cache/ # Translations *.mo *.pot # Django stuff: *.log local_settings.py db.sqlite3 db.sqlite3-journal # Flask stuff: instance/ .webassets-cache # Scrapy stuff: .scrapy # Sphinx documentation docs/_build/ # PyBuilder target/ # Jupyter Notebook .ipynb_checkpoints # IPython profile_default/ ipython_config.py # pyenv .python-version # pipenv # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. # However, in case of collaboration, if having platform-specific dependencies or dependencies # having no cross-platform support, pipenv may install dependencies that don't work, or not # install all needed dependencies. #Pipfile.lock # PEP 582; used by e.g. github.com/David-OConnor/pyflow __pypackages__/ # Celery stuff celerybeat-schedule celerybeat.pid # SageMath parsed files *.sage.py # Environments .env .venv env/ venv/ ENV/ env.bak/ venv.bak/ # Spyder project settings .spyderproject .spyproject # Rope project settings .ropeproject # mkdocs documentation /site # mypy .mypy_cache/ .dmypy.json dmypy.json # Pyre type checker .pyre/ ================================================ FILE: Dockerfile ================================================ # Installs maskcam on a BalenaOS container (devkit or Photon) FROM balenalib/jetson-nano-ubuntu:20210201 # Don't prompt with any configuration questions ENV DEBIAN_FRONTEND noninteractive # Switch the nvidia apt source repos and # install some utilities RUN \ apt-get update && apt-get install -y \ lbzip2 wget tar python3 git ENV UDEV=1 # Download and install BSP binaries for L4T 32.4.2 # This is mostly from Balena's Alan Boris at: # https://github.com/balena-io-playground/jetson-nano-sample-new/blob/master/CUDA/Dockerfile RUN apt-get update && apt-get install -y wget tar python3 libegl1 && \ wget https://developer.nvidia.com/embedded/L4T/r32_Release_v4.2/t210ref_release_aarch64/Tegra210_Linux_R32.4.2_aarch64.tbz2 && \ tar xf Tegra210_Linux_R32.4.2_aarch64.tbz2 && \ cd Linux_for_Tegra && \ sed -i 's/config.tbz2\"/config.tbz2\" --exclude=etc\/hosts --exclude=etc\/hostname/g' apply_binaries.sh && \ 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 && \ sed -i 's/LC_ALL=C chroot . mount -t proc none \/proc/ /g' nv_tegra/nv-apply-debs.sh && \ sed -i 's/umount ${L4T_ROOTFS_DIR}\/proc/ /g' nv_tegra/nv-apply-debs.sh && \ sed -i 's/chroot . \// /g' nv_tegra/nv-apply-debs.sh && \ ./apply_binaries.sh -r / --target-overlay && cd .. && \ rm -rf Tegra210_Linux_R32.4.2_aarch64.tbz2 && \ rm -rf Linux_for_Tegra && \ echo "/usr/lib/aarch64-linux-gnu/tegra" > /etc/ld.so.conf.d/nvidia-tegra.conf && \ echo "/usr/lib/aarch64-linux-gnu/tegra-egl" > /etc/ld.so.conf.d/nvidia-tegra-egl.conf && ldconfig # Install GStreamer and remove unnecessary files RUN apt-get install -y \ libssl1.0.0 \ libgstreamer1.0-0 \ gstreamer1.0-tools \ gstreamer1.0-plugins-good \ gstreamer1.0-plugins-bad \ gstreamer1.0-plugins-ugly \ gstreamer1.0-libav \ libgstrtspserver-1.0-0 \ libjansson4=2.11-1 \ cuda-toolkit-10-2 && \ ldconfig RUN \ rm -rf /usr/src/nvidia/graphics_demos \ /usr/local/cuda-10.2/samples \ /usr/local/cuda-10.2/doc # Install DeepStream RUN apt-get install -y deepstream-5.0 && \ rm -rf /opt/nvidia/deepstream/deepstream-5.0/samples \ /usr/lib/aarch64-linux-gnu/libcudnn_static_v8.a \ /usr/lib/aarch64-linux-gnu/libcudnn_cnn_infer_static_v8.a \ /usr/lib/aarch64-linux-gnu/libnvinfer_static.a \ /usr/lib/aarch64-linux-gnu/libcudnn_adv_infer_static_v8.a \ /usr/lib/aarch64-linux-gnu/libcublas_static.a \ /usr/lib/aarch64-linux-gnu/libcudnn_adv_train_static_v8.a \ /usr/lib/aarch64-linux-gnu/libcudnn_ops_infer_static_v8.a \ /usr/lib/aarch64-linux-gnu/libcublasLt_static.a \ /usr/lib/aarch64-linux-gnu/libcudnn_cnn_train_static_v8.a \ /usr/lib/aarch64-linux-gnu/libcudnn_ops_train_static_v8.a \ /usr/lib/aarch64-linux-gnu/libmyelin_compiler_static.a \ /usr/lib/aarch64-linux-gnu/libmyelin_executor_static.a \ /usr/lib/aarch64-linux-gnu/libnvinfer_plugin_static.a && \ ldconfig # Install system-level python3 packages RUN apt-get update && apt-get install -y \ gir1.2-gst-rtsp-server-1.0 \ python3-pip \ python3-opencv \ python3-libnvinfer \ python3-scipy \ cython3 \ python3-sklearn \ python-gi-dev \ unzip && ldconfig # These system-level packages don't provide egg-info files, add them manually so that pip knows COPY docker/opencv_python-3.2.0.egg-info /usr/lib/python3/dist-packages/ COPY docker/scikit-learn-0.19.1.egg-info /usr/lib/python3/dist-packages/ # Install gst-python (python bindings for GStreamer) RUN \ 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" && \ export GST_LIBS="-lgstreamer-1.0 -lgobject-2.0 -lglib-2.0" && \ git clone https://github.com/GStreamer/gst-python.git && \ cd gst-python && git checkout 1a8f48a && \ ./autogen.sh PYTHON=python3 && \ ./configure PYTHON=python3 && \ make && make install # Install pyds (python bindings for DeepStream) RUN cd /opt/nvidia/deepstream/deepstream-5.0/lib && python3 setup.py install # Upgrade here to avoid re-running on code changes RUN pip3 install --upgrade pip # ---- Below steps are run before copying full maskcam code to allow layer caching ---- # Compile YOLOv4 plugin for DeepStream COPY deepstream_plugin_yolov4 /deepstream_plugin_yolov4 ENV CUDA_VER=10.2 RUN cd /deepstream_plugin_yolov4 && make # Get TensorRT engine (pretrained YOLOv4-tiny) # Model trained on smaller dataset # RUN wget -P / https://maskcam.s3.us-east-2.amazonaws.com/facemask_y4tiny_1024_608_fp16.trt # Model trained on bigger dataset, merged with MAFA, WiderFace, Kaggle Medical Masks and FDDB RUN wget -P / https://maskcam.s3.us-east-2.amazonaws.com/maskcam_y4t_1024_608_fp16.trt # RUN wget -P / https://maskcam.s3.us-east-2.amazonaws.com/maskcam_y4t_1120_640_fp16.trt # Install requirements with pinned versions COPY requirements.txt /maskcam_requirements.txt RUN pip3 install -r /maskcam_requirements.txt # ---- Note: all layers below this will be re-generated each time code changes ---- # Copy full maskcam code COPY . /opt/maskcam_1.0/ WORKDIR /opt/maskcam_1.0 # Move pre-copied files to their maskcam location # NOTE: Ignoring errors with `exit 0` to avoid breaking on balena livepush RUN rm -r deepstream_plugin_yolov4 && mv /deepstream_plugin_yolov4 . ; exit 0 RUN mv /*.trt yolo/ ; exit 0 # Preload library to avoids Gst errors "cannot allocate memory in static TLS block" ENV LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libgomp.so.1 # Un-pinned versions of maskcam requirements (comment pip3 install above before this) # RUN pip3 install -r requirements.in -c docker/constraints.docker RUN chmod +x docker/start.sh RUN chmod +x maskcam_run.py CMD ["docker/start.sh"] ================================================ FILE: LICENSE.md ================================================ MIT License Copyright (c) 2020, 2021 Berkeley Design Technology, Inc.. All rights reserved. Permission 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: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE 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. ================================================ FILE: README.md ================================================ # MaskCam

MaskCam 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. MaskCam 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. The 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. We 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) MaskCam 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! ## Table of contents - [Start Here!](#start-here) - [Running MaskCam from a Container on a Jetson Nano Developer Kit](#running-maskcam-from-a-container-on-a-jetson-nano-developer-kit) - [Viewing the Live Video Stream](#viewing-the-live-video-stream) - [Setting Device Configuration Parameters](#setting-device-configuration-parameters) - [MQTT Server Setup](#mqtt-and-web-server-setup) - [Running the MQTT Broker and Web Server](#running-the-mqtt-broker-and-web-server) - [Setup a Device with Your Server](#setup-a-device-with-your-server) - [Checking MQTT Connection](#checking-mqtt-connection) - [Working With the MaskCam Container](#working-with-the-maskcam-container) - [Development Mode: Manually Running MaskCam](#development-mode-manually-running-maskcam) - [Debugging: Running MaskCam Modules as Standalone Processes](#debugging-running-maskcam-modules-as-standalone-processes) - [Additional Information](#additional-information) - [Running on Jetson Nano Developer Kit Using BalenaOS](#running-on-jetson-nano-developer-kit-using-balenaos) - [Custom Container Development](#custom-container-development) - [Building From Source on Jetson Nano Developer Kit](#building-from-source-on-jetson-nano-developer-kit) - [Using Your Own Detection Model](#using-your-own-detection-model) - [Installing MaskCam Manually (Without a Container)](#installing-maskcam-manually-without-a-container) - [Running on Jetson Nano with Photon Carrier Board](#running-on-jetson-nano-with-photon-carrier-board) - [Useful Development Scripts](#useful-development-scripts) - [Troubleshooting Common Errors](#troubleshooting-common-errors) ## Start Here! ### Running MaskCam from a Container on a Jetson Nano Developer Kit The easiest and fastest way to get MaskCam running on your Jetson Nano Dev Kit is using our pre-built containers. You will need: 1. A Jetson Nano Dev Kit running JetPack 4.4.1 or 4.5 2. 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. 3. A USB webcam attached to your Nano 4. 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/). First, the MaskCam container needs to be downloaded from Docker Hub. On your Nano, run: ``` # This will take 10 minutes or more to download sudo docker pull maskcam/maskcam-beta ``` Find 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. Make 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. ``` # Connect USB camera before running this! sudo docker run --runtime nvidia --privileged --rm -it --env MASKCAM_DEVICE_ADDRESS= -p 1883:1883 -p 8080:8080 -p 8554:8554 maskcam/maskcam-beta ``` The 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. Otherwise, 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! ### Viewing the Live Video Stream If you scroll through the logs and don't see any errors, you should find a message like: ```Streaming at rtsp://aaa.bbb.ccc.ddd:8554/maskcam``` where `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. You 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.

This 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). If you encounter any errors running the live stream, check the [Troubleshooting](#troubleshooting-common-errors) section for tips on resolving errors. ### Setting Device Configuration Parameters MaskCam 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). This 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: ``` # Run with MASKCAM_INPUT and MASKCAM_DEVICE_ADDRESS sudo 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 ``` Another 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. As 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 two addtional environment variables, `MQTT_BROKER_IP` and `MQTT_DEVICE_NAME`. This allows your device to find the MQTT server and identify itself: ``` # Run with MQTT_BROKER_IP, MQTT_DEVICE_NAME, and MASKCAM_DEVICE_ADDRESS sudo 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 ``` *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.* ## MQTT and Web Server Setup ### Running the MQTT Broker and Web Server MaskCam 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.

You 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. The 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. After installing docker-compose, clone this repo: ``` git clone https://github.com/bdtinc/maskcam.git ``` Go to the `server/` folder, which has all the needed components implemented on four containers: the Mosquitto broker, backend API, database, and Streamlit frontend. These containers are configured using environment variables, so create the `.env` files by copying the default templates: ``` cd server cp database.env.template database.env cp frontend.env.template frontend.env cp backend.env.template backend.env ``` The 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: ``` POSTGRES_USER=postgres POSTGRES_PASSWORD=some_password POSTGRES_DB=maskcam ``` *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. After editing the database environment file, you're ready to build all the containers and run them with a single command: ``` sudo docker-compose up -d ``` Wait 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). Next, open a web browser and enter the server IP to visit the frontend webpage: ``` http://:8501/ ``` If 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. *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. ### Setup a Device With Your Server Once 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.) ``` # Run with MQTT_BROKER_IP, MQTT_DEVICE_NAME, and MASKCAM_DEVICE_ADDRESS sudo 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 ``` And 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. Check the next section if the MQTT connection is not established from the device to the server. ### Checking MQTT Connection If you're running the MQTT broker on a machine in your local network, make sure its IP is accessible from the Jetson device: ``` ping ``` *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...` If you're setting up a remote server and using its public IP to connect from your device, chances are you're not setting properly the port `1883` to be opened for inbound and outbound traffic. If you want to check the port is correctly configured, use `nc` from a local machine or your jetson: ``` nc -vz 1883 ``` Remember 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). ## Working With the MaskCam Container ### Development Mode: Manually Running MaskCam If you want to play around with the code, you probably don't want the container to automatically start running the `maskcam_run.py` script. The easiest way to achieve that, is by defining the environment variable `DEV_MODE=1`: ``` docker run --runtime nvidia --privileged --rm -it --env DEV_MODE=1 -p 1883:1883 -p 8080:8080 -p 8554:8554 maskcam/maskcam-beta ``` This 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 of its sub-modules as standalone processes: ``` # e.g: Run with a different input instead of default `/dev/video0` ./maskcam_run.py v4l2:///dev/video1 # e.g: Disable tracker to visualize raw detections and scores MASKCAM_DISABLE_TRACKER=1 ./maskcam_run.py ``` ### Debugging: Running MaskCam Modules as Standalone Processes The script `maskcam_run.py`, which is the main entrypoint for the MaskCam software, has two roles: - Handles all the MQTT communication (send stats and receive commands) - Orchestrates all other processes that live under `maskcam/maskcam_*.py`. But you can actually run any of those modules as standalone processes, which can be easier for debugging. You 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: ``` # e.g: Run only the static file server process python3 -m maskcam.maskcam_fileserver # e.g: Serve another directory to test python3 -m maskcam.maskcam_fileserver /tmp # e.g: Run only the inference and streaming processes python3 -m maskcam.maskcam_streaming & # Hit enter until you get a prompt and then: python3 -m maskcam.maskcam_inference ``` **Note:** In the last example, `maskcam_streaming` is running on background, so it will not terminate if you press `Ctrl+C` (only `maskcam_inference` will, since it's running on the foreground). To check that the streaming is still running and then bring it to foreground to terminate it, run: ``` jobs fg %1 # Now you can hit Ctrl+C to terminate streaming ``` ## Additional Information Further 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. ### Running on Jetson Nano Developer Kit Using BalenaOS [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). ### Custom Container Development MaskCam 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. #### Building From Source on Jetson Nano Developer Kit Please 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. #### Using Your Own Detection Model Please 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. ### Installing MaskCam Manually (Without a Container) MaskCam 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). ### Running on Jetson Nano with Photon Carrier Board For 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). ### Useful Development Scripts During 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). ## Troubleshooting Common Errors If you run into any errors or issues while working with MaskCam, this section gives common errors and their solutions. MaskCam 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. #### Error: camera not connected/not recognized If 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. #### Error: not running in privileged mode In this case, you'll see a bunch of annoying messages like: ``` Error: Can't initialize nvrm channel Couldn't create ddkvic Session: Cannot allocate memory nvbuf_utils: Could not create Default NvBufferSession ``` You'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). #### Error: reason not negotiated/camera capabilities If you get an error that looks like: `v4l-camera-source / reason not-negotiated` Then 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 ...` Find 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)). #### Error: Streaming or file server are not accessible (nothing else seems to fail) Make 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: ``` fileserver-port=8080 streaming-port=8554 mqtt-broker-port=1883 ``` These 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. #### Other Errors Sometimes 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. ## Questions? Need Help? Email us at maskcam@bdti.com, and be sure to check out our [independent report on the development of MaskCam](https://bdti.com/maskcam)! ================================================ FILE: deepstream_plugin_yolov4/Makefile ================================================ ################################################################################ # Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ CUDA_VER?= ifeq ($(CUDA_VER),) $(error "CUDA_VER is not set") endif CC:= g++ NVCC:=/usr/local/cuda-$(CUDA_VER)/bin/nvcc CFLAGS:= -Wall -std=c++11 -shared -fPIC -Wno-error=deprecated-declarations CFLAGS+= -I../../includes -I/usr/local/cuda-$(CUDA_VER)/include -I/opt/nvidia/deepstream/deepstream-5.0/sources/includes LIBS:= -lnvinfer_plugin -lnvinfer -lnvparsers -L/usr/local/cuda-$(CUDA_VER)/lib64 -lcudart -lcublas -lstdc++fs -L/opt/nvidia/deepstream/deepstream-5.0/lib LFLAGS:= -shared -Wl,--start-group $(LIBS) -Wl,--end-group INCS:= $(wildcard *.h) SRCFILES:= nvdsinfer_yolo_engine.cpp \ nvdsparsebbox_Yolo.cpp \ yoloPlugins.cpp \ trt_utils.cpp \ yolo.cpp \ kernels.cu TARGET_LIB:= libnvdsinfer_custom_impl_Yolo.so TARGET_OBJS:= $(SRCFILES:.cpp=.o) TARGET_OBJS:= $(TARGET_OBJS:.cu=.o) all: $(TARGET_LIB) %.o: %.cpp $(INCS) Makefile $(CC) -c -o $@ $(CFLAGS) $< %.o: %.cu $(INCS) Makefile $(NVCC) -c -o $@ --compiler-options '-fPIC' $< $(TARGET_LIB) : $(TARGET_OBJS) $(CC) -o $@ $(TARGET_OBJS) $(LFLAGS) clean: rm -rf $(TARGET_LIB) ================================================ FILE: deepstream_plugin_yolov4/README.md ================================================ ## YOLOv4 plugin for DeepStream This plugin was obtained from: https://github.com/Tianxiaomo/pytorch-YOLOv4/tree/master/DeepStream It must be compiled locally on the jetson device: ``` export CUDA_VER=10.2 make ``` ================================================ FILE: deepstream_plugin_yolov4/kernels.cu ================================================ /* * Copyright (c) 2018-2019 NVIDIA Corporation. All rights reserved. * * NVIDIA Corporation and its licensors retain all intellectual property * and proprietary rights in and to this software, related documentation * and any modifications thereto. Any use, reproduction, disclosure or * distribution of this software and related documentation without an express * license agreement from NVIDIA Corporation is strictly prohibited. * */ #include #include #include #include #include inline __device__ float sigmoidGPU(const float& x) { return 1.0f / (1.0f + __expf(-x)); } __global__ void gpuYoloLayerV3(const float* input, float* output, const uint gridSize, const uint numOutputClasses, const uint numBBoxes) { uint x_id = blockIdx.x * blockDim.x + threadIdx.x; uint y_id = blockIdx.y * blockDim.y + threadIdx.y; uint z_id = blockIdx.z * blockDim.z + threadIdx.z; if ((x_id >= gridSize) || (y_id >= gridSize) || (z_id >= numBBoxes)) { return; } const int numGridCells = gridSize * gridSize; const int bbindex = y_id * gridSize + x_id; output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)] = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)]); output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)] = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)]); output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)] = __expf(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)]); output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)] = __expf(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)]); output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)] = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)]); for (uint i = 0; i < numOutputClasses; ++i) { output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))] = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))]); } } cudaError_t cudaYoloLayerV3(const void* input, void* output, const uint& batchSize, const uint& gridSize, const uint& numOutputClasses, const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream); cudaError_t cudaYoloLayerV3(const void* input, void* output, const uint& batchSize, const uint& gridSize, const uint& numOutputClasses, const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream) { dim3 threads_per_block(16, 16, 4); dim3 number_of_blocks((gridSize / threads_per_block.x) + 1, (gridSize / threads_per_block.y) + 1, (numBBoxes / threads_per_block.z) + 1); for (unsigned int batch = 0; batch < batchSize; ++batch) { gpuYoloLayerV3<<>>( reinterpret_cast(input) + (batch * outputSize), reinterpret_cast(output) + (batch * outputSize), gridSize, numOutputClasses, numBBoxes); } return cudaGetLastError(); } ================================================ FILE: deepstream_plugin_yolov4/nvdsinfer_yolo_engine.cpp ================================================ /* * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. * * Permission 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: * * The above copyright notice and this permission notice shall be included in * all copies or substantial portions of the Software. * * THE 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. */ #include "nvdsinfer_custom_impl.h" #include "nvdsinfer_context.h" #include "yoloPlugins.h" #include "yolo.h" #include #define USE_CUDA_ENGINE_GET_API 1 static bool getYoloNetworkInfo (NetworkInfo &networkInfo, const NvDsInferContextInitParams* initParams) { std::string yoloCfg = initParams->customNetworkConfigFilePath; std::string yoloType; std::transform (yoloCfg.begin(), yoloCfg.end(), yoloCfg.begin(), [] (uint8_t c) { return std::tolower (c);}); if (yoloCfg.find("yolov2") != std::string::npos) { if (yoloCfg.find("yolov2-tiny") != std::string::npos) yoloType = "yolov2-tiny"; else yoloType = "yolov2"; } else if (yoloCfg.find("yolov3") != std::string::npos) { if (yoloCfg.find("yolov3-tiny") != std::string::npos) yoloType = "yolov3-tiny"; else yoloType = "yolov3"; } else { std::cerr << "Yolo type is not defined from config file name:" << yoloCfg << std::endl; return false; } networkInfo.networkType = yoloType; networkInfo.configFilePath = initParams->customNetworkConfigFilePath; networkInfo.wtsFilePath = initParams->modelFilePath; networkInfo.deviceType = (initParams->useDLA ? "kDLA" : "kGPU"); networkInfo.inputBlobName = "data"; if (networkInfo.configFilePath.empty() || networkInfo.wtsFilePath.empty()) { std::cerr << "Yolo config file or weights file is NOT specified." << std::endl; return false; } if (!fileExists(networkInfo.configFilePath) || !fileExists(networkInfo.wtsFilePath)) { std::cerr << "Yolo config file or weights file is NOT exist." << std::endl; return false; } return true; } #if !USE_CUDA_ENGINE_GET_API IModelParser* NvDsInferCreateModelParser( const NvDsInferContextInitParams* initParams) { NetworkInfo networkInfo; if (!getYoloNetworkInfo(networkInfo, initParams)) { return nullptr; } return new Yolo(networkInfo); } #else extern "C" bool NvDsInferYoloCudaEngineGet(nvinfer1::IBuilder * const builder, const NvDsInferContextInitParams * const initParams, nvinfer1::DataType dataType, nvinfer1::ICudaEngine *& cudaEngine); extern "C" bool NvDsInferYoloCudaEngineGet(nvinfer1::IBuilder * const builder, const NvDsInferContextInitParams * const initParams, nvinfer1::DataType dataType, nvinfer1::ICudaEngine *& cudaEngine) { NetworkInfo networkInfo; if (!getYoloNetworkInfo(networkInfo, initParams)) { return false; } Yolo yolo(networkInfo); cudaEngine = yolo.createEngine (builder); if (cudaEngine == nullptr) { std::cerr << "Failed to build cuda engine on " << networkInfo.configFilePath << std::endl; return false; } return true; } #endif ================================================ FILE: deepstream_plugin_yolov4/nvdsparsebbox_Yolo.cpp ================================================ /* * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. * * Permission 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: * * The above copyright notice and this permission notice shall be included in * all copies or substantial portions of the Software. * * THE 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. */ #include #include #include #include #include #include #include #include "nvdsinfer_custom_impl.h" #include "trt_utils.h" static const int NUM_CLASSES_YOLO = 4; // This is just for checkings, keep updated by hand extern "C" bool NvDsInferParseCustomYoloV3( std::vector const &outputLayersInfo, NvDsInferNetworkInfo const &networkInfo, NvDsInferParseDetectionParams const &detectionParams, std::vector &objectList); extern "C" bool NvDsInferParseCustomYoloV3Tiny( std::vector const &outputLayersInfo, NvDsInferNetworkInfo const &networkInfo, NvDsInferParseDetectionParams const &detectionParams, std::vector &objectList); extern "C" bool NvDsInferParseCustomYoloV2( std::vector const &outputLayersInfo, NvDsInferNetworkInfo const &networkInfo, NvDsInferParseDetectionParams const &detectionParams, std::vector &objectList); extern "C" bool NvDsInferParseCustomYoloV2Tiny( std::vector const &outputLayersInfo, NvDsInferNetworkInfo const &networkInfo, NvDsInferParseDetectionParams const &detectionParams, std::vector &objectList); extern "C" bool NvDsInferParseCustomYoloTLT( std::vector const &outputLayersInfo, NvDsInferNetworkInfo const &networkInfo, NvDsInferParseDetectionParams const &detectionParams, std::vector &objectList); extern "C" bool NvDsInferParseCustomYoloV4( std::vector const &outputLayersInfo, NvDsInferNetworkInfo const &networkInfo, NvDsInferParseDetectionParams const &detectionParams, std::vector &objectList); /* This is a sample bounding box parsing function for the sample YoloV3 detector model */ static NvDsInferParseObjectInfo convertBBox(const float &bx, const float &by, const float &bw, const float &bh, const int &stride, const uint &netW, const uint &netH) { NvDsInferParseObjectInfo b; // Restore coordinates to network input resolution float xCenter = bx * stride; float yCenter = by * stride; float x0 = xCenter - bw / 2; float y0 = yCenter - bh / 2; float x1 = x0 + bw; float y1 = y0 + bh; x0 = clamp(x0, 0, netW); y0 = clamp(y0, 0, netH); x1 = clamp(x1, 0, netW); y1 = clamp(y1, 0, netH); b.left = x0; b.width = clamp(x1 - x0, 0, netW); b.top = y0; b.height = clamp(y1 - y0, 0, netH); return b; } static void addBBoxProposal(const float bx, const float by, const float bw, const float bh, const uint stride, const uint &netW, const uint &netH, const int maxIndex, const float maxProb, std::vector &binfo) { NvDsInferParseObjectInfo bbi = convertBBox(bx, by, bw, bh, stride, netW, netH); if (bbi.width < 1 || bbi.height < 1) return; bbi.detectionConfidence = maxProb; bbi.classId = maxIndex; binfo.push_back(bbi); } static std::vector decodeYoloV2Tensor( const float *detections, const std::vector &anchors, const uint gridSizeW, const uint gridSizeH, const uint stride, const uint numBBoxes, const uint numOutputClasses, const uint &netW, const uint &netH) { std::vector binfo; for (uint y = 0; y < gridSizeH; ++y) { for (uint x = 0; x < gridSizeW; ++x) { for (uint b = 0; b < numBBoxes; ++b) { const float pw = anchors[b * 2]; const float ph = anchors[b * 2 + 1]; const int numGridCells = gridSizeH * gridSizeW; const int bbindex = y * gridSizeW + x; const float bx = x + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 0)]; const float by = y + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 1)]; const float bw = pw * exp(detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 2)]); const float bh = ph * exp(detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 3)]); const float objectness = detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 4)]; float maxProb = 0.0f; int maxIndex = -1; for (uint i = 0; i < numOutputClasses; ++i) { float prob = (detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + (5 + i))]); if (prob > maxProb) { maxProb = prob; maxIndex = i; } } maxProb = objectness * maxProb; addBBoxProposal(bx, by, bw, bh, stride, netW, netH, maxIndex, maxProb, binfo); } } } return binfo; } static std::vector decodeYoloV3Tensor( const float *detections, const std::vector &mask, const std::vector &anchors, const uint gridSizeW, const uint gridSizeH, const uint stride, const uint numBBoxes, const uint numOutputClasses, const uint &netW, const uint &netH) { std::vector binfo; for (uint y = 0; y < gridSizeH; ++y) { for (uint x = 0; x < gridSizeW; ++x) { for (uint b = 0; b < numBBoxes; ++b) { const float pw = anchors[mask[b] * 2]; const float ph = anchors[mask[b] * 2 + 1]; const int numGridCells = gridSizeH * gridSizeW; const int bbindex = y * gridSizeW + x; const float bx = x + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 0)]; const float by = y + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 1)]; const float bw = pw * detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 2)]; const float bh = ph * detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 3)]; const float objectness = detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 4)]; float maxProb = 0.0f; int maxIndex = -1; for (uint i = 0; i < numOutputClasses; ++i) { float prob = (detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + (5 + i))]); if (prob > maxProb) { maxProb = prob; maxIndex = i; } } maxProb = objectness * maxProb; addBBoxProposal(bx, by, bw, bh, stride, netW, netH, maxIndex, maxProb, binfo); } } } return binfo; } static inline std::vector SortLayers(const std::vector &outputLayersInfo) { std::vector outLayers; for (auto const &layer : outputLayersInfo) { outLayers.push_back(&layer); } std::sort(outLayers.begin(), outLayers.end(), [](const NvDsInferLayerInfo *a, const NvDsInferLayerInfo *b) { return a->inferDims.d[1] < b->inferDims.d[1]; }); return outLayers; } static bool NvDsInferParseYoloV3( std::vector const &outputLayersInfo, NvDsInferNetworkInfo const &networkInfo, NvDsInferParseDetectionParams const &detectionParams, std::vector &objectList, const std::vector &anchors, const std::vector> &masks) { const uint kNUM_BBOXES = 3; const std::vector sortedLayers = SortLayers(outputLayersInfo); if (sortedLayers.size() != masks.size()) { std::cerr << "ERROR: yoloV3 output layer.size: " << sortedLayers.size() << " does not match mask.size: " << masks.size() << std::endl; return false; } if (NUM_CLASSES_YOLO != detectionParams.numClassesConfigured) { std::cerr << "WARNING: Num classes mismatch. Configured:" << detectionParams.numClassesConfigured << ", detected by network: " << NUM_CLASSES_YOLO << std::endl; } std::vector objects; for (uint idx = 0; idx < masks.size(); ++idx) { const NvDsInferLayerInfo &layer = *sortedLayers[idx]; // 255 x Grid x Grid assert(layer.inferDims.numDims == 3); const uint gridSizeH = layer.inferDims.d[1]; const uint gridSizeW = layer.inferDims.d[2]; const uint stride = DIVUP(networkInfo.width, gridSizeW); assert(stride == DIVUP(networkInfo.height, gridSizeH)); std::vector outObjs = decodeYoloV3Tensor((const float *)(layer.buffer), masks[idx], anchors, gridSizeW, gridSizeH, stride, kNUM_BBOXES, NUM_CLASSES_YOLO, networkInfo.width, networkInfo.height); objects.insert(objects.end(), outObjs.begin(), outObjs.end()); } objectList = objects; return true; } static NvDsInferParseObjectInfo convertBBoxYoloV4(const float &bx1, const float &by1, const float &bx2, const float &by2, const uint &netW, const uint &netH) { NvDsInferParseObjectInfo b; // Restore coordinates to network input resolution float x1 = bx1 * netW; float y1 = by1 * netH; float x2 = bx2 * netW; float y2 = by2 * netH; x1 = clamp(x1, 0, netW); y1 = clamp(y1, 0, netH); x2 = clamp(x2, 0, netW); y2 = clamp(y2, 0, netH); b.left = x1; b.width = clamp(x2 - x1, 0, netW); b.top = y1; b.height = clamp(y2 - y1, 0, netH); return b; } static void addBBoxProposalYoloV4(const float bx, const float by, const float bw, const float bh, const uint &netW, const uint &netH, const int maxIndex, const float maxProb, std::vector &binfo) { NvDsInferParseObjectInfo bbi = convertBBoxYoloV4(bx, by, bw, bh, netW, netH); if (bbi.width < 1 || bbi.height < 1) return; bbi.detectionConfidence = maxProb; bbi.classId = maxIndex; binfo.push_back(bbi); } static std::vector decodeYoloV4Tensor( const float *boxes, const float *scores, const uint num_bboxes, NvDsInferParseDetectionParams const &detectionParams, const uint &netW, const uint &netH) { std::vector binfo; uint bbox_location = 0; uint score_location = 0; for (uint b = 0; b < num_bboxes; ++b) { float bx1 = boxes[bbox_location]; float by1 = boxes[bbox_location + 1]; float bx2 = boxes[bbox_location + 2]; float by2 = boxes[bbox_location + 3]; float maxProb = 0.0f; int maxIndex = -1; for (uint c = 0; c < detectionParams.numClassesConfigured; ++c) { float prob = scores[score_location + c]; if (prob > maxProb) { maxProb = prob; maxIndex = c; } } if (maxProb > detectionParams.perClassPreclusterThreshold[maxIndex]) { addBBoxProposalYoloV4(bx1, by1, bx2, by2, netW, netH, maxIndex, maxProb, binfo); } bbox_location += 4; score_location += detectionParams.numClassesConfigured; } return binfo; } /* C-linkage to prevent name-mangling */ static bool NvDsInferParseYoloV4( std::vector const &outputLayersInfo, NvDsInferNetworkInfo const &networkInfo, NvDsInferParseDetectionParams const &detectionParams, std::vector &objectList) { if (NUM_CLASSES_YOLO != detectionParams.numClassesConfigured) { std::cerr << "WARNING: Num classes mismatch. Configured:" << detectionParams.numClassesConfigured << ", detected by network: " << NUM_CLASSES_YOLO << std::endl; } std::vector objects; const NvDsInferLayerInfo &boxes = outputLayersInfo[0]; // num_boxes x 4 const NvDsInferLayerInfo &scores = outputLayersInfo[1]; // num_boxes x num_classes const NvDsInferLayerInfo &subbox = outputLayersInfo[2]; //* printf("%d\n", subbox.inferDims.numDims); // 3 dimensional: [num_boxes, 1, 4] assert(boxes.inferDims.numDims == 3); // 2 dimensional: [num_boxes, num_classes] assert(scores.inferDims.numDims == 2); // The second dimension should be num_classes assert(detectionParams.numClassesConfigured == scores.inferDims.d[1]); uint num_bboxes = boxes.inferDims.d[0]; // std::cout << "Network Info: " << networkInfo.height << " " << networkInfo.width << std::endl; std::vector outObjs = decodeYoloV4Tensor( (const float *)(boxes.buffer), (const float *)(scores.buffer), num_bboxes, detectionParams, networkInfo.width, networkInfo.height); objects.insert(objects.end(), outObjs.begin(), outObjs.end()); objectList = objects; return true; } extern "C" bool NvDsInferParseCustomYoloV4( std::vector const &outputLayersInfo, NvDsInferNetworkInfo const &networkInfo, NvDsInferParseDetectionParams const &detectionParams, std::vector &objectList) { return NvDsInferParseYoloV4( outputLayersInfo, networkInfo, detectionParams, objectList); } extern "C" bool NvDsInferParseCustomYoloV3( std::vector const &outputLayersInfo, NvDsInferNetworkInfo const &networkInfo, NvDsInferParseDetectionParams const &detectionParams, std::vector &objectList) { static const std::vector kANCHORS = { 10.0, 13.0, 16.0, 30.0, 33.0, 23.0, 30.0, 61.0, 62.0, 45.0, 59.0, 119.0, 116.0, 90.0, 156.0, 198.0, 373.0, 326.0}; static const std::vector> kMASKS = { {6, 7, 8}, {3, 4, 5}, {0, 1, 2}}; return NvDsInferParseYoloV3( outputLayersInfo, networkInfo, detectionParams, objectList, kANCHORS, kMASKS); } extern "C" bool NvDsInferParseCustomYoloV3Tiny( std::vector const &outputLayersInfo, NvDsInferNetworkInfo const &networkInfo, NvDsInferParseDetectionParams const &detectionParams, std::vector &objectList) { static const std::vector kANCHORS = { 10, 14, 23, 27, 37, 58, 81, 82, 135, 169, 344, 319}; static const std::vector> kMASKS = { {3, 4, 5}, //{0, 1, 2}}; // as per output result, select {1,2,3} {1, 2, 3}}; return NvDsInferParseYoloV3( outputLayersInfo, networkInfo, detectionParams, objectList, kANCHORS, kMASKS); } static bool NvDsInferParseYoloV2( std::vector const &outputLayersInfo, NvDsInferNetworkInfo const &networkInfo, NvDsInferParseDetectionParams const &detectionParams, std::vector &objectList) { // copy anchor data from yolov2.cfg file std::vector anchors = {0.57273, 0.677385, 1.87446, 2.06253, 3.33843, 5.47434, 7.88282, 3.52778, 9.77052, 9.16828}; const uint kNUM_BBOXES = 5; if (outputLayersInfo.empty()) { std::cerr << "Could not find output layer in bbox parsing" << std::endl; ; return false; } const NvDsInferLayerInfo &layer = outputLayersInfo[0]; if (NUM_CLASSES_YOLO != detectionParams.numClassesConfigured) { std::cerr << "WARNING: Num classes mismatch. Configured:" << detectionParams.numClassesConfigured << ", detected by network: " << NUM_CLASSES_YOLO << std::endl; } assert(layer.inferDims.numDims == 3); const uint gridSizeH = layer.inferDims.d[1]; const uint gridSizeW = layer.inferDims.d[2]; const uint stride = DIVUP(networkInfo.width, gridSizeW); assert(stride == DIVUP(networkInfo.height, gridSizeH)); for (auto &anchor : anchors) { anchor *= stride; } std::vector objects = decodeYoloV2Tensor((const float *)(layer.buffer), anchors, gridSizeW, gridSizeH, stride, kNUM_BBOXES, NUM_CLASSES_YOLO, networkInfo.width, networkInfo.height); objectList = objects; return true; } extern "C" bool NvDsInferParseCustomYoloV2( std::vector const &outputLayersInfo, NvDsInferNetworkInfo const &networkInfo, NvDsInferParseDetectionParams const &detectionParams, std::vector &objectList) { return NvDsInferParseYoloV2( outputLayersInfo, networkInfo, detectionParams, objectList); } extern "C" bool NvDsInferParseCustomYoloV2Tiny( std::vector const &outputLayersInfo, NvDsInferNetworkInfo const &networkInfo, NvDsInferParseDetectionParams const &detectionParams, std::vector &objectList) { return NvDsInferParseYoloV2( outputLayersInfo, networkInfo, detectionParams, objectList); } extern "C" bool NvDsInferParseCustomYoloTLT( std::vector const &outputLayersInfo, NvDsInferNetworkInfo const &networkInfo, NvDsInferParseDetectionParams const &detectionParams, std::vector &objectList) { if (outputLayersInfo.size() != 4) { std::cerr << "Mismatch in the number of output buffers." << "Expected 4 output buffers, detected in the network :" << outputLayersInfo.size() << std::endl; return false; } const int topK = 200; const int *keepCount = static_cast(outputLayersInfo.at(0).buffer); const float *boxes = static_cast(outputLayersInfo.at(1).buffer); const float *scores = static_cast(outputLayersInfo.at(2).buffer); const float *cls = static_cast(outputLayersInfo.at(3).buffer); for (int i = 0; (i < keepCount[0]) && (objectList.size() <= topK); ++i) { const float *loc = &boxes[0] + (i * 4); const float *conf = &scores[0] + i; const float *cls_id = &cls[0] + i; if (conf[0] > 1.001) continue; if ((loc[0] < 0) || (loc[1] < 0) || (loc[2] < 0) || (loc[3] < 0)) continue; if ((loc[0] > networkInfo.width) || (loc[2] > networkInfo.width) || (loc[1] > networkInfo.height) || (loc[3] > networkInfo.width)) continue; if ((loc[2] < loc[0]) || (loc[3] < loc[1])) continue; if (((loc[3] - loc[1]) > networkInfo.height) || ((loc[2] - loc[0]) > networkInfo.width)) continue; NvDsInferParseObjectInfo curObj{static_cast(cls_id[0]), loc[0], loc[1], (loc[2] - loc[0]), (loc[3] - loc[1]), conf[0]}; objectList.push_back(curObj); } return true; } /* Check that the custom function has been defined correctly */ CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV4); CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV3); CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV3Tiny); CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV2); CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV2Tiny); CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloTLT); ================================================ FILE: deepstream_plugin_yolov4/trt_utils.cpp ================================================ /* * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. * * Permission 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: * * The above copyright notice and this permission notice shall be included in * all copies or substantial portions of the Software. * * THE 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. */ #include "trt_utils.h" #include #include #include #include #include #include #include "NvInferPlugin.h" static void leftTrim(std::string& s) { s.erase(s.begin(), std::find_if(s.begin(), s.end(), [](int ch) { return !isspace(ch); })); } static void rightTrim(std::string& s) { s.erase(std::find_if(s.rbegin(), s.rend(), [](int ch) { return !isspace(ch); }).base(), s.end()); } std::string trim(std::string s) { leftTrim(s); rightTrim(s); return s; } float clamp(const float val, const float minVal, const float maxVal) { assert(minVal <= maxVal); return std::min(maxVal, std::max(minVal, val)); } bool fileExists(const std::string fileName, bool verbose) { if (!std::experimental::filesystem::exists(std::experimental::filesystem::path(fileName))) { if (verbose) std::cout << "File does not exist : " << fileName << std::endl; return false; } return true; } std::vector loadWeights(const std::string weightsFilePath, const std::string& networkType) { assert(fileExists(weightsFilePath)); std::cout << "Loading pre-trained weights..." << std::endl; std::ifstream file(weightsFilePath, std::ios_base::binary); assert(file.good()); std::string line; if (networkType == "yolov2") { // Remove 4 int32 bytes of data from the stream belonging to the header file.ignore(4 * 4); } else if ((networkType == "yolov3") || (networkType == "yolov3-tiny") || (networkType == "yolov2-tiny")) { // Remove 5 int32 bytes of data from the stream belonging to the header file.ignore(4 * 5); } else { std::cout << "Invalid network type" << std::endl; assert(0); } std::vector weights; char floatWeight[4]; while (!file.eof()) { file.read(floatWeight, 4); assert(file.gcount() == 4); weights.push_back(*reinterpret_cast(floatWeight)); if (file.peek() == std::istream::traits_type::eof()) break; } std::cout << "Loading weights of " << networkType << " complete!" << std::endl; std::cout << "Total Number of weights read : " << weights.size() << std::endl; return weights; } std::string dimsToString(const nvinfer1::Dims d) { std::stringstream s; assert(d.nbDims >= 1); for (int i = 0; i < d.nbDims - 1; ++i) { s << std::setw(4) << d.d[i] << " x"; } s << std::setw(4) << d.d[d.nbDims - 1]; return s.str(); } void displayDimType(const nvinfer1::Dims d) { std::cout << "(" << d.nbDims << ") "; for (int i = 0; i < d.nbDims; ++i) { switch (d.type[i]) { case nvinfer1::DimensionType::kSPATIAL: std::cout << "kSPATIAL "; break; case nvinfer1::DimensionType::kCHANNEL: std::cout << "kCHANNEL "; break; case nvinfer1::DimensionType::kINDEX: std::cout << "kINDEX "; break; case nvinfer1::DimensionType::kSEQUENCE: std::cout << "kSEQUENCE "; break; } } std::cout << std::endl; } int getNumChannels(nvinfer1::ITensor* t) { nvinfer1::Dims d = t->getDimensions(); assert(d.nbDims == 3); return d.d[0]; } uint64_t get3DTensorVolume(nvinfer1::Dims inputDims) { assert(inputDims.nbDims == 3); return inputDims.d[0] * inputDims.d[1] * inputDims.d[2]; } nvinfer1::ILayer* netAddMaxpool(int layerIdx, std::map& block, nvinfer1::ITensor* input, nvinfer1::INetworkDefinition* network) { assert(block.at("type") == "maxpool"); assert(block.find("size") != block.end()); assert(block.find("stride") != block.end()); int size = std::stoi(block.at("size")); int stride = std::stoi(block.at("stride")); nvinfer1::IPoolingLayer* pool = network->addPooling(*input, nvinfer1::PoolingType::kMAX, nvinfer1::DimsHW{size, size}); assert(pool); std::string maxpoolLayerName = "maxpool_" + std::to_string(layerIdx); pool->setStride(nvinfer1::DimsHW{stride, stride}); pool->setPaddingMode(nvinfer1::PaddingMode::kSAME_UPPER); pool->setName(maxpoolLayerName.c_str()); return pool; } nvinfer1::ILayer* netAddConvLinear(int layerIdx, std::map& block, std::vector& weights, std::vector& trtWeights, int& weightPtr, int& inputChannels, nvinfer1::ITensor* input, nvinfer1::INetworkDefinition* network) { assert(block.at("type") == "convolutional"); assert(block.find("batch_normalize") == block.end()); assert(block.at("activation") == "linear"); assert(block.find("filters") != block.end()); assert(block.find("pad") != block.end()); assert(block.find("size") != block.end()); assert(block.find("stride") != block.end()); int filters = std::stoi(block.at("filters")); int padding = std::stoi(block.at("pad")); int kernelSize = std::stoi(block.at("size")); int stride = std::stoi(block.at("stride")); int pad; if (padding) pad = (kernelSize - 1) / 2; else pad = 0; // load the convolution layer bias nvinfer1::Weights convBias{nvinfer1::DataType::kFLOAT, nullptr, filters}; float* val = new float[filters]; for (int i = 0; i < filters; ++i) { val[i] = weights[weightPtr]; weightPtr++; } convBias.values = val; trtWeights.push_back(convBias); // load the convolutional layer weights int size = filters * inputChannels * kernelSize * kernelSize; nvinfer1::Weights convWt{nvinfer1::DataType::kFLOAT, nullptr, size}; val = new float[size]; for (int i = 0; i < size; ++i) { val[i] = weights[weightPtr]; weightPtr++; } convWt.values = val; trtWeights.push_back(convWt); nvinfer1::IConvolutionLayer* conv = network->addConvolution( *input, filters, nvinfer1::DimsHW{kernelSize, kernelSize}, convWt, convBias); assert(conv != nullptr); std::string convLayerName = "conv_" + std::to_string(layerIdx); conv->setName(convLayerName.c_str()); conv->setStride(nvinfer1::DimsHW{stride, stride}); conv->setPadding(nvinfer1::DimsHW{pad, pad}); return conv; } nvinfer1::ILayer* netAddConvBNLeaky(int layerIdx, std::map& block, std::vector& weights, std::vector& trtWeights, int& weightPtr, int& inputChannels, nvinfer1::ITensor* input, nvinfer1::INetworkDefinition* network) { assert(block.at("type") == "convolutional"); assert(block.find("batch_normalize") != block.end()); assert(block.at("batch_normalize") == "1"); assert(block.at("activation") == "leaky"); assert(block.find("filters") != block.end()); assert(block.find("pad") != block.end()); assert(block.find("size") != block.end()); assert(block.find("stride") != block.end()); bool batchNormalize, bias; if (block.find("batch_normalize") != block.end()) { batchNormalize = (block.at("batch_normalize") == "1"); bias = false; } else { batchNormalize = false; bias = true; } // all conv_bn_leaky layers assume bias is false assert(batchNormalize == true && bias == false); UNUSED(batchNormalize); UNUSED(bias); int filters = std::stoi(block.at("filters")); int padding = std::stoi(block.at("pad")); int kernelSize = std::stoi(block.at("size")); int stride = std::stoi(block.at("stride")); int pad; if (padding) pad = (kernelSize - 1) / 2; else pad = 0; /***** CONVOLUTION LAYER *****/ /*****************************/ // batch norm weights are before the conv layer // load BN biases (bn_biases) std::vector bnBiases; for (int i = 0; i < filters; ++i) { bnBiases.push_back(weights[weightPtr]); weightPtr++; } // load BN weights std::vector bnWeights; for (int i = 0; i < filters; ++i) { bnWeights.push_back(weights[weightPtr]); weightPtr++; } // load BN running_mean std::vector bnRunningMean; for (int i = 0; i < filters; ++i) { bnRunningMean.push_back(weights[weightPtr]); weightPtr++; } // load BN running_var std::vector bnRunningVar; for (int i = 0; i < filters; ++i) { // 1e-05 for numerical stability bnRunningVar.push_back(sqrt(weights[weightPtr] + 1.0e-5)); weightPtr++; } // load Conv layer weights (GKCRS) int size = filters * inputChannels * kernelSize * kernelSize; nvinfer1::Weights convWt{nvinfer1::DataType::kFLOAT, nullptr, size}; float* val = new float[size]; for (int i = 0; i < size; ++i) { val[i] = weights[weightPtr]; weightPtr++; } convWt.values = val; trtWeights.push_back(convWt); nvinfer1::Weights convBias{nvinfer1::DataType::kFLOAT, nullptr, 0}; trtWeights.push_back(convBias); nvinfer1::IConvolutionLayer* conv = network->addConvolution( *input, filters, nvinfer1::DimsHW{kernelSize, kernelSize}, convWt, convBias); assert(conv != nullptr); std::string convLayerName = "conv_" + std::to_string(layerIdx); conv->setName(convLayerName.c_str()); conv->setStride(nvinfer1::DimsHW{stride, stride}); conv->setPadding(nvinfer1::DimsHW{pad, pad}); /***** BATCHNORM LAYER *****/ /***************************/ size = filters; // create the weights nvinfer1::Weights shift{nvinfer1::DataType::kFLOAT, nullptr, size}; nvinfer1::Weights scale{nvinfer1::DataType::kFLOAT, nullptr, size}; nvinfer1::Weights power{nvinfer1::DataType::kFLOAT, nullptr, size}; float* shiftWt = new float[size]; for (int i = 0; i < size; ++i) { shiftWt[i] = bnBiases.at(i) - ((bnRunningMean.at(i) * bnWeights.at(i)) / bnRunningVar.at(i)); } shift.values = shiftWt; float* scaleWt = new float[size]; for (int i = 0; i < size; ++i) { scaleWt[i] = bnWeights.at(i) / bnRunningVar[i]; } scale.values = scaleWt; float* powerWt = new float[size]; for (int i = 0; i < size; ++i) { powerWt[i] = 1.0; } power.values = powerWt; trtWeights.push_back(shift); trtWeights.push_back(scale); trtWeights.push_back(power); // Add the batch norm layers nvinfer1::IScaleLayer* bn = network->addScale( *conv->getOutput(0), nvinfer1::ScaleMode::kCHANNEL, shift, scale, power); assert(bn != nullptr); std::string bnLayerName = "batch_norm_" + std::to_string(layerIdx); bn->setName(bnLayerName.c_str()); /***** ACTIVATION LAYER *****/ /****************************/ nvinfer1::ITensor* bnOutput = bn->getOutput(0); nvinfer1::IActivationLayer* leaky = network->addActivation( *bnOutput, nvinfer1::ActivationType::kLEAKY_RELU); leaky->setAlpha(0.1); assert(leaky != nullptr); std::string leakyLayerName = "leaky_" + std::to_string(layerIdx); leaky->setName(leakyLayerName.c_str()); return leaky; } nvinfer1::ILayer* netAddUpsample(int layerIdx, std::map& block, std::vector& weights, std::vector& trtWeights, int& inputChannels, nvinfer1::ITensor* input, nvinfer1::INetworkDefinition* network) { assert(block.at("type") == "upsample"); nvinfer1::Dims inpDims = input->getDimensions(); assert(inpDims.nbDims == 3); assert(inpDims.d[1] == inpDims.d[2]); int h = inpDims.d[1]; int w = inpDims.d[2]; int stride = std::stoi(block.at("stride")); // add pre multiply matrix as a constant nvinfer1::Dims preDims{3, {1, stride * h, w}, {nvinfer1::DimensionType::kCHANNEL, nvinfer1::DimensionType::kSPATIAL, nvinfer1::DimensionType::kSPATIAL}}; int size = stride * h * w; nvinfer1::Weights preMul{nvinfer1::DataType::kFLOAT, nullptr, size}; float* preWt = new float[size]; /* (2*h * w) [ [1, 0, ..., 0], [1, 0, ..., 0], [0, 1, ..., 0], [0, 1, ..., 0], ..., ..., [0, 0, ..., 1], [0, 0, ..., 1] ] */ for (int i = 0, idx = 0; i < h; ++i) { for (int s = 0; s < stride; ++s) { for (int j = 0; j < w; ++j, ++idx) { preWt[idx] = (i == j) ? 1.0 : 0.0; } } } preMul.values = preWt; trtWeights.push_back(preMul); nvinfer1::IConstantLayer* preM = network->addConstant(preDims, preMul); assert(preM != nullptr); std::string preLayerName = "preMul_" + std::to_string(layerIdx); preM->setName(preLayerName.c_str()); // add post multiply matrix as a constant nvinfer1::Dims postDims{3, {1, h, stride * w}, {nvinfer1::DimensionType::kCHANNEL, nvinfer1::DimensionType::kSPATIAL, nvinfer1::DimensionType::kSPATIAL}}; size = stride * h * w; nvinfer1::Weights postMul{nvinfer1::DataType::kFLOAT, nullptr, size}; float* postWt = new float[size]; /* (h * 2*w) [ [1, 1, 0, 0, ..., 0, 0], [0, 0, 1, 1, ..., 0, 0], ..., ..., [0, 0, 0, 0, ..., 1, 1] ] */ for (int i = 0, idx = 0; i < h; ++i) { for (int j = 0; j < stride * w; ++j, ++idx) { postWt[idx] = (j / stride == i) ? 1.0 : 0.0; } } postMul.values = postWt; trtWeights.push_back(postMul); nvinfer1::IConstantLayer* post_m = network->addConstant(postDims, postMul); assert(post_m != nullptr); std::string postLayerName = "postMul_" + std::to_string(layerIdx); post_m->setName(postLayerName.c_str()); // add matrix multiply layers for upsampling nvinfer1::IMatrixMultiplyLayer* mm1 = network->addMatrixMultiply(*preM->getOutput(0), nvinfer1::MatrixOperation::kNONE, *input, nvinfer1::MatrixOperation::kNONE); assert(mm1 != nullptr); std::string mm1LayerName = "mm1_" + std::to_string(layerIdx); mm1->setName(mm1LayerName.c_str()); nvinfer1::IMatrixMultiplyLayer* mm2 = network->addMatrixMultiply(*mm1->getOutput(0), nvinfer1::MatrixOperation::kNONE, *post_m->getOutput(0), nvinfer1::MatrixOperation::kNONE); assert(mm2 != nullptr); std::string mm2LayerName = "mm2_" + std::to_string(layerIdx); mm2->setName(mm2LayerName.c_str()); return mm2; } void printLayerInfo(std::string layerIndex, std::string layerName, std::string layerInput, std::string layerOutput, std::string weightPtr) { std::cout << std::setw(6) << std::left << layerIndex << std::setw(15) << std::left << layerName; std::cout << std::setw(20) << std::left << layerInput << std::setw(20) << std::left << layerOutput; std::cout << std::setw(6) << std::left << weightPtr << std::endl; } ================================================ FILE: deepstream_plugin_yolov4/trt_utils.h ================================================ /* * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. * * Permission 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: * * The above copyright notice and this permission notice shall be included in * all copies or substantial portions of the Software. * * THE 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. */ #ifndef __TRT_UTILS_H__ #define __TRT_UTILS_H__ #include #include #include #include #include #include #include #include "NvInfer.h" #define UNUSED(expr) (void)(expr) #define DIVUP(n, d) ((n) + (d)-1) / (d) std::string trim(std::string s); float clamp(const float val, const float minVal, const float maxVal); bool fileExists(const std::string fileName, bool verbose = true); std::vector loadWeights(const std::string weightsFilePath, const std::string& networkType); std::string dimsToString(const nvinfer1::Dims d); void displayDimType(const nvinfer1::Dims d); int getNumChannels(nvinfer1::ITensor* t); uint64_t get3DTensorVolume(nvinfer1::Dims inputDims); // Helper functions to create yolo engine nvinfer1::ILayer* netAddMaxpool(int layerIdx, std::map& block, nvinfer1::ITensor* input, nvinfer1::INetworkDefinition* network); nvinfer1::ILayer* netAddConvLinear(int layerIdx, std::map& block, std::vector& weights, std::vector& trtWeights, int& weightPtr, int& inputChannels, nvinfer1::ITensor* input, nvinfer1::INetworkDefinition* network); nvinfer1::ILayer* netAddConvBNLeaky(int layerIdx, std::map& block, std::vector& weights, std::vector& trtWeights, int& weightPtr, int& inputChannels, nvinfer1::ITensor* input, nvinfer1::INetworkDefinition* network); nvinfer1::ILayer* netAddUpsample(int layerIdx, std::map& block, std::vector& weights, std::vector& trtWeights, int& inputChannels, nvinfer1::ITensor* input, nvinfer1::INetworkDefinition* network); void printLayerInfo(std::string layerIndex, std::string layerName, std::string layerInput, std::string layerOutput, std::string weightPtr); #endif ================================================ FILE: deepstream_plugin_yolov4/yolo.cpp ================================================ /* * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. * * Permission 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: * * The above copyright notice and this permission notice shall be included in * all copies or substantial portions of the Software. * * THE 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. */ #include "yolo.h" #include "yoloPlugins.h" #include #include #include Yolo::Yolo(const NetworkInfo& networkInfo) : m_NetworkType(networkInfo.networkType), // yolov3 m_ConfigFilePath(networkInfo.configFilePath), // yolov3.cfg m_WtsFilePath(networkInfo.wtsFilePath), // yolov3.weights m_DeviceType(networkInfo.deviceType), // kDLA, kGPU m_InputBlobName(networkInfo.inputBlobName), // data m_InputH(0), m_InputW(0), m_InputC(0), m_InputSize(0) {} Yolo::~Yolo() { destroyNetworkUtils(); } nvinfer1::ICudaEngine *Yolo::createEngine (nvinfer1::IBuilder* builder) { assert (builder); std::vector weights = loadWeights(m_WtsFilePath, m_NetworkType); std::vector trtWeights; nvinfer1::INetworkDefinition *network = builder->createNetwork(); if (parseModel(*network) != NVDSINFER_SUCCESS) { network->destroy(); return nullptr; } // Build the engine std::cout << "Building the TensorRT Engine..." << std::endl; nvinfer1::ICudaEngine * engine = builder->buildCudaEngine(*network); if (engine) { std::cout << "Building complete!" << std::endl; } else { std::cerr << "Building engine failed!" << std::endl; } // destroy network->destroy(); return engine; } NvDsInferStatus Yolo::parseModel(nvinfer1::INetworkDefinition& network) { destroyNetworkUtils(); m_ConfigBlocks = parseConfigFile(m_ConfigFilePath); parseConfigBlocks(); std::vector weights = loadWeights(m_WtsFilePath, m_NetworkType); // build yolo network std::cout << "Building Yolo network..." << std::endl; NvDsInferStatus status = buildYoloNetwork(weights, network); if (status == NVDSINFER_SUCCESS) { std::cout << "Building yolo network complete!" << std::endl; } else { std::cerr << "Building yolo network failed!" << std::endl; } return status; } NvDsInferStatus Yolo::buildYoloNetwork( std::vector& weights, nvinfer1::INetworkDefinition& network) { int weightPtr = 0; int channels = m_InputC; nvinfer1::ITensor* data = network.addInput(m_InputBlobName.c_str(), nvinfer1::DataType::kFLOAT, nvinfer1::DimsCHW{static_cast(m_InputC), static_cast(m_InputH), static_cast(m_InputW)}); assert(data != nullptr && data->getDimensions().nbDims > 0); nvinfer1::ITensor* previous = data; std::vector tensorOutputs; uint outputTensorCount = 0; // build the network using the network API for (uint i = 0; i < m_ConfigBlocks.size(); ++i) { // check if num. of channels is correct assert(getNumChannels(previous) == channels); std::string layerIndex = "(" + std::to_string(tensorOutputs.size()) + ")"; if (m_ConfigBlocks.at(i).at("type") == "net") { printLayerInfo("", "layer", " inp_size", " out_size", "weightPtr"); } else if (m_ConfigBlocks.at(i).at("type") == "convolutional") { std::string inputVol = dimsToString(previous->getDimensions()); nvinfer1::ILayer* out; std::string layerType; // check if batch_norm enabled if (m_ConfigBlocks.at(i).find("batch_normalize") != m_ConfigBlocks.at(i).end()) { out = netAddConvBNLeaky(i, m_ConfigBlocks.at(i), weights, m_TrtWeights, weightPtr, channels, previous, &network); layerType = "conv-bn-leaky"; } else { out = netAddConvLinear(i, m_ConfigBlocks.at(i), weights, m_TrtWeights, weightPtr, channels, previous, &network); layerType = "conv-linear"; } previous = out->getOutput(0); assert(previous != nullptr); channels = getNumChannels(previous); std::string outputVol = dimsToString(previous->getDimensions()); tensorOutputs.push_back(out->getOutput(0)); printLayerInfo(layerIndex, layerType, inputVol, outputVol, std::to_string(weightPtr)); } else if (m_ConfigBlocks.at(i).at("type") == "shortcut") { assert(m_ConfigBlocks.at(i).at("activation") == "linear"); assert(m_ConfigBlocks.at(i).find("from") != m_ConfigBlocks.at(i).end()); int from = stoi(m_ConfigBlocks.at(i).at("from")); std::string inputVol = dimsToString(previous->getDimensions()); // check if indexes are correct assert((i - 2 >= 0) && (i - 2 < tensorOutputs.size())); assert((i + from - 1 >= 0) && (i + from - 1 < tensorOutputs.size())); assert(i + from - 1 < i - 2); nvinfer1::IElementWiseLayer* ew = network.addElementWise( *tensorOutputs[i - 2], *tensorOutputs[i + from - 1], nvinfer1::ElementWiseOperation::kSUM); assert(ew != nullptr); std::string ewLayerName = "shortcut_" + std::to_string(i); ew->setName(ewLayerName.c_str()); previous = ew->getOutput(0); assert(previous != nullptr); std::string outputVol = dimsToString(previous->getDimensions()); tensorOutputs.push_back(ew->getOutput(0)); printLayerInfo(layerIndex, "skip", inputVol, outputVol, " -"); } else if (m_ConfigBlocks.at(i).at("type") == "yolo") { nvinfer1::Dims prevTensorDims = previous->getDimensions(); assert(prevTensorDims.d[1] == prevTensorDims.d[2]); TensorInfo& curYoloTensor = m_OutputTensors.at(outputTensorCount); curYoloTensor.gridSize = prevTensorDims.d[1]; curYoloTensor.stride = m_InputW / curYoloTensor.gridSize; m_OutputTensors.at(outputTensorCount).volume = curYoloTensor.gridSize * curYoloTensor.gridSize * (curYoloTensor.numBBoxes * (5 + curYoloTensor.numClasses)); std::string layerName = "yolo_" + std::to_string(i); curYoloTensor.blobName = layerName; nvinfer1::IPluginV2* yoloPlugin = new YoloLayerV3(m_OutputTensors.at(outputTensorCount).numBBoxes, m_OutputTensors.at(outputTensorCount).numClasses, m_OutputTensors.at(outputTensorCount).gridSize); assert(yoloPlugin != nullptr); nvinfer1::IPluginV2Layer* yolo = network.addPluginV2(&previous, 1, *yoloPlugin); assert(yolo != nullptr); yolo->setName(layerName.c_str()); std::string inputVol = dimsToString(previous->getDimensions()); previous = yolo->getOutput(0); assert(previous != nullptr); previous->setName(layerName.c_str()); std::string outputVol = dimsToString(previous->getDimensions()); network.markOutput(*previous); channels = getNumChannels(previous); tensorOutputs.push_back(yolo->getOutput(0)); printLayerInfo(layerIndex, "yolo", inputVol, outputVol, std::to_string(weightPtr)); ++outputTensorCount; } else if (m_ConfigBlocks.at(i).at("type") == "region") { nvinfer1::Dims prevTensorDims = previous->getDimensions(); assert(prevTensorDims.d[1] == prevTensorDims.d[2]); TensorInfo& curRegionTensor = m_OutputTensors.at(outputTensorCount); curRegionTensor.gridSize = prevTensorDims.d[1]; curRegionTensor.stride = m_InputW / curRegionTensor.gridSize; m_OutputTensors.at(outputTensorCount).volume = curRegionTensor.gridSize * curRegionTensor.gridSize * (curRegionTensor.numBBoxes * (5 + curRegionTensor.numClasses)); std::string layerName = "region_" + std::to_string(i); curRegionTensor.blobName = layerName; nvinfer1::plugin::RegionParameters RegionParameters{ static_cast(curRegionTensor.numBBoxes), 4, static_cast(curRegionTensor.numClasses), nullptr}; std::string inputVol = dimsToString(previous->getDimensions()); nvinfer1::IPluginV2* regionPlugin = createRegionPlugin(RegionParameters); assert(regionPlugin != nullptr); nvinfer1::IPluginV2Layer* region = network.addPluginV2(&previous, 1, *regionPlugin); assert(region != nullptr); region->setName(layerName.c_str()); previous = region->getOutput(0); assert(previous != nullptr); previous->setName(layerName.c_str()); std::string outputVol = dimsToString(previous->getDimensions()); network.markOutput(*previous); channels = getNumChannels(previous); tensorOutputs.push_back(region->getOutput(0)); printLayerInfo(layerIndex, "region", inputVol, outputVol, std::to_string(weightPtr)); std::cout << "Anchors are being converted to network input resolution i.e. Anchors x " << curRegionTensor.stride << " (stride)" << std::endl; for (auto& anchor : curRegionTensor.anchors) anchor *= curRegionTensor.stride; ++outputTensorCount; } else if (m_ConfigBlocks.at(i).at("type") == "reorg") { std::string inputVol = dimsToString(previous->getDimensions()); nvinfer1::IPluginV2* reorgPlugin = createReorgPlugin(2); assert(reorgPlugin != nullptr); nvinfer1::IPluginV2Layer* reorg = network.addPluginV2(&previous, 1, *reorgPlugin); assert(reorg != nullptr); std::string layerName = "reorg_" + std::to_string(i); reorg->setName(layerName.c_str()); previous = reorg->getOutput(0); assert(previous != nullptr); std::string outputVol = dimsToString(previous->getDimensions()); channels = getNumChannels(previous); tensorOutputs.push_back(reorg->getOutput(0)); printLayerInfo(layerIndex, "reorg", inputVol, outputVol, std::to_string(weightPtr)); } // route layers (single or concat) else if (m_ConfigBlocks.at(i).at("type") == "route") { std::string strLayers = m_ConfigBlocks.at(i).at("layers"); std::vector idxLayers; size_t lastPos = 0, pos = 0; while ((pos = strLayers.find(',', lastPos)) != std::string::npos) { int vL = std::stoi(trim(strLayers.substr(lastPos, pos - lastPos))); idxLayers.push_back (vL); lastPos = pos + 1; } if (lastPos < strLayers.length()) { std::string lastV = trim(strLayers.substr(lastPos)); if (!lastV.empty()) { idxLayers.push_back (std::stoi(lastV)); } } assert (!idxLayers.empty()); std::vector concatInputs; for (int idxLayer : idxLayers) { if (idxLayer < 0) { idxLayer = tensorOutputs.size() + idxLayer; } assert (idxLayer >= 0 && idxLayer < (int)tensorOutputs.size()); concatInputs.push_back (tensorOutputs[idxLayer]); } nvinfer1::IConcatenationLayer* concat = network.addConcatenation(concatInputs.data(), concatInputs.size()); assert(concat != nullptr); std::string concatLayerName = "route_" + std::to_string(i - 1); concat->setName(concatLayerName.c_str()); // concatenate along the channel dimension concat->setAxis(0); previous = concat->getOutput(0); assert(previous != nullptr); std::string outputVol = dimsToString(previous->getDimensions()); // set the output volume depth channels = getNumChannels(previous); tensorOutputs.push_back(concat->getOutput(0)); printLayerInfo(layerIndex, "route", " -", outputVol, std::to_string(weightPtr)); } else if (m_ConfigBlocks.at(i).at("type") == "upsample") { std::string inputVol = dimsToString(previous->getDimensions()); nvinfer1::ILayer* out = netAddUpsample(i - 1, m_ConfigBlocks[i], weights, m_TrtWeights, channels, previous, &network); previous = out->getOutput(0); std::string outputVol = dimsToString(previous->getDimensions()); tensorOutputs.push_back(out->getOutput(0)); printLayerInfo(layerIndex, "upsample", inputVol, outputVol, " -"); } else if (m_ConfigBlocks.at(i).at("type") == "maxpool") { std::string inputVol = dimsToString(previous->getDimensions()); nvinfer1::ILayer* out = netAddMaxpool(i, m_ConfigBlocks.at(i), previous, &network); previous = out->getOutput(0); assert(previous != nullptr); std::string outputVol = dimsToString(previous->getDimensions()); tensorOutputs.push_back(out->getOutput(0)); printLayerInfo(layerIndex, "maxpool", inputVol, outputVol, std::to_string(weightPtr)); } else { std::cout << "Unsupported layer type --> \"" << m_ConfigBlocks.at(i).at("type") << "\"" << std::endl; assert(0); } } if ((int)weights.size() != weightPtr) { std::cout << "Number of unused weights left : " << weights.size() - weightPtr << std::endl; assert(0); } std::cout << "Output yolo blob names :" << std::endl; for (auto& tensor : m_OutputTensors) { std::cout << tensor.blobName << std::endl; } int nbLayers = network.getNbLayers(); std::cout << "Total number of yolo layers: " << nbLayers << std::endl; return NVDSINFER_SUCCESS; } std::vector> Yolo::parseConfigFile (const std::string cfgFilePath) { assert(fileExists(cfgFilePath)); std::ifstream file(cfgFilePath); assert(file.good()); std::string line; std::vector> blocks; std::map block; while (getline(file, line)) { if (line.size() == 0) continue; if (line.front() == '#') continue; line = trim(line); if (line.front() == '[') { if (block.size() > 0) { blocks.push_back(block); block.clear(); } std::string key = "type"; std::string value = trim(line.substr(1, line.size() - 2)); block.insert(std::pair(key, value)); } else { int cpos = line.find('='); std::string key = trim(line.substr(0, cpos)); std::string value = trim(line.substr(cpos + 1)); block.insert(std::pair(key, value)); } } blocks.push_back(block); return blocks; } void Yolo::parseConfigBlocks() { for (auto block : m_ConfigBlocks) { if (block.at("type") == "net") { assert((block.find("height") != block.end()) && "Missing 'height' param in network cfg"); assert((block.find("width") != block.end()) && "Missing 'width' param in network cfg"); assert((block.find("channels") != block.end()) && "Missing 'channels' param in network cfg"); m_InputH = std::stoul(block.at("height")); m_InputW = std::stoul(block.at("width")); m_InputC = std::stoul(block.at("channels")); assert(m_InputW == m_InputH); m_InputSize = m_InputC * m_InputH * m_InputW; } else if ((block.at("type") == "region") || (block.at("type") == "yolo")) { assert((block.find("num") != block.end()) && std::string("Missing 'num' param in " + block.at("type") + " layer").c_str()); assert((block.find("classes") != block.end()) && std::string("Missing 'classes' param in " + block.at("type") + " layer") .c_str()); assert((block.find("anchors") != block.end()) && std::string("Missing 'anchors' param in " + block.at("type") + " layer") .c_str()); TensorInfo outputTensor; std::string anchorString = block.at("anchors"); while (!anchorString.empty()) { int npos = anchorString.find_first_of(','); if (npos != -1) { float anchor = std::stof(trim(anchorString.substr(0, npos))); outputTensor.anchors.push_back(anchor); anchorString.erase(0, npos + 1); } else { float anchor = std::stof(trim(anchorString)); outputTensor.anchors.push_back(anchor); break; } } if ((m_NetworkType == "yolov3") || (m_NetworkType == "yolov3-tiny")) { assert((block.find("mask") != block.end()) && std::string("Missing 'mask' param in " + block.at("type") + " layer") .c_str()); std::string maskString = block.at("mask"); while (!maskString.empty()) { int npos = maskString.find_first_of(','); if (npos != -1) { uint mask = std::stoul(trim(maskString.substr(0, npos))); outputTensor.masks.push_back(mask); maskString.erase(0, npos + 1); } else { uint mask = std::stoul(trim(maskString)); outputTensor.masks.push_back(mask); break; } } } outputTensor.numBBoxes = outputTensor.masks.size() > 0 ? outputTensor.masks.size() : std::stoul(trim(block.at("num"))); outputTensor.numClasses = std::stoul(block.at("classes")); m_OutputTensors.push_back(outputTensor); } } } void Yolo::destroyNetworkUtils() { // deallocate the weights for (uint i = 0; i < m_TrtWeights.size(); ++i) { if (m_TrtWeights[i].count > 0) free(const_cast(m_TrtWeights[i].values)); } m_TrtWeights.clear(); } ================================================ FILE: deepstream_plugin_yolov4/yolo.h ================================================ /* * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. * * Permission 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: * * The above copyright notice and this permission notice shall be included in * all copies or substantial portions of the Software. * * THE 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. */ #ifndef _YOLO_H_ #define _YOLO_H_ #include #include #include #include #include "NvInfer.h" #include "trt_utils.h" #include "nvdsinfer_custom_impl.h" /** * Holds all the file paths required to build a network. */ struct NetworkInfo { std::string networkType; std::string configFilePath; std::string wtsFilePath; std::string deviceType; std::string inputBlobName; }; /** * Holds information about an output tensor of the yolo network. */ struct TensorInfo { std::string blobName; uint stride{0}; uint gridSize{0}; uint numClasses{0}; uint numBBoxes{0}; uint64_t volume{0}; std::vector masks; std::vector anchors; int bindingIndex{-1}; float* hostBuffer{nullptr}; }; class Yolo : public IModelParser { public: Yolo(const NetworkInfo& networkInfo); ~Yolo() override; bool hasFullDimsSupported() const override { return false; } const char* getModelName() const override { return m_ConfigFilePath.empty() ? m_NetworkType.c_str() : m_ConfigFilePath.c_str(); } NvDsInferStatus parseModel(nvinfer1::INetworkDefinition& network) override; nvinfer1::ICudaEngine *createEngine (nvinfer1::IBuilder* builder); protected: const std::string m_NetworkType; const std::string m_ConfigFilePath; const std::string m_WtsFilePath; const std::string m_DeviceType; const std::string m_InputBlobName; std::vector m_OutputTensors; std::vector> m_ConfigBlocks; uint m_InputH; uint m_InputW; uint m_InputC; uint64_t m_InputSize; // TRT specific members std::vector m_TrtWeights; private: NvDsInferStatus buildYoloNetwork( std::vector& weights, nvinfer1::INetworkDefinition& network); std::vector> parseConfigFile( const std::string cfgFilePath); void parseConfigBlocks(); void destroyNetworkUtils(); }; #endif // _YOLO_H_ ================================================ FILE: deepstream_plugin_yolov4/yoloPlugins.cpp ================================================ /* * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. * * Permission 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: * * The above copyright notice and this permission notice shall be included in * all copies or substantial portions of the Software. * * THE 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. */ #include "yoloPlugins.h" #include "NvInferPlugin.h" #include #include #include namespace { template void write(char*& buffer, const T& val) { *reinterpret_cast(buffer) = val; buffer += sizeof(T); } template void read(const char*& buffer, T& val) { val = *reinterpret_cast(buffer); buffer += sizeof(T); } } //namespace // Forward declaration of cuda kernels cudaError_t cudaYoloLayerV3 ( const void* input, void* output, const uint& batchSize, const uint& gridSize, const uint& numOutputClasses, const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream); YoloLayerV3::YoloLayerV3 (const void* data, size_t length) { const char *d = static_cast(data); read(d, m_NumBoxes); read(d, m_NumClasses); read(d, m_GridSize); read(d, m_OutputSize); }; YoloLayerV3::YoloLayerV3 ( const uint& numBoxes, const uint& numClasses, const uint& gridSize) : m_NumBoxes(numBoxes), m_NumClasses(numClasses), m_GridSize(gridSize) { assert(m_NumBoxes > 0); assert(m_NumClasses > 0); assert(m_GridSize > 0); m_OutputSize = m_GridSize * m_GridSize * (m_NumBoxes * (4 + 1 + m_NumClasses)); }; nvinfer1::Dims YoloLayerV3::getOutputDimensions( int index, const nvinfer1::Dims* inputs, int nbInputDims) { assert(index == 0); assert(nbInputDims == 1); return inputs[0]; } bool YoloLayerV3::supportsFormat ( nvinfer1::DataType type, nvinfer1::PluginFormat format) const { return (type == nvinfer1::DataType::kFLOAT && format == nvinfer1::PluginFormat::kNCHW); } void YoloLayerV3::configureWithFormat ( const nvinfer1::Dims* inputDims, int nbInputs, const nvinfer1::Dims* outputDims, int nbOutputs, nvinfer1::DataType type, nvinfer1::PluginFormat format, int maxBatchSize) { assert(nbInputs == 1); assert (format == nvinfer1::PluginFormat::kNCHW); assert(inputDims != nullptr); } int YoloLayerV3::enqueue( int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) { CHECK(cudaYoloLayerV3( inputs[0], outputs[0], batchSize, m_GridSize, m_NumClasses, m_NumBoxes, m_OutputSize, stream)); return 0; } size_t YoloLayerV3::getSerializationSize() const { return sizeof(m_NumBoxes) + sizeof(m_NumClasses) + sizeof(m_GridSize) + sizeof(m_OutputSize); } void YoloLayerV3::serialize(void* buffer) const { char *d = static_cast(buffer); write(d, m_NumBoxes); write(d, m_NumClasses); write(d, m_GridSize); write(d, m_OutputSize); } nvinfer1::IPluginV2* YoloLayerV3::clone() const { return new YoloLayerV3 (m_NumBoxes, m_NumClasses, m_GridSize); } REGISTER_TENSORRT_PLUGIN(YoloLayerV3PluginCreator); ================================================ FILE: deepstream_plugin_yolov4/yoloPlugins.h ================================================ /* * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. * * Permission 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: * * The above copyright notice and this permission notice shall be included in * all copies or substantial portions of the Software. * * THE 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. */ #ifndef __YOLO_PLUGINS__ #define __YOLO_PLUGINS__ #include #include #include #include #include #include "NvInferPlugin.h" #define CHECK(status) \ { \ if (status != 0) \ { \ std::cout << "Cuda failure: " << cudaGetErrorString(status) << " in file " << __FILE__ \ << " at line " << __LINE__ << std::endl; \ abort(); \ } \ } namespace { const char* YOLOV3LAYER_PLUGIN_VERSION {"1"}; const char* YOLOV3LAYER_PLUGIN_NAME {"YoloLayerV3_TRT"}; } // namespace class YoloLayerV3 : public nvinfer1::IPluginV2 { public: YoloLayerV3 (const void* data, size_t length); YoloLayerV3 (const uint& numBoxes, const uint& numClasses, const uint& gridSize); const char* getPluginType () const override { return YOLOV3LAYER_PLUGIN_NAME; } const char* getPluginVersion () const override { return YOLOV3LAYER_PLUGIN_VERSION; } int getNbOutputs () const override { return 1; } nvinfer1::Dims getOutputDimensions ( int index, const nvinfer1::Dims* inputs, int nbInputDims) override; bool supportsFormat ( nvinfer1::DataType type, nvinfer1::PluginFormat format) const override; void configureWithFormat ( const nvinfer1::Dims* inputDims, int nbInputs, const nvinfer1::Dims* outputDims, int nbOutputs, nvinfer1::DataType type, nvinfer1::PluginFormat format, int maxBatchSize) override; int initialize () override { return 0; } void terminate () override {} size_t getWorkspaceSize (int maxBatchSize) const override { return 0; } int enqueue ( int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) override; size_t getSerializationSize() const override; void serialize (void* buffer) const override; void destroy () override { delete this; } nvinfer1::IPluginV2* clone() const override; void setPluginNamespace (const char* pluginNamespace)override { m_Namespace = pluginNamespace; } virtual const char* getPluginNamespace () const override { return m_Namespace.c_str(); } private: uint m_NumBoxes {0}; uint m_NumClasses {0}; uint m_GridSize {0}; uint64_t m_OutputSize {0}; std::string m_Namespace {""}; }; class YoloLayerV3PluginCreator : public nvinfer1::IPluginCreator { public: YoloLayerV3PluginCreator () {} ~YoloLayerV3PluginCreator () {} const char* getPluginName () const override { return YOLOV3LAYER_PLUGIN_NAME; } const char* getPluginVersion () const override { return YOLOV3LAYER_PLUGIN_VERSION; } const nvinfer1::PluginFieldCollection* getFieldNames() override { std::cerr<< "YoloLayerV3PluginCreator::getFieldNames is not implemented" << std::endl; return nullptr; } nvinfer1::IPluginV2* createPlugin ( const char* name, const nvinfer1::PluginFieldCollection* fc) override { std::cerr<< "YoloLayerV3PluginCreator::getFieldNames is not implemented.\n"; return nullptr; } nvinfer1::IPluginV2* deserializePlugin ( const char* name, const void* serialData, size_t serialLength) override { std::cout << "Deserialize yoloLayerV3 plugin: " << name << std::endl; return new YoloLayerV3(serialData, serialLength); } void setPluginNamespace(const char* libNamespace) override { m_Namespace = libNamespace; } const char* getPluginNamespace() const override { return m_Namespace.c_str(); } private: std::string m_Namespace {""}; }; #endif // __YOLO_PLUGINS__ ================================================ FILE: docker/constraints.docker ================================================ scikit-learn==0.19.1 opencv-python==3.2.0 ================================================ FILE: docker/opencv_python-3.2.0.egg-info ================================================ Metadata-Version: 1.1 Name: opencv-python Version: 3.2.0 Summary: stub file to make pip aware of apt installed opencv Home-page: UNKNOWN Author: UNKNOWN Author-email: UNKNOWN License: UNKNOWN Description: UNKNOWN Platform: UNKNOWN Classifier: Programming Language :: Python :: 3 ================================================ FILE: docker/scikit-learn-0.19.1.egg-info ================================================ Metadata-Version: 1.1 Name: scikit-learn Version: 0.19.1 Summary: stub file to make pip aware of apt installed package Home-page: UNKNOWN Author: UNKNOWN Author-email: UNKNOWN License: UNKNOWN Description: UNKNOWN Platform: UNKNOWN Classifier: Programming Language :: Python :: 3 ================================================ FILE: docker/start.sh ================================================ #!/usr/bin/env bash nvargus-daemon & echo echo "Provided input:" echo " - MASKCAM_INPUT = $MASKCAM_INPUT" echo "Device Address:" echo " - MASKCAM_DEVICE_ADDRESS = $MASKCAM_DEVICE_ADDRESS" echo "Development mode:" echo " - DEV_MODE = $DEV_MODE" echo echo "MQTT configuration:" echo " - MQTT_BROKER_IP = $MQTT_BROKER_IP" echo " - MQTT_DEVICE_NAME = $MQTT_DEVICE_NAME" echo if [[ $DEV_MODE -eq 1 ]]; then echo "Development mode enabled, exec maskcam_run.py manually" /bin/bash else ./maskcam_run.py $MASKCAM_INPUT fi ================================================ FILE: docker-compose.yml ================================================ version: '2' # this file only exists for local development # pushes to override the restart function services: mc_maskcam: restart: no build: . privileged: true ports: - "1883:1883" - "8080:8080" - "8554:8554" ================================================ FILE: docs/BalenaOS-DevKit-Nano-Setup.md ================================================ # BalenaOS Developer Kit Nano Setup Instructions for MaskCam BalenaOS is a very light weight distribution designed for running containers on edge devices. It has a number of advantages for fleet deployment and management, especially when combined with balena's balenaCloud mangament system. Explaining the details of how to set up balenaCloud applications is beyond the scope of this document, but you can test MaskCam on balenaOS using a local development environment setup. Except for installing balenaOS and using a slightly modified launch command, this process is essentially the same as the Jetson Nano Development kit instructions from [the main README](https://github.com/bdtinc/maskcam#running-maskcam-from-a-container-on-a-jetson-nano-developer-kit). If you want to use balenaCloud instead (i.e: see your device in the web dashboard), and you're willing to take some time to push the container to your own account, check [Using balenaCloud](#using-balenacloud) at the end of this section. In any case, this will require a Jetson Nano Development Kit, a 32 GB or higher Micro-SD card, and another computer (referred to here as main system) on the same network. ### Installing balenaOS As mentioned, this procedure will not link your device with a balenaCloud account, but instead it will enable local development. First, go to https://www.balena.io/os/, scroll to the Download section, and download the development version for Nvidia Jetson Nano SD-CARD. Next, go to https://www.balena.io/etcher/ and install balenaEtcher. In balenaEtcher, simply select the zip file you downloaded, and after inserting the sd card into your main system select it, then press the 'Flash!' icon. After the flashing process is completed, place the sd card into your Jetson Nano Development Kit, ensure the network cable is plugged into the device and power up the Jetson. ### Installing balena CLI Use [these instructions](https://github.com/balena-io/balena-cli/blob/master/INSTALL.md) to install the balena CLI tool. ### Connecting to your Jetson First, in a terminal on your main system run the command: ``` sudo balena scan ``` Note the ip address in the result. Next connect to your Jetson: ``` balena ssh ``` At this point you are in a console as root user on your Jetson running balenaOS. The commands from this point on are exactly the same as the instructions for running using JetPack on the Nano Developer Kit with the following differences. 1. The `docker` command is replaced by `balena` 2. Do not use the `--runtime nvidia` switch. It is automatic on balenaOS for Jetson and you will get errors if you include it. So issuing the following commands will run MaskCam: ``` $ balena pull maskcam/maskcam-beta $ balena run --privileged --rm -it --env MASKCAM_DEVICE_ADDRESS= -p 1883:1883 -p 8080:8080 -p 8554:8554 maskcam/maskcam-beta ``` Note that setting `MASKCAM_DEVICE_ADDRESS` is optional, and you can also set other configuration parameters exactly as indicated in the [device configuration](https://github.com/bdtinc/maskcam#setting-device-configuration-parameters) section of the main docs. ### Using balenaCloud You can create a free balenaCloud account that will allow you to link up to 10 devices, in order to test some of the most useful features that this platform provides. You'll need to create an App, install the balena CLI and then follow these instructions in order to deploy the maskcam container to your app: https://www.balena.io/docs/learn/deploy/deployment/ For a simple use case, you can just use the `balena push myApp` command from the root directory of this project (it will take a long time while it builds and pushes the whole image), but you should familiarize yourself with the platform and use the deployment method that better fits your needs. ================================================ FILE: docs/BalenaOS-Photon-Nano-Setup.md ================================================ # BalenaOS Setup Instructions for MaskCam on Jetson Nano with Photon Carrier Board BalenaOS is a very light weight distribution designed for running containers on edge devices. It has a number of advantages for fleet deployment and management, especially when combined with balena's balenaCloud mangament system. Explaining the details of how to set up balenaCloud applications is beyond the scope of this document, but you can test MaskCam on balenaOS using a local development environment setup. When using a Jetson Nano with a Photon carrier board (i.e. a "Photon Nano"), the process for installing balenaOS is different than with a Developer Kit. The production Jetson Nano module does not have an SD card slot, so balenaOS has to be directly flashed onto the device over USB, rather than using balenaEtcher. Fortunately, balena has created a flashing tool called [jetson-flash](https://github.com/balena-os/jetson-flash) that allows you to do this. ### Setting up jetson-flash To flash the balenaOS image onto the Photon Nano, we need to use Balena's jetson-flash tool. The instructions here show how to install and use jetson-flash on an Ubuntu v18.04 PC. The tool also requires NodeJS >= v10, which can be installed on Ubuntu using [these installation instructions](https://github.com/nodesource/distributions/blob/master/README.md#installation-instructions). First, clone the jetson-flash repository using: ``` git clone https://github.com/balena-os/jetson-flash.git ``` Next, go to the [balenaOS download page](https://www.balena.io/os/#download) and download the CTI Photon Nano Development image. Unzip the image, and move it to the jetson-flash directory. Then, from inside the `jetson-flash` directory, issue the following command to download the NodeJS package dependencies. ``` npm install ``` Now jetson-flash is ready to be used to flash the OS image onto the Photon Nano. ### Flashing balenaOS onto the Jetson Nano over USB Before flashing the OS, the Photon Nano has to be powered on and put into Force Recovery as shown in the [Photon manual](https://connecttech.com/ftp/pdf/CTIM_NGX002_Manual.pdf). Starting with a Photon carrier board that has a Jetson Nano module installed, plug in 12V power to the barrel jack on the carrier board. Then, press and hold SW2 for at least 10 seconds. Plug a micro-USB cord from the Ubuntu PC to P13 on the bottom side of the Photon carrier board. Verify the board is in Force Recovery mode by issuing `lsusb` and checking that an Nvidia device is listed. If it isn't, try re-connecting the USB cable and repeating the process to put the board in Force Recovery mode. Begin flashing by issuing the following command, where `balena.img` is replaced with the filename for the image that was downloaded and extracted previously. ``` sudo ./bin/cmd.js -f balena.img -m jetson-nano-emmc ``` This will initiate the flashing process, which takes about 10 minutes. Once it's complete, unplug the micro-USB cable and power cycle the Photon carrier board. Plug an Ethernet cable into the Photon to connect it to a local network. ### Installing balena CLI Use [these instructions](https://github.com/balena-io/balena-cli/blob/master/INSTALL.md) to install the balena CLI tool on your Ubuntu PC. ### Connecting to your Jetson On the Ubuntu PC, open a terminal and run: ``` sudo balena scan ``` This will report the IP address of your Photon Nano. Use this IP address with the following command to connect to your Jetson: ``` balena ssh ``` At this point you are in a console as root user on your Jetson running balenaOS. Now, we just need to download the docker container and run it! On balenaOS, "docker" is replaced by "balena", as shown in the following command. Issue the two commands below to download and run MaskCam: ``` $ balena pull maskcam/maskcam-beta $ balena run --privileged --rm -it --env MASKCAM_DEVICE_ADDRESS= -p 1883:1883 -p 8080:8080 -p 8554:8554 maskcam/maskcam-beta ``` Note that setting `MASKCAM_DEVICE_ADDRESS` is optional, and you can also set other configuration parameters exactly as indicated in the [device configuration](https://github.com/bdtinc/maskcam#setting-device-configuration-parameters) section of the main docs. ### Using balenaCloud You can create a free balenaCloud account that will allow you to link up to 10 devices, in order to test some of the most useful features that this platform provides. You'll need to create an App, install the balena CLI and then follow these instructions in order to deploy the maskcam container to your app: https://www.balena.io/docs/learn/deploy/deployment/ For a simple use case, you can just use the `balena push myApp` command from the root directory of this project (it will take a long time while it builds and pushes the whole image), but you should familiarize yourself with the platform and use the deployment method that better fits your needs. ================================================ FILE: docs/Custom-Container-Development.md ================================================ # Custom Container Development The MaskCam code in this repository can be used as a starting point for developing your own smart camera application. If you'd like to develop a custom application (for example, a dog detector that counts how many dogs walk past your house and reports the count to a server), you can build your own container that has the custom code, files, and packages used for your unique application. This page gives instructions on how to build a custom container, rather than downloading our pre-built container from Docker. This page is split in to two sections: - [How to Build Your Own Container from Source on the Jetson Nano](#how-to-build-your-own-container-from-source-on-the-jetson-nano) - [How to Use Your Own Detection Model](#how-to-use-your-own-detection-model) ## How to Build Your Own Container from Source on the Jetson Nano The easiest way to get Maskcam running or set up for development purposes, is by using a container like the one provided in the main [Dockerfile](Dockerfile), which provides the right versions of the OS (Ubuntu 18.04 / Bionic Beaver) and all the system level packages required (mainly NVIDIA L4T packages, GStreamer and DeepStream among others). For development, you could make modifications to the code or the container definition, and then rebuild locally using: ``` docker build . -t maskcam_custom ``` The above building step could be executed in the target Jetson Nano device (easier), or in another development environment (i.e: pushing the result to [Docker Hub](https://hub.docker.com/) and then pulling from device). Either way, once the image is ready on the device, remember to run the container using the `--runtime nvidia` and `--privileged` flags (to access the camera device), and mapping the used ports (MQTT -1883-, static file serving -8080- and streaming -8554-, as defined in [maskcam_config.txt](maskcam_config.txt)): ``` docker run --runtime nvidia --privileged --rm -it -p 1883:1883 -p 8080:8080 -p 8554:8554 maskcam_custom ``` If you still want to better understand some of the [Dockerfile](Dockerfile) steps, or you need to run without a container and are willing to deal with version conflicts, please see the dependencies manual installation and building instructions at [docs/Manual-Dependency-Installation.md](docs/Manual-Dependencies-Installation.md) ## How to Use Your Own Detection Model As mentioned above, MaskCam is a reference design for smart camera applications that need to perform computer vision tasks on the edge. Specifically, those involving **Object Detection** (for which you'll need a TensorRT engine) and **Tracking** (for which we use [Norfair](https://github.com/tryolabs/norfair)). Depending on the degree of similarity with this particular use case, you might need to just change the configuration file or some parts of the source code. ### Changing the DeepStream model If you train a new model that is compatible with DeepStream, and has exactly the same (or a subset of the) object classes that are used in this project (`mask`, `no_mask`, `not_visible`, `misplaced`), then you only need to edit the configuration file. In particular, you should change only the corresponding parts of the [maskcam_config.txt](maskcam_config.txt) file, which are under the `[property]` section, and make them match your app's configuration parameters (usually under a file `config_infer_primary.txt` in NVIDIA sample apps). You should not need to change any of the `[face-processor]`, `[mqtt]` or `[maskcam]` sections of the config file, in order to use a new compatible model. Also, note that the `interval` parameter of that section will be ignored when `inference-interval-auto` is enabled. As an example, you'll find there's commented code showing how to use a `Detectnet_v2` model like the one trained using the [NVIDIA facemask app](https://github.com/NVIDIA-AI-IOT/face-mask-detection), but after converting the label names as mentioned above. Check the [DeepStream docs](https://docs.nvidia.com/metropolis/deepstream/dev-guide/text/DS_using_custom_model.html) for more information about how to convert a model in order to use it with DeepStream (in particular, the `nvinfer` GStreamer plugin). Remember to include your new model engine file in the [Dockerfile](Dockerfile) before building the container! ### Changing the object labels If your custom model does not have exactly the same label names, you should edit the [maskcam_inference.py](maskcam/maskcam_inference.py) file, and change the constants `LABEL_MASK`, `LABEL_NO_MASK`, `LABEL_MISPLACED` and `LABEL_NOT_VISIBLE`, to match your needs. If your application has nothing to do with detecting face masks, then you'll probably need to change many other parts of the source code for this application, but a good place to start is the `FaceMaskProcessor` class definition, used in the same inference file, which contains all the code related to the DeepStream pipeline. ================================================ FILE: docs/Manual-Dependencies-Installation.md ================================================ ## Manual installation and building of dependencies These instructions are aimed to manually recreate a native environment similar to the one produced by the [Dockerfile](Dockerfile). They are tested on **Ubuntu 18.04 (Bionic Beaver)** with **Jetpack 4.4.1**. 1. Make sure these packages are installed at system level (other required packages are not listed here since they're included with Jetpack, check the [Dockerfile](Dockerfile) for a complete list): ``` sudo apt install git, python3-pip, python3-opencv python3-libnvinfer python-gi-dev cuda-toolkit-10-2 ``` 2. Clone this repo: ``` git clone .git ``` 3. Copy any `.egg-info` file under `docker/` to the python's `dist-packages` dir, so that system-level installed packages are visible by Pypi: ``` sudo cp docker/*.egg-info /usr/lib/python3/dist-packages/ ``` 4. Install the requirements listed on `requirements.txt`: ``` pip3 install -r requirements.txt ``` If any version above fails or you want to ignore the pinned versions for some reason, try: ``` # Only run this if you don't want to use the pinned versions pip3 install -r requirements.in -c docker/constraints.docker ``` 5. Install Nvidia DeepStream: Aside from the system requirements of th previous step, you also need to install [DeepStream 5.0](https://docs.nvidia.com/metropolis/deepstream/dev-guide/text/DS_Quickstart.html#jetson-setup) (no need to install Kafka protocol adaptor) and also make sure to install the corresponding **python bindings** for GStreamer [gst-python](https://docs.nvidia.com/metropolis/deepstream/dev-guide/text/DS_Python_Sample_Apps.html#python-bindings), and for DeepStream [pyds](https://docs.nvidia.com/metropolis/deepstream/dev-guide/text/DS_Python_Sample_Apps.html#metadata-access). 6. Compile YOLOv4 plugin for DeepStream: After installing DeepStream, compile the YOLOv4 plugin for DeepStream: ``` cd /deepstream_plugin_yolov4 export CUDA_VER=10.2 make ``` If all went well, you should see a library `libnvdsinfer_custom_impl_Yolo.so` in that directory. 7. Download TensorRT engine file from [here](https://maskcam.s3.us-east-2.amazonaws.com/facemask_y4tiny_1024_608_fp16.trt) and save it as `yolo/facemask_y4tiny_1024_608_fp16.trt`. 8. Now you should be ready to run. By default, the device `/dev/video0` will be used, but other devices can be set as first argument: ```bash # Use default input camera /dev/video0 python3 maskcam_run.py # Equivalent as above: python3 maskcam_run.py v4l2:///dev/video0 # Process an mp4 file instead (no network functions, MQTT and static file server disabled) python3 maskcam_run.py file:///path/to/video.mp4 # Read from Raspi2 camera using device-id python3 maskcam_run.py argus:///0 ``` Check the main [README.md](README) for more parameters that can be configured before running, using environment variables. ================================================ FILE: docs/Useful-Development-Scripts.md ================================================ # Useful development scripts These scripts are intended to be used by developers. They require some knowledge on the subject they're used for. ## Running TensorRT engine on images This script will run the engine on a folder of images, and generate another folder for the images with the bounding boxes drawn, and the detection score. To run this script, you need basically the same general instructions that the regular installation, except that you don't need DeepStream and you do need OpenCV instead. Usage: ``` cd yolo/ python3 run_yolo_images.py path/to/input/folder path/to/output/folder ``` ## Debugging MQTT communication If you want to see the raw messages that the MQTT broker receives, and be able to send custom messages to the device (at your own risk), there's a script `maskcam/mqtt_commander.py`, which may be useful for debugging on your local computer or from the Jetson device itself. The script connects to the MQTT broker and sniffs all the communication to/from any device to the broker. ``` export MQTT_BROKER_IP= export MQTT_DEVICE_NAME= python3 -m maskcam.mqtt_commander ``` ## Convert weights generated using the original darknet implementation to TRT 1. Clone the pytorch implementation of YOLOv4: ``` git clone git@github.com:Tianxiaomo/pytorch-YOLOv4.git ``` 2. Convert the Darknet model to ONNX using the script in `tool/darknet2onnx.py`, e.g: ``` PYTHONPATH='pytorch-YOLOv4:$PYTHONPATH' python3 pytorch-YOLOv4/tool/darknet2onnx.py yolo/facemask-yolov4-tiny.cfg yolo/facemask-yolov4-tiny_best.weights ``` 3. Convert the ONNX model to TRT (on the Jetson Nano, `trtexec` can be found under `/usr/src/tensorrt/bin/trtexec`): ``` /usr/src/tensorrt/bin/trtexec --fp16 --onnx=../yolo/yolov4_1_3_608_608_static.onnx --explicitBatch --saveEngine=tensorrt_fp16.trt ``` ================================================ FILE: maskcam/common.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ CODEC_MP4 = "MP4" CODEC_H265 = "H265" CODEC_H264 = "H264" USBCAM_PROTOCOL = "v4l2://" # Invented by us since there's no URI for this RASPICAM_PROTOCOL = "argus://" # Invented by us since there's no URI for this CONFIG_FILE = "maskcam_config.txt" # Also used in nvinfer element # Available commands (to send internally, between processes or via MQTT) CMD_FILE_SAVE = "save_file" CMD_STREAMING_START = "streaming_start" CMD_STREAMING_STOP = "streaming_stop" CMD_INFERENCE_RESTART = "inference_restart" CMD_FILESERVER_RESTART = "fileserver_restart" CMD_STATUS_REQUEST = "status_request" ================================================ FILE: maskcam/config.py ================================================ import os import configparser from maskcam.common import CONFIG_FILE from maskcam.prints import print_common as print config = configparser.ConfigParser() config.read(CONFIG_FILE) config.sections() # Environment variables overriding config file values # Each row is: (ENV_VAR_NAME, (config-section, config-param)) ENV_CONFIG_OVERRIDES = ( ("MASKCAM_INPUT", ("maskcam", "default-input")), # Redundant with start.sh script ("MASKCAM_DEVICE_ADDRESS", ("maskcam", "device-address")), ("MASKCAM_DETECTION_THRESHOLD", ("face-processor", "detection-threshold")), ("MASKCAM_VOTING_THRESHOLD", ("face-processor", "voting-threshold")), ("MASKCAM_MIN_FACE_SIZE", ("face-processor", "min-face-size")), ("MASKCAM_DISABLE_TRACKER", ("face-processor", "disable-tracker")), ("MASKCAM_ALERT_MIN_VISIBLE_PEOPLE", ("maskcam", "alert-min-visible-people")), ("MASKCAM_ALERT_MAX_TOTAL_PEOPLE", ("maskcam", "alert-max-total-people")), ("MASKCAM_ALERT_NO_MASK_FRACTION", ("maskcam", "alert-no-mask-fraction")), ("MASKCAM_STATISTICS_PERIOD", ("maskcam", "statistics-period")), ("MASKCAM_TIMEOUT_INFERENCE_RESTART", ("maskcam", "timeout-inference-restart")), ("MASKCAM_CAMERA_FRAMERATE", ("maskcam", "camera-framerate")), ("MASKCAM_CAMERA_FLIP_METHOD", ("maskcam", "camera-flip-method")), ("MASKCAM_OUTPUT_VIDEO_WIDTH", ("maskcam", "output-video-width")), ("MASKCAM_OUTPUT_VIDEO_HEIGHT", ("maskcam", "output-video-height")), ("MASKCAM_INFERENCE_INTERVAL_AUTO", ("maskcam", "inference-interval-auto")), ("MASKCAM_INFERENCE_MAX_FPS", ("maskcam", "inference-max-fps")), ("MASKCAM_INFERENCE_LOG_INTERVAL", ("maskcam", "inference-log-interval")), ("MASKCAM_STREAMING_START_DEFAULT", ("maskcam", "streaming-start-default")), ("MASKCAM_STREAMING_PORT", ("maskcam", "streaming-port")), ("MASKCAM_FILESERVER_ENABLED", ("maskcam", "fileserver-enabled")), ("MASKCAM_FILESERVER_FORCE_SAVE", ("maskcam", "fileserver-force-save")), ("MASKCAM_FILESERVER_VIDEO_PERIOD", ("maskcam", "fileserver-video-period")), ("MASKCAM_FILESERVER_VIDEO_DURATION", ("maskcam", "fileserver-video-duration")), ("MASKCAM_FILESERVER_HDD_DIR", ("maskcam", "fileserver-hdd-dir")), ("MQTT_BROKER_IP", ("mqtt", "mqtt-broker-ip")), ("MQTT_BROKER_PORT", ("mqtt", "mqtt-broker-port")), ("MQTT_DEVICE_NAME", ("mqtt", "mqtt-device-name")), ("MQTT_DEVICE_DESCRIPTION", ("mqtt", "mqtt-device-description")), ) # Apply overrides for env_var, config_param in ENV_CONFIG_OVERRIDES: override_value = os.environ.get(env_var, None) if override_value is not None: config[config_param[0]][config_param[1]] = override_value def print_config_overrides(): # Leave prints separated so that it can be executed on demand # by one single process instead of each import for env_var, config_param in ENV_CONFIG_OVERRIDES: override_value = os.environ.get(env_var, None) if override_value is not None: print(f"\nConfig override {env_var}={override_value}") ================================================ FILE: maskcam/maskcam_filesave.py ================================================ #!/usr/bin/env python3 ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ import os import gi import pyds import sys import time import signal import platform import threading import multiprocessing as mp from datetime import datetime gi.require_version("Gst", "1.0") gi.require_version("GstBase", "1.0") gi.require_version("GstRtspServer", "1.0") from gi.repository import GLib, Gst, GstRtspServer, GstBase from .prints import print_filesave as print from .common import CODEC_MP4, CODEC_H264, CODEC_H265, CONFIG_FILE from .utils import glib_cb_restart from .config import config, print_config_overrides e_interrupt = None def make_elm_or_print_err(factoryname, name, printedname, detail=""): """Creates an element with Gst Element Factory make. Return the element if successfully created, otherwise print to stderr and return None. """ print("Creating", printedname) elm = Gst.ElementFactory.make(factoryname, name) if not elm: print("Unable to create " + printedname, error=True) if detail: print(detail) return elm def sigint_handler(sig, frame): # This function is not used if e_external_interrupt is provided print("[red]Ctrl+C pressed. Interrupting file-save...[/red]") e_interrupt.set() def main( config: dict, output_filename: str, udp_port: int, e_external_interrupt: mp.Event = None, ): global e_interrupt codec = config["maskcam"]["codec"] streaming_clock_rate = int(config["maskcam"]["streaming-clock-rate"]) udp_capabilities = f"application/x-rtp,media=video,encoding-name=(string){codec},clock-rate={streaming_clock_rate}" # Standard GStreamer initialization # GObject.threads_init() # Doesn't seem necessary (see https://pygobject.readthedocs.io/en/latest/guide/threading.html) Gst.init(None) # Create gstreamer elements # Create Pipeline element that will form a connection of other elements print( "[green]Creating:[/green] file-saving pipeline " f"UDP(port:{udp_port})->File({output_filename})" ) pipeline = Gst.Pipeline() if not pipeline: print("Unable to create Pipeline", error=True) udpsrc = make_elm_or_print_err("udpsrc", "udpsrc", "UDP Source") udpsrc.set_property("port", udp_port) udpsrc.set_property("buffer-size", 524288) udpsrc.set_property("caps", Gst.Caps.from_string(udp_capabilities)) rtpjitterbuffer = make_elm_or_print_err( "rtpjitterbuffer", "rtpjitterbuffer", "RTP Jitter Buffer" ) # Default mode is 1 (slave), acts as a live source and gets laggy rtpjitterbuffer.set_property("mode", 4) # caps_udp = make_elm_or_print_err("capsfilter", "caps_udp", "UDP RTP capabilities") # caps_udp.set_property("caps", Gst.Caps.from_string(udp_capabilities)) if codec == CODEC_MP4: print("Creating MPEG-4 payload decoder") rtpdepay = make_elm_or_print_err("rtpmp4vpay", "rtpdepay", "RTP MPEG-4 Payload Decoder") codeparser = make_elm_or_print_err("mpeg4videoparse", "mpeg4-parser", "Code Parser") elif codec == CODEC_H264: print("Creating H264 payload decoder") rtpdepay = make_elm_or_print_err("rtph264depay", "rtpdepay", "RTP H264 Payload Decoder") codeparser = make_elm_or_print_err("h264parse", "h264-parser", "Code Parser") else: # Default: H265 (recommended) print("Creating H265 payload decoder") rtpdepay = make_elm_or_print_err("rtph265depay", "rtpdepay", "RTP H265 Payload Decoder") codeparser = make_elm_or_print_err("h265parse", "h265-parser", "Code Parser") # Workaround for this issue: https://gitlab.freedesktop.org/gstreamer/gst-plugins-good/-/issues/410 GstBase.BaseParse.set_pts_interpolation(codeparser, True) container = make_elm_or_print_err("qtmux", "qtmux", "Container") filesink = make_elm_or_print_err("filesink", "filesink", "File Sink") filesink.set_property("location", output_filename) # filesink.set_property("sync", False) # filesink.set_property("async", False) pipeline.add(udpsrc) pipeline.add(rtpjitterbuffer) # pipeline.add(caps_udp) pipeline.add(rtpdepay) pipeline.add(codeparser) pipeline.add(container) pipeline.add(filesink) # Pipeline Links udpsrc.link(rtpjitterbuffer) rtpjitterbuffer.link(rtpdepay) # caps_udp.link(rtpdepay) rtpdepay.link(codeparser) codeparser.link(container) container.link(filesink) # GLib loop required for RTSP server g_loop = GLib.MainLoop() g_context = g_loop.get_context() # GStreamer message bus bus = pipeline.get_bus() if e_external_interrupt is None: # Use threading instead of mp.Event() for sigint_handler, see: # https://bugs.python.org/issue41606 e_interrupt = threading.Event() signal.signal(signal.SIGINT, sigint_handler) print("[green bold]Press Ctrl+C to save video and exit[/green bold]") else: # If there's an external interrupt, don't capture SIGINT e_interrupt = e_external_interrupt # Periodic gloop interrupt (see utils.glib_cb_restart) t_check = 50 GLib.timeout_add(t_check, glib_cb_restart, t_check) # Custom event loop, allows saving file on Ctrl+C press running = True # start play back and listen to events pipeline.set_state(Gst.State.PLAYING) print("[green]Playing:[/green] file-saving pipeline UDP->File\n") while running: g_context.iteration(may_block=True) message = bus.pop() if message is not None: t = message.type if t == Gst.MessageType.EOS: print(f"File saved: [yellow]{output_filename}[/yellow]") running = False elif t == Gst.MessageType.WARNING: err, debug = message.parse_warning() print("%s: %s" % (err, debug), warning=True) elif t == Gst.MessageType.ERROR: err, debug = message.parse_error() print("%s: %s" % (err, debug), error=True) running = False if e_interrupt.is_set(): print("Interruption received. Sending EOS to generate video file.") # This will allow the filesink to create a readable mp4 file container.send_event(Gst.Event.new_eos()) e_interrupt.clear() print("File-saver main loop ending.") # cleanup pipeline.set_state(Gst.State.NULL) if __name__ == "__main__": # Print any ENV var config override to avoid confusions print_config_overrides() # Check arguments output_filename = None udp_port = None if len(sys.argv) > 1: output_filename = sys.argv[1] if len(sys.argv) > 2: udp_port = int(sys.argv[2]) if not output_filename: output_dir = config["maskcam"]["fileserver-hdd-dir"] output_filename = f"{output_dir}/{datetime.today().strftime('%Y%m%d_%H%M%S')}.mp4" if not udp_port: # Use first listed in config udp_port = int(config["maskcam"]["udp-ports-filesave"].split(",")[0]) print(f"Output file: {output_filename}") sys.exit(main(config=config, output_filename=output_filename, udp_port=udp_port)) ================================================ FILE: maskcam/maskcam_fileserver.py ================================================ #!/usr/bin/env python3 ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ import os import sys import time import socket import threading import multiprocessing as mp from datetime import datetime from http.server import SimpleHTTPRequestHandler from socketserver import TCPServer, ThreadingTCPServer from .config import config, print_config_overrides from .utils import get_ip_address from .prints import print_fileserver as print class Handler(SimpleHTTPRequestHandler): # Needed to set extensions_map pass def start_server(httpd_server): httpd_server.serve_forever(poll_interval=0.5) def cb_handle_error(request, client_address): # Not important, happens very often but nothing actually fails print(f"Static file server: File request interrupted [client: {client_address}]") def main(config, directory=None, e_external_interrupt: mp.Event = None): if directory is None: directory = config["maskcam"]["fileserver-hdd-dir"] directory = os.fspath(directory) print(f"Serving static files from directory: [yellow]{directory}[/yellow]") port = int(config["maskcam"]["fileserver-port"]) # Create dir if doesn't exist os.system(f"mkdir -p {directory}") os.chdir(directory) # easiest way # Force download mp4 files Handler.extensions_map[".mp4"] = "application/octet-stream" print(f"[green]Static server STARTED[/green] at http://{get_ip_address()}:{port}") with ThreadingTCPServer(("", port), Handler) as httpd: httpd.handle_error = cb_handle_error s = threading.Thread(target=start_server, args=(httpd,)) s.start() try: if e_external_interrupt is not None: e_external_interrupt.wait() # blocking else: s.join() # blocking except KeyboardInterrupt: print("Ctrl+C pressed") print("Shutting down static file server") httpd.shutdown() httpd.server_close() s.join(timeout=1) if s.is_alive(): print("Server thread did not stop", warning=True) else: print("Server shut down correctly") print(f"Server alive threads: {threading.enumerate()}") if __name__ == "__main__": # Print any ENV var config override to avoid confusions print_config_overrides() # Input source directory = sys.argv[1] if len(sys.argv) > 1 else None main(config, directory=directory) ================================================ FILE: maskcam/maskcam_inference.py ================================================ #!/usr/bin/env python3 ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ import os import gi import pyds import sys import ipdb import time import signal import platform import threading import numpy as np import multiprocessing as mp from rich.console import Console from datetime import datetime, timezone gi.require_version("Gst", "1.0") gi.require_version("GstRtspServer", "1.0") from gi.repository import GLib, Gst, GstRtspServer from norfair.tracker import Tracker, Detection from .config import config, print_config_overrides from .prints import print_inference as print from .common import ( CODEC_MP4, CODEC_H264, CODEC_H265, USBCAM_PROTOCOL, RASPICAM_PROTOCOL, CONFIG_FILE, ) from .utils import glib_cb_restart, load_udp_ports_filesaving # YOLO labels. See obj.names file LABEL_MASK = "mask" LABEL_NO_MASK = "no_mask" # YOLOv4: no_mask LABEL_MISPLACED = "misplaced" LABEL_NOT_VISIBLE = "not_visible" FRAMES_LOG_INTERVAL = int(config["maskcam"]["inference-log-interval"]) # Global vars frame_number = 0 start_time = None end_time = None console = Console() e_interrupt = None class FaceMaskProcessor: def __init__( self, th_detection=0, th_vote=0, min_face_size=0, tracker_period=1, disable_tracker=False ): self.people_votes = {} self.current_people = set() self.th_detection = th_detection self.th_vote = th_vote self.tracker_period = tracker_period self.min_face_size = min_face_size self.disable_detection_validation = False self.min_votes = 5 self.max_votes = 50 self.color_mask = (0.0, 1.0, 0.0) # green self.color_no_mask = (1.0, 0.0, 0.0) # red self.color_unknown = (1.0, 1.0, 0.0) # yellow self.draw_raw_detections = disable_tracker self.draw_tracked_people = not disable_tracker self.stats_lock = threading.Lock() # Norfair Tracker if disable_tracker: self.tracker = None else: self.tracker = Tracker( distance_function=self.keypoints_distance, detection_threshold=self.th_detection, distance_threshold=1, point_transience=8, hit_inertia_min=15, hit_inertia_max=45, ) def keypoints_distance(self, detected_pose, tracked_pose): detected_points = detected_pose.points estimated_pose = tracked_pose.estimate min_box_size = min( max( detected_points[1][0] - detected_points[0][0], # x2 - x1 detected_points[1][1] - detected_points[0][1], # y2 - y1 1, ), max( estimated_pose[1][0] - estimated_pose[0][0], # x2 - x1 estimated_pose[1][1] - estimated_pose[0][1], # y2 - y1 1, ), ) mean_distance_normalized = ( np.mean(np.linalg.norm(detected_points - estimated_pose, axis=1)) / min_box_size ) return mean_distance_normalized def validate_detection(self, box_points, score, label): if self.disable_detection_validation: return True box_width = box_points[1][0] - box_points[0][0] box_height = box_points[1][1] - box_points[0][1] return min(box_width, box_height) >= self.min_face_size and score >= self.th_detection def add_detection(self, person_id, label, score): # This function is called from streaming thread with self.stats_lock: self.current_people.add(person_id) if person_id not in self.people_votes: self.people_votes[person_id] = 0 if score > self.th_vote: if label == LABEL_MASK: self.people_votes[person_id] += 1 elif label == LABEL_NO_MASK or LABEL_MISPLACED: self.people_votes[person_id] -= 1 # max_votes limit self.people_votes[person_id] = np.clip( self.people_votes[person_id], -self.max_votes, self.max_votes ) def get_person_label(self, person_id): person_votes = self.people_votes[person_id] if abs(person_votes) >= self.min_votes: color = self.color_mask if person_votes > 0 else self.color_no_mask label = "mask" if person_votes > 0 else "no mask" else: color = self.color_unknown label = "not visible" return f"{person_id}|{label}({abs(person_votes)})", color def get_instant_statistics(self, refresh=True): """ Get statistics only including people that appeared on camera since last refresh """ instant_stats = self.get_statistics(filter_ids=self.current_people) if refresh: with self.stats_lock: self.current_people = set() return instant_stats def get_statistics(self, filter_ids=None): with self.stats_lock: if filter_ids is not None: filtered_people = { id: votes for id, votes in self.people_votes.items() if id in filter_ids } else: filtered_people = self.people_votes total_people = len(filtered_people) total_classified = 0 total_mask = 0 for person_id in filtered_people: person_votes = filtered_people[person_id] if abs(person_votes) >= self.min_votes: total_classified += 1 if person_votes > 0: total_mask += 1 return total_people, total_classified, total_mask def cb_add_statistics(cb_args): stats_period, stats_queue, face_processor = cb_args people_total, people_classified, people_mask = face_processor.get_instant_statistics( refresh=True ) people_no_mask = people_classified - people_mask # stats_queue is an mp.Queue optionally provided externally (in main()) stats_queue.put_nowait( { "people_total": people_total, "people_with_mask": people_mask, "people_without_mask": people_no_mask, "timestamp": datetime.timestamp(datetime.now(timezone.utc)), } ) # Next report timeout GLib.timeout_add_seconds(stats_period, cb_add_statistics, cb_args) def sigint_handler(sig, frame): # This function is not used if e_external_interrupt is provided print("[red]Ctrl+C pressed. Interrupting inference...[/red]") e_interrupt.set() def is_aarch64(): return platform.uname()[4] == "aarch64" def draw_detection(display_meta, n_draw, box_points, detection_label, color): # print(f"Drawing {n_draw} | {detection_label}") # print(box_points) rect = display_meta.rect_params[n_draw] ((x1, y1), (x2, y2)) = box_points rect.left = x1 rect.top = y1 rect.width = x2 - x1 rect.height = y2 - y1 # print(f"{x1} {y1}, {x2} {y2}") # Bug: bg color is always green # rect.has_bg_color = True # rect.bg_color.set(0.5, 0.5, 0.5, 0.6) # RGBA rect.border_color.set(*color, 1.0) rect.border_width = 2 label = display_meta.text_params[n_draw] label.x_offset = x1 label.y_offset = y2 label.font_params.font_name = "Verdana" label.font_params.font_size = 9 label.font_params.font_color.set(0, 0, 0, 1.0) # Black # label.display_text = f"{person.id} | {detection_p:.2f}" label.display_text = detection_label label.set_bg_clr = True label.text_bg_clr.set(*color, 0.5) display_meta.num_rects = n_draw + 1 display_meta.num_labels = n_draw + 1 def cb_buffer_probe(pad, info, cb_args): global frame_number global start_time face_processor, e_ready = cb_args gst_buffer = info.get_buffer() if not gst_buffer: print("Unable to get GstBuffer", error=True) return # Set e_ready event to notify the pipeline is working (e.g: for orchestrator) if e_ready is not None and not e_ready.is_set(): print("Inference pipeline setting [green]e_ready[/green]") e_ready.set() # Retrieve batch metadata from the gst_buffer # Note that pyds.gst_buffer_get_nvds_batch_meta() expects the # C address of gst_buffer as input, which is obtained with hash(gst_buffer) batch_meta = pyds.gst_buffer_get_nvds_batch_meta(hash(gst_buffer)) l_frame = batch_meta.frame_meta_list while l_frame is not None: try: # Note that l_frame.data needs a cast to pyds.NvDsFrameMeta # The casting is done by pyds.glist_get_nvds_frame_meta() # The casting also keeps ownership of the underlying memory # in the C code, so the Python garbage collector will leave # it alone. # frame_meta = pyds.glist_get_nvds_frame_meta(l_frame.data) frame_meta = pyds.NvDsFrameMeta.cast(l_frame.data) except StopIteration: break frame_number = frame_meta.frame_num # num_detections = frame_meta.num_obj_meta l_obj = frame_meta.obj_meta_list detections = [] obj_meta_list = [] while l_obj is not None: try: # Casting l_obj.data to pyds.NvDsObjectMeta # obj_meta=pyds.glist_get_nvds_object_meta(l_obj.data) obj_meta = pyds.NvDsObjectMeta.cast(l_obj.data) except StopIteration: break obj_meta_list.append(obj_meta) obj_meta.rect_params.border_color.set(0.0, 0.0, 1.0, 0.0) box = obj_meta.rect_params # print(f"{obj_meta.obj_label} | {obj_meta.confidence}") box_points = ( (box.left, box.top), (box.left + box.width, box.top + box.height), ) box_p = obj_meta.confidence box_label = obj_meta.obj_label if face_processor.validate_detection(box_points, box_p, box_label): det_data = {"label": box_label, "p": box_p} detections.append( Detection( np.array(box_points), data=det_data, ) ) # print(f"Added detection: {det_data}") try: l_obj = l_obj.next except StopIteration: break # Remove all object meta to avoid drawing. Do this outside while since we're modifying list for obj_meta in obj_meta_list: # Remove this to avoid drawing label texts pyds.nvds_remove_obj_meta_from_frame(frame_meta, obj_meta) obj_meta_list = None # Each meta object carries max 16 rects/labels/etc. max_drawings_per_meta = 16 # This is hardcoded, not documented if face_processor.tracker is not None: # Track, count and draw tracked people tracked_people = face_processor.tracker.update( detections, period=face_processor.tracker_period ) # Filter out people with no live points (don't draw) drawn_people = [person for person in tracked_people if person.live_points.any()] if face_processor.draw_tracked_people: for n_person, person in enumerate(drawn_people): points = person.estimate box_points = points.clip(0).astype(int) # Update mask votes face_processor.add_detection( person.id, person.last_detection.data["label"], person.last_detection.data["p"], ) label, color = face_processor.get_person_label(person.id) # Index of this person's drawing in the current meta n_draw = n_person % max_drawings_per_meta if n_draw == 0: # Initialize meta # Acquiring a display meta object. The memory ownership remains in # the C code so downstream plugins can still access it. Otherwise # the garbage collector will claim it when this probe function exits. display_meta = pyds.nvds_acquire_display_meta_from_pool(batch_meta) pyds.nvds_add_display_meta_to_frame(frame_meta, display_meta) draw_detection(display_meta, n_draw, box_points, label, color) # Raw detections if face_processor.draw_raw_detections: for n_detection, detection in enumerate(detections): points = detection.points box_points = points.clip(0).astype(int) label = detection.data["label"] if label == LABEL_MASK: color = face_processor.color_mask elif label == LABEL_NO_MASK or label == LABEL_MISPLACED: color = face_processor.color_no_mask else: color = face_processor.color_unknown label = f"{label} | {detection.data['p']:.2f}" n_draw = n_detection % max_drawings_per_meta if n_draw == 0: # Initialize meta # Acquiring a display meta object. The memory ownership remains in # the C code so downstream plugins can still access it. Otherwise # the garbage collector will claim it when this probe function exits. display_meta = pyds.nvds_acquire_display_meta_from_pool(batch_meta) pyds.nvds_add_display_meta_to_frame(frame_meta, display_meta) draw_detection(display_meta, n_draw, box_points, label, color) # Using pyds.get_string() to get display_text as string # print(pyds.get_string(py_nvosd_text_params.display_text)) # print(".", end="", flush=True) # print("") if not frame_number % FRAMES_LOG_INTERVAL: print(f"Processed {frame_number} frames...") try: l_frame = l_frame.next except StopIteration: break # Start timer at the end of first frame processing if start_time is None: start_time = time.time() return Gst.PadProbeReturn.OK def cb_newpad(decodebin, decoder_src_pad, data): print("In cb_newpad\n") caps = decoder_src_pad.get_current_caps() gststruct = caps.get_structure(0) gstname = gststruct.get_name() source_bin = data features = caps.get_features(0) # Need to check if the pad created by the decodebin is for video and not # audio. print("gstname=", gstname) if gstname.find("video") != -1: # Link the decodebin pad only if decodebin has picked nvidia # decoder plugin nvdec_*. We do this by checking if the pad caps contain # NVMM memory features. print("features=", features) if features.contains("memory:NVMM"): # Get the source bin ghost pad bin_ghost_pad = source_bin.get_static_pad("src") if not bin_ghost_pad.set_target(decoder_src_pad): print("Failed to link decoder src pad to source bin ghost pad", error=True) else: print("Decodebin did not pick nvidia decoder plugin", error=True) def decodebin_child_added(child_proxy, Object, name, user_data): print(f"Decodebin child added: {name}") if name.find("decodebin") != -1: Object.connect("child-added", decodebin_child_added, user_data) if is_aarch64() and name.find("nvv4l2decoder") != -1: Object.set_property("bufapi-version", True) def create_source_bin(index, uri): print("Creating source bin") # Create a source GstBin to abstract this bin's content from the rest of the # pipeline bin_name = "source-bin-%02d" % index print(bin_name) nbin = Gst.Bin.new(bin_name) if not nbin: print("Unable to create source bin", error=True) # Source element for reading from the uri. # We will use decodebin and let it figure out the container format of the # stream and the codec and plug the appropriate demux and decode plugins. uri_decode_bin = Gst.ElementFactory.make("uridecodebin", "uri-decode-bin") if not uri_decode_bin: print("Unable to create uri decode bin", error=True) # We set the input uri to the source element uri_decode_bin.set_property("uri", uri) # Connect to the "pad-added" signal of the decodebin which generates a # callback once a new pad for raw data has beed created by the decodebin uri_decode_bin.connect("pad-added", cb_newpad, nbin) uri_decode_bin.connect("child-added", decodebin_child_added, nbin) # We need to create a ghost pad for the source bin which will act as a proxy # for the video decoder src pad. The ghost pad will not have a target right # now. Once the decode bin creates the video decoder and generates the # cb_newpad callback, we will set the ghost pad target to the video decoder # src pad. Gst.Bin.add(nbin, uri_decode_bin) bin_pad = nbin.add_pad(Gst.GhostPad.new_no_target("src", Gst.PadDirection.SRC)) if not bin_pad: print("Failed to add ghost pad in source bin", error=True) return None return nbin def make_elm_or_print_err(factoryname, name, printedname): """Creates an element with Gst Element Factory make. Return the element if successfully created, otherwise print to stderr and return None. """ print("Creating", printedname) elm = Gst.ElementFactory.make(factoryname, name) if not elm: print("Unable to create ", printedname, error=True) show_troubleshooting() return elm def show_troubleshooting(): # On Jetson, there is a problem with the encoder failing to initialize # due to limitation on TLS usage. To work around this, preload libgomp. # Add a reminder here in case the user forgets. print( """ [yellow]TROUBLESHOOTING HELP[/yellow] [yellow]If the error is like: v4l-camera-source / reason not-negotiated[/yellow] [green]Solution:[/green] configure camera capabilities Run the script under utils/gst_capabilities.sh and find the lines with type video/x-raw ... Find a suitable framerate=X/1 (with X being an integer like 24, 15, etc.) Then edit config_maskcam.txt and change the line: camera-framerate=X Or configure using --env MASKCAM_CAMERA_FRAMERATE=X (see README) [yellow]If the error is like: /usr/lib/aarch64-linux-gnu/libgomp.so.1: cannot allocate memory in static TLS block[/yellow] [green]Solution:[/green] preload the offending library export LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libgomp.so.1 [yellow]END HELP[/yellow] """ ) def main( config: dict, input_filename: str, output_filename: str = None, e_external_interrupt: mp.Event = None, stats_queue: mp.Queue = None, e_ready: mp.Event = None, ): global frame_number global start_time global end_time global e_interrupt # Load all udp ports to output video udp_ports = {int(config["maskcam"]["udp-port-streaming"])} load_udp_ports_filesaving(config, udp_ports) codec = config["maskcam"]["codec"] stats_period = int(config["maskcam"]["statistics-period"]) # Original: 1920x1080, bdti_resized: 1024x576, yolo-input: 1024x608 output_width = int(config["maskcam"]["output-video-width"]) output_height = int(config["maskcam"]["output-video-height"]) output_bitrate = 6000000 # Nice for h264@1024x576: 4000000 # Two types of camera supported: USB or Raspi usbcam_input = USBCAM_PROTOCOL in input_filename raspicam_input = RASPICAM_PROTOCOL in input_filename camera_input = usbcam_input or raspicam_input if camera_input: camera_framerate = int(config["maskcam"]["camera-framerate"]) camera_flip_method = int(config["maskcam"]["camera-flip-method"]) # Set nvinfer.interval (number of frames to skip inference and use tracker instead) if camera_input and int(config["maskcam"]["inference-interval-auto"]): max_fps = int(config["maskcam"]["inference-max-fps"]) skip_inference = camera_framerate // max_fps print(f"Auto calculated frames to skip inference: {skip_inference}") else: skip_inference = int(config["property"]["interval"]) print(f"Configured frames to skip inference: {skip_inference}") # FaceMask initialization face_tracker_period = skip_inference + 1 # tracker_period=skipped + inference frame(1) face_detection_threshold = float(config["face-processor"]["detection-threshold"]) face_voting_threshold = float(config["face-processor"]["voting-threshold"]) face_min_face_size = int(config["face-processor"]["min-face-size"]) face_disable_tracker = int(config["face-processor"]["disable-tracker"]) face_processor = FaceMaskProcessor( th_detection=face_detection_threshold, th_vote=face_voting_threshold, min_face_size=face_min_face_size, tracker_period=face_tracker_period, disable_tracker=face_disable_tracker, ) # Standard GStreamer initialization Gst.init(None) # Create gstreamer elements # Create Pipeline element that will form a connection of other elements print("Creating Pipeline \n ") pipeline = Gst.Pipeline() if not pipeline: print("Unable to create Pipeline", error=True) if camera_input: if usbcam_input: input_device = input_filename[len(USBCAM_PROTOCOL) :] source = make_elm_or_print_err("v4l2src", "v4l2-camera-source", "Camera input") source.set_property("device", input_device) nvvidconvsrc = make_elm_or_print_err( "nvvideoconvert", "convertor_src2", "Convertor src 2" ) # Input camera configuration # Use ./gst_capabilities.sh to get the list of available capabilities from /dev/video0 camera_capabilities = f"video/x-raw, framerate={camera_framerate}/1" elif raspicam_input: input_device = input_filename[len(RASPICAM_PROTOCOL) :] source = make_elm_or_print_err( "nvarguscamerasrc", "nv-argus-camera-source", "RaspiCam input" ) source.set_property("sensor-id", int(input_device)) source.set_property("bufapi-version", 1) # Special camera_capabilities for raspicam camera_capabilities = f"video/x-raw(memory:NVMM),framerate={camera_framerate}/1" nvvidconvsrc = make_elm_or_print_err("nvvidconv", "convertor_flip", "Convertor flip") nvvidconvsrc.set_property("flip-method", camera_flip_method) # Misterious converting sequence from deepstream_test_1_usb.py caps_camera = make_elm_or_print_err("capsfilter", "camera_src_caps", "Camera caps filter") caps_camera.set_property( "caps", Gst.Caps.from_string(camera_capabilities), ) vidconvsrc = make_elm_or_print_err("videoconvert", "convertor_src1", "Convertor src 1") caps_vidconvsrc = make_elm_or_print_err( "capsfilter", "nvmm_caps", "NVMM caps for input stream" ) caps_vidconvsrc.set_property("caps", Gst.Caps.from_string("video/x-raw(memory:NVMM)")) else: source_bin = create_source_bin(0, input_filename) # Create nvstreammux instance to form batches from one or more sources. streammux = make_elm_or_print_err("nvstreammux", "Stream-muxer", "NvStreamMux") streammux.set_property("width", output_width) streammux.set_property("height", output_height) streammux.set_property("enable-padding", True) # Keeps aspect ratio, but adds black margin streammux.set_property("batch-size", 1) streammux.set_property("batched-push-timeout", 4000000) # Adding this element after muxer will cause detections to get delayed # videorate = make_elm_or_print_err("videorate", "Vide-rate", "Video Rate") # Inference element: object detection using TRT engine pgie = make_elm_or_print_err("nvinfer", "primary-inference", "pgie") pgie.set_property("config-file-path", CONFIG_FILE) pgie.set_property("interval", skip_inference) # Use convertor to convert from NV12 to RGBA as required by nvosd convert_pre_osd = make_elm_or_print_err( "nvvideoconvert", "convert_pre_osd", "Converter NV12->RGBA" ) # OSD: to draw on the RGBA buffer nvosd = make_elm_or_print_err("nvdsosd", "onscreendisplay", "OSD (nvosd)") nvosd.set_property("process-mode", 2) # 0: CPU Mode, 1: GPU (only dGPU), 2: VIC (Jetson only) # nvosd.set_property("display-bbox", False) # Bug: Removes all squares nvosd.set_property("display-clock", False) nvosd.set_property("display-text", True) # Needed for any text # Finally encode and save the osd output queue = make_elm_or_print_err("queue", "queue", "Queue") convert_post_osd = make_elm_or_print_err( "nvvideoconvert", "convert_post_osd", "Converter RGBA->NV12" ) # Video capabilities: check format and GPU/CPU location capsfilter = make_elm_or_print_err("capsfilter", "capsfilter", "capsfilter") if codec == CODEC_MP4: # Not hw accelerated caps = Gst.Caps.from_string("video/x-raw, format=I420") else: # hw accelerated caps = Gst.Caps.from_string("video/x-raw(memory:NVMM), format=I420") capsfilter.set_property("caps", caps) # Encoder: H265 has more efficient compression if codec == CODEC_MP4: print("Creating MPEG-4 stream") encoder = make_elm_or_print_err("avenc_mpeg4", "encoder", "Encoder") codeparser = make_elm_or_print_err("mpeg4videoparse", "mpeg4-parser", "Code Parser") rtppay = make_elm_or_print_err("rtpmp4vpay", "rtppay", "RTP MPEG-44 Payload") elif codec == CODEC_H264: print("Creating H264 stream") encoder = make_elm_or_print_err("nvv4l2h264enc", "encoder", "Encoder") encoder.set_property("preset-level", 1) encoder.set_property("bufapi-version", 1) codeparser = make_elm_or_print_err("h264parse", "h264-parser", "Code Parser") rtppay = make_elm_or_print_err("rtph264pay", "rtppay", "RTP H264 Payload") else: # Default: H265 (recommended) print("Creating H265 stream") encoder = make_elm_or_print_err("nvv4l2h265enc", "encoder", "Encoder") encoder.set_property("preset-level", 1) encoder.set_property("bufapi-version", 1) codeparser = make_elm_or_print_err("h265parse", "h265-parser", "Code Parser") rtppay = make_elm_or_print_err("rtph265pay", "rtppay", "RTP H265 Payload") encoder.set_property("insert-sps-pps", 1) encoder.set_property("bitrate", output_bitrate) splitter_file_udp = make_elm_or_print_err("tee", "tee_file_udp", "Splitter file/UDP") # UDP streaming queue_udp = make_elm_or_print_err("queue", "queue_udp", "UDP queue") multiudpsink = make_elm_or_print_err("multiudpsink", "multi udpsink", "Multi UDP Sink") # udpsink.set_property("host", "127.0.0.1") # udpsink.set_property("port", udp_port) # Comma separated list of clients, don't add spaces :S client_list = [f"127.0.0.1:{udp_port}" for udp_port in udp_ports] multiudpsink.set_property("clients", ",".join(client_list)) multiudpsink.set_property("async", False) multiudpsink.set_property("sync", True) if output_filename is not None: queue_file = make_elm_or_print_err("queue", "queue_file", "File save queue") # codeparser already created above depending on codec container = make_elm_or_print_err("qtmux", "qtmux", "Container") filesink = make_elm_or_print_err("filesink", "filesink", "File Sink") filesink.set_property("location", output_filename) else: # Fake sink, no save fakesink = make_elm_or_print_err("fakesink", "fakesink", "Fake Sink") # Add elements to the pipeline if camera_input: pipeline.add(source) pipeline.add(caps_camera) pipeline.add(vidconvsrc) pipeline.add(nvvidconvsrc) pipeline.add(caps_vidconvsrc) else: pipeline.add(source_bin) pipeline.add(streammux) pipeline.add(pgie) pipeline.add(convert_pre_osd) pipeline.add(nvosd) pipeline.add(queue) pipeline.add(convert_post_osd) pipeline.add(capsfilter) pipeline.add(encoder) pipeline.add(splitter_file_udp) if output_filename is not None: pipeline.add(queue_file) pipeline.add(codeparser) pipeline.add(container) pipeline.add(filesink) else: pipeline.add(fakesink) # Output to UDP pipeline.add(queue_udp) pipeline.add(rtppay) pipeline.add(multiudpsink) print("Linking elements in the Pipeline \n") # Pipeline Links if camera_input: source.link(caps_camera) caps_camera.link(vidconvsrc) vidconvsrc.link(nvvidconvsrc) nvvidconvsrc.link(caps_vidconvsrc) srcpad = caps_vidconvsrc.get_static_pad("src") else: srcpad = source_bin.get_static_pad("src") sinkpad = streammux.get_request_pad("sink_0") if not srcpad or not sinkpad: print("Unable to get file source or mux sink pads", error=True) srcpad.link(sinkpad) streammux.link(pgie) pgie.link(convert_pre_osd) convert_pre_osd.link(nvosd) nvosd.link(queue) queue.link(convert_post_osd) convert_post_osd.link(capsfilter) capsfilter.link(encoder) encoder.link(splitter_file_udp) # Split stream to file and rtsp tee_file = splitter_file_udp.get_request_pad("src_%u") tee_udp = splitter_file_udp.get_request_pad("src_%u") # Output to File or fake sinks if output_filename is not None: tee_file.link(queue_file.get_static_pad("sink")) queue_file.link(codeparser) codeparser.link(container) container.link(filesink) else: tee_file.link(fakesink.get_static_pad("sink")) # Output to UDP tee_udp.link(queue_udp.get_static_pad("sink")) queue_udp.link(rtppay) rtppay.link(multiudpsink) # Lets add probe to get informed of the meta data generated, we add probe to # the sink pad of the osd element, since by that time, the buffer would have # had got all the metadata. osdsinkpad = nvosd.get_static_pad("sink") if not osdsinkpad: print("Unable to get sink pad of nvosd", error=True) cb_args = (face_processor, e_ready) osdsinkpad.add_probe(Gst.PadProbeType.BUFFER, cb_buffer_probe, cb_args) # GLib loop required for RTSP server g_loop = GLib.MainLoop() g_context = g_loop.get_context() # GStreamer message bus bus = pipeline.get_bus() if e_external_interrupt is None: # Use threading instead of mp.Event() for sigint_handler, see: # https://bugs.python.org/issue41606 e_interrupt = threading.Event() signal.signal(signal.SIGINT, sigint_handler) print("[green bold]Press Ctrl+C to stop pipeline[/green bold]") else: # If there's an external interrupt, don't capture SIGINT e_interrupt = e_external_interrupt # start play back and listen to events pipeline.set_state(Gst.State.PLAYING) # After setting pipeline to PLAYING, stop it even on exceptions try: time_start_playing = time.time() # Timer to add statistics to queue if stats_queue is not None: cb_args = stats_period, stats_queue, face_processor GLib.timeout_add_seconds(stats_period, cb_add_statistics, cb_args) # Periodic gloop interrupt (see utils.glib_cb_restart) t_check = 100 GLib.timeout_add(t_check, glib_cb_restart, t_check) # Custom event loop running = True while running: g_context.iteration(may_block=True) message = bus.pop() if message is not None: t = message.type if t == Gst.MessageType.EOS: print("End-of-stream\n") running = False elif t == Gst.MessageType.WARNING: err, debug = message.parse_warning() print(f"{err}: {debug}", warning=True) elif t == Gst.MessageType.ERROR: err, debug = message.parse_error() print(f"{err}: {debug}", error=True) show_troubleshooting() running = False if e_interrupt.is_set(): # Send EOS to container to generate a valid mp4 file if output_filename is not None: container.send_event(Gst.Event.new_eos()) multiudpsink.send_event(Gst.Event.new_eos()) else: pipeline.send_event(Gst.Event.new_eos()) # fakesink EOS won't work end_time = time.time() print("Inference main loop ending.") pipeline.set_state(Gst.State.NULL) # Profiling display if start_time is not None and end_time is not None: total_time = end_time - start_time total_frames = frame_number inference_frames = total_frames // (skip_inference + 1) print() print(f"[bold yellow] ---- Profiling ---- [/bold yellow]") print(f"Inference frames: {inference_frames} | Processed frames: {total_frames}") print(f"Time from time_start_playing: {end_time - time_start_playing:.2f} seconds") print(f"Total time skipping first inference: {total_time:.2f} seconds") print(f"Avg. time/frame: {total_time/total_frames:.4f} secs") print(f"[bold yellow]FPS: {total_frames/total_time:.1f} frames/second[/bold yellow]\n") if skip_inference != 0: print( "[red]NOTE: FPS calculated skipping inference every" f" interval={skip_inference} frames[/red]" ) if output_filename is not None: print(f"Output file saved: [green bold]{output_filename}[/green bold]") except: console.print_exception() pipeline.set_state(Gst.State.NULL) if __name__ == "__main__": print_config_overrides() # Check input arguments output_filename = None if len(sys.argv) > 1: input_filename = sys.argv[1] print(f"Provided input source: {input_filename}") if len(sys.argv) > 2: output_filename = sys.argv[2] print(f"Save output file: [green]{output_filename}[/green]") else: input_filename = config["maskcam"]["default-input"] print(f"Using input from config file: {input_filename}") sys.exit( main( config=config, input_filename=input_filename, output_filename=output_filename, ) ) ================================================ FILE: maskcam/maskcam_streaming.py ================================================ #!/usr/bin/env python3 ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ import gi import pyds import sys import time import signal import platform import threading import multiprocessing as mp from datetime import datetime gi.require_version("Gst", "1.0") gi.require_version("GstRtspServer", "1.0") from gi.repository import GLib, Gst, GstRtspServer from .config import config, print_config_overrides from .prints import print_streaming as print from .utils import get_ip_address, glib_cb_restart, get_streaming_address from .common import CODEC_MP4, CODEC_H264, CODEC_H265, CONFIG_FILE e_interrupt = None def sigint_handler(sig, frame): # This function is not used if e_external_interrupt is provided print("[red]Ctrl+C pressed. Interrupting streaming...[/red]") e_interrupt.set() def main(config, e_external_interrupt: mp.Event = None): global e_interrupt udp_port = int(config["maskcam"]["udp-port-streaming"]) codec = config["maskcam"]["codec"] # Streaming address: rtsp://:/ rtsp_port = int(config["maskcam"]["streaming-port"]) rtsp_address = config["maskcam"]["streaming-path"] streaming_clock_rate = int(config["maskcam"]["streaming-clock-rate"]) # udp_capabilities = f"application/x-rtp,media=video,encoding-name={codec},payload=96" print(f"Codec: {codec}") # Standard GStreamer initialization Gst.init(None) # Start streaming server = GstRtspServer.RTSPServer.new() server.props.service = str(rtsp_port) server.attach(None) factory = GstRtspServer.RTSPMediaFactory.new() factory.set_launch( f"( udpsrc name=pay0 port={udp_port} buffer-size=524288" f' caps="application/x-rtp, media=video, clock-rate={streaming_clock_rate},' f' encoding-name=(string){codec}, payload=96 " )' ) factory.set_shared(True) server.get_mount_points().add_factory(rtsp_address, factory) streaming_address = get_streaming_address(get_ip_address(), rtsp_port, rtsp_address) print(f"\n\n[green bold]Streaming[/green bold] at {streaming_address}\n\n") # GLib loop required for RTSP server g_loop = GLib.MainLoop() g_context = g_loop.get_context() if e_external_interrupt is None: # Use threading instead of mp.Event() for sigint_handler, see: # https://bugs.python.org/issue41606 e_interrupt = threading.Event() signal.signal(signal.SIGINT, sigint_handler) print("[green bold]Press Ctrl+C to stop pipeline[/green bold]") else: # If there's an external interrupt, don't capture SIGINT e_interrupt = e_external_interrupt # Periodic gloop interrupt (see utils.glib_cb_restart) t_check = 100 GLib.timeout_add(t_check, glib_cb_restart, t_check) while not e_interrupt.is_set(): g_context.iteration(may_block=True) print("Ending streaming") if __name__ == "__main__": # Print any config override by env variables to avoid confusions print_config_overrides() main(config) ================================================ FILE: maskcam/mqtt_commander.py ================================================ #!/usr/bin/env python3 ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ import sys import json import time from rich import print from .mqtt_common import mqtt_send_msg, mqtt_connect_broker from .mqtt_common import MQTT_BROKER_IP, MQTT_BROKER_PORT, MQTT_DEVICE_NAME from .mqtt_common import ( MQTT_TOPIC_ALERTS, MQTT_TOPIC_FILES, MQTT_TOPIC_HELLO, MQTT_TOPIC_STATS, MQTT_TOPIC_COMMANDS, ) from .common import ( CMD_FILE_SAVE, CMD_STREAMING_START, CMD_STREAMING_STOP, CMD_INFERENCE_RESTART, ) def show_message(mqtt_client, userdata, message): print(f"Message received in topic: [yellow]{message.topic}[/yellow]") print(json.loads(message.payload.decode())) if MQTT_BROKER_IP is None or MQTT_DEVICE_NAME is None: print( "\n[red]MQTT is DISABLED[/red]" " since MQTT_BROKER_IP or MQTT_DEVICE_NAME env vars are not defined\n" ) sys.exit(0) # Subscribe to some topics print("\n[blue]Available topics:[/blue]") print(MQTT_TOPIC_ALERTS) print(MQTT_TOPIC_FILES) print(MQTT_TOPIC_HELLO) print(MQTT_TOPIC_STATS) print(MQTT_TOPIC_COMMANDS) topics_subscribe = [] while True: topic = input("\nSubscribe to topic (empty to continue): ") if topic == "": break topics_subscribe.append((topic, 2)) # Use qos=2 # Connect to client and subscribe mqtt_client = mqtt_connect_broker( client_id="commander", broker_ip=MQTT_BROKER_IP, broker_port=MQTT_BROKER_PORT, subscribe_to=topics_subscribe, ) mqtt_client.on_message = show_message time.sleep(1) # Wait to print connection messages # Send commands print("\n[blue]Available commands:[/blue]") print(CMD_FILE_SAVE) print(CMD_STREAMING_START) print(CMD_STREAMING_STOP) print(CMD_INFERENCE_RESTART) while True: cmd = input("\nSend command to device (q to exit):\n") if cmd == "q": break payload = {"device_id": MQTT_DEVICE_NAME, "command": cmd} mqtt_send_msg(mqtt_client, MQTT_TOPIC_COMMANDS, payload) ================================================ FILE: maskcam/mqtt_common.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ import os import json from multiprocessing import Queue from typing import Callable, List from paho.mqtt import client as paho_mqtt_client from .config import config from .prints import print_mqtt as print # MQTT topics MQTT_TOPIC_HELLO = "hello" MQTT_TOPIC_UPDATE = "device-status" MQTT_TOPIC_STATS = "receive-from-jetson" MQTT_TOPIC_ALERTS = "alerts" MQTT_TOPIC_FILES = "video-files" MQTT_TOPIC_COMMANDS = "commands" config_broker_ip = config["mqtt"]["mqtt-broker-ip"].strip() config_device_name = config["mqtt"]["mqtt-device-name"].strip() # Must come defined or MQTT gets disabled MQTT_BROKER_IP = None if config_broker_ip and config_broker_ip != "0": MQTT_BROKER_IP = config_broker_ip MQTT_DEVICE_NAME = None if config_device_name and config_device_name != "0": MQTT_DEVICE_NAME = config_device_name MQTT_BROKER_PORT = int(config["mqtt"]["mqtt-broker-port"]) MQTT_DEVICE_DESCRIPTION = config["mqtt"]["mqtt-device-description"] mqtt_msg_queue = Queue(maxsize=100) # 100 mqtt messages stored max def mqtt_send_queue(mqtt_client): success = True while not mqtt_msg_queue.empty() and success: q_msg = mqtt_msg_queue.get_nowait() print(f"Sending enqueued message to topic: {q_msg['topic']}") success = mqtt_send_msg(mqtt_client, q_msg["topic"], q_msg["message"]) return success def mqtt_connect_broker( client_id: str, broker_ip: str, broker_port: int, subscribe_to: List[List] = None, cb_success: Callable = None, ) -> paho_mqtt_client: def cb_on_connect(client, userdata, flags, code): if code == 0: print("[green]Connected to MQTT Broker[/green]") if subscribe_to: print("Subscribing to topics:") print(subscribe_to) client.subscribe(subscribe_to) # Always re-suscribe after reconnecting if cb_success is not None: cb_success(client) if not mqtt_send_queue(client): print(f"Failed to send MQTT message queue after connecting", warning=True) else: print(f"Failed to connect to MQTT[/red], return code {code}", warning=True) def cb_on_disconnect(client, userdata, code): print(f"Disconnected from MQTT Broker, code: {code}") client = paho_mqtt_client.Client(client_id) client.on_connect = cb_on_connect client.on_disconnect = cb_on_disconnect client.connect(broker_ip, broker_port) client.loop_start() return client def mqtt_send_msg(mqtt_client, topic, message, enqueue=True): if mqtt_client is None: print(f"MQTT not connected. Skipping message to topic: {topic}") return False # Check previous enqueued msgs mqtt_send_queue(mqtt_client) result = mqtt_client.publish(topic, json.dumps(message)) if result[0] == 0: print(f"{topic} | MQTT message [green]SENT[/green]") print(message) return True else: if enqueue: if not mqtt_msg_queue.full(): print(f"{topic} | MQTT message [yellow]ENQUEUED[/yellow]") mqtt_msg_queue.put_nowait({"topic": topic, "message": message}) else: print(f"{topic} | MQTT message [red]DROPPED: FULL QUEUE[/red]", error=True) else: print(f"{topic} | MQTT message [yellow]DISCARDED[/yellow]", warning=True) return False ================================================ FILE: maskcam/prints.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ import logging from rich.logging import RichHandler logging.basicConfig( level="NOTSET", format="%(message)s", datefmt="|", # Not needed w/balena, use [%X] otherwise handlers=[RichHandler(markup=True)], ) log = logging.getLogger("rich") def print_process( color, process_name, *args, error=False, warning=False, exception=False, **kwargs ): msg = " ".join([str(arg) for arg in args]) # Concatenate all incoming strings or objects rich_msg = f"[{color}]{process_name}[/{color}] | {msg}" if error: log.error(rich_msg) elif warning: log.warning(rich_msg) elif exception: log.exception(rich_msg) else: log.info(rich_msg) def print_run(*args, **kwargs): print_process("blue", "maskcam-run", *args, **kwargs) def print_fileserver(*args, **kwargs): print_process("dark_violet", "file-server", *args, **kwargs) def print_filesave(*args, **kwargs): print_process("dark_magenta", "file-save", *args, **kwargs) def print_streaming(*args, **kwargs): print_process("dark_green", "streaming", *args, **kwargs) def print_inference(*args, **kwargs): print_process("bright_yellow", "inference", *args, **kwargs) def print_mqtt(*args, **kwargs): print_process("bright_green", "mqtt", *args, **kwargs) def print_common(*args, **kwargs): print_process("white", "common", *args, **kwargs) ================================================ FILE: maskcam/utils.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from .config import config from gi.repository import GLib ADDRESS_UNKNOWN_LABEL = "" def get_ip_address(): result_value = config["maskcam"]["device-address"].strip() if not result_value or result_value == "0": result_value = ADDRESS_UNKNOWN_LABEL return result_value def get_streaming_address(host_address, rtsp_port, rtsp_path): return f"rtsp://{host_address}:{rtsp_port}{rtsp_path}" def format_tdelta(time_delta): # Format to show timedelta objects as string if time_delta is None: return "N/A" return f"{time_delta}".split(".")[0] # Remove nanoseconds def glib_cb_restart(t_restart): # Timer to avoid GLoop locking infinitely # We want to run g_context.iteration(may_block=True) # since may_block=False will use high CPU, # and adding sleeps lags event processing. # But we want to check periodically for other events GLib.timeout_add(t_restart, glib_cb_restart, t_restart) def load_udp_ports_filesaving(config, udp_ports_pool): for port in config["maskcam"]["udp-ports-filesave"].split(","): udp_ports_pool.add(int(port)) return udp_ports_pool ================================================ FILE: maskcam_config.txt ================================================ # NOTE: Some values might be overriden via ENV vars (check maskcam/config.py) [face-processor] # Detections with score below this threshold will be discarded detection-threshold=0.1 # Only vote mask/no_mask when detection score is above this voting-threshold=0.75 # Smaller detections (in pixels) will be discarded min-face-size=8 # Disable tracker to draw raw detections and set thresholds above disable-tracker=0 [mqtt] # These are just placeholders, to enable MQTT define these env variables: # MQTT_BROKER_IP and MQTT_DEVICE_NAME mqtt-broker-ip=0 mqtt-device-name=0 mqtt-broker-port=1883 mqtt-device-description=MaskCam @ Jetson Nano [maskcam] # Alert conditions # Minimum people to even calculate no-mask-fraction alert-min-visible-people=1 # More than this fraction of people without mask will raise an alarm alert-no-mask-fraction=0.25 # More than this people detected will raise alarm despite no-mask-fraction alert-max-total-people=10 # Time to send statistics in seconds. Set smaller than fileserver-video-period statistics-period=15 # Time (in seconds) to restart statistics (and the whole Deepstream inference process) # Set to 0 to disable / 24hs = 86400 seconds timeout-inference-restart=86400 inference-log-interval=300 # Other valid inputs: # - CSI cameras like RaspiCam: # -> argus://0 # - Any file: # -> file:///absolute/path/to/file.mp4 default-input=v4l2:///dev/video0 # Output/streaming video resolution. 1024x576 keeps 4k aspect ratio of 1.777 output-video-width=1024 output-video-height=576 # Run utils/gst_capabilities.sh and find video/x-raw entries camera-framerate=30 # Only used for argus:// inputs camera-flip-method=0 # Auto-calculate nvinfer's `interval` based on `camera-framerate` and `inference-max-fps` # to avoid delaying the pipeline. This will override the fixed `interval` parameter below # E.g: if framerate=30 and max-fps=14, # -> will set interval=2 so that inference runs only 1/3 of incoming frames inference-interval-auto=1 # Set this value to the actual FPS bottleneck of the model. Only used if inference-interval-auto. # e.g: run the model on a video file (instead of live camera source) to determine model's FPS on your device inference-max-fps=14 udp-port-streaming=5400 # 2 ports for overlapping file-save processes udp-ports-filesave=5401,5402 streaming-start-default=1 streaming-port=8554 streaming-path=/maskcam streaming-clock-rate=90000 # Supported: MP4, H264, H265 # Recommended H264 for stability on video save codec=H264 # Sequentially saving videos fileserver-enabled=1 fileserver-port=8080 fileserver-video-period=30 fileserver-video-duration=35 fileserver-force-save=0 fileserver-ram-dir=/dev/shm # Use /tmp/* to clean saved videos on system reboot fileserver-hdd-dir=/tmp/saved_videos # IP or domain address that this device will show in info messages (logs and web frontend, for streaming and file downloading) # Recommended: use env variable MASKCAM_DEVICE_ADDRESS to set this device-address=0 [property] interval=0 gpu-id=0 # Was: net-scale-factor=0.0039215697906911373 #0=RGB, 1=BGR model-color-format=0 # YOLOv4 # model-engine-file=yolo/facemask_y4tiny_1024_608_fp16.trt # model-engine-file=yolo/maskcam_y4t_1184_672_fp16.trt # model-engine-file=yolo/maskcam_y4t_1120_640_fp16.trt model-engine-file=yolo/maskcam_y4t_1024_608_fp16.trt labelfile-path=yolo/data/obj.names custom-lib-path=deepstream_plugin_yolov4/libnvdsinfer_custom_impl_Yolo.so # Detectnet_v2 # tlt-encoded-model=detectnet_v2/resnet18_detector.etlt # tlt-model-key=tlt_encode # labelfile-path=detectnet_v2/labels.txt # input-dims=3;544;960;0 # where c = number of channels, h = height of the model input, w = width of model input, 0: implies CHW format. # uff-input-blob-name=input_1 # output-blob-names=output_cov/Sigmoid;output_bbox/BiasAdd num-detected-classes=4 ## 0=FP32, 1=INT8, 2=FP16 mode network-mode=2 gie-unique-id=1 network-type=0 # is-classifier=0 ## 0=Group Rectangles, 1=DBSCAN, 2=NMS, 3= DBSCAN+NMS Hybrid, 4 = None(No clustering) # Default: 2 cluster-mode=2 # Skip inference these frames maintain-aspect-ratio=0 parse-bbox-func-name=NvDsInferParseCustomYoloV4 engine-create-func-name=NvDsInferYoloCudaEngineGet scaling-filter=1 scaling-compute-hw=1 #output-blob-names=2012 # Async mode doesn't make sense with our custom python tracker classifier-async-mode=0 [class-attrs-all] nms-iou-threshold=0.2 # Default: 0.4 pre-cluster-threshold=0.4 ================================================ FILE: maskcam_run.py ================================================ #!/usr/bin/env python3 ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ import os import sys import json import shutil import signal import threading import multiprocessing as mp # Avoids random hangs in child processes (https://pythonspeed.com/articles/python-multiprocessing/) mp.set_start_method("spawn") # noqa from rich.console import Console from datetime import datetime, timedelta from maskcam.prints import print_run as print from maskcam.config import config, print_config_overrides from maskcam.common import USBCAM_PROTOCOL, RASPICAM_PROTOCOL from maskcam.common import ( CMD_FILE_SAVE, CMD_STREAMING_START, CMD_STREAMING_STOP, CMD_INFERENCE_RESTART, CMD_FILESERVER_RESTART, CMD_STATUS_REQUEST, ) from maskcam.utils import ( get_ip_address, ADDRESS_UNKNOWN_LABEL, load_udp_ports_filesaving, get_streaming_address, format_tdelta, ) from maskcam.mqtt_common import mqtt_connect_broker, mqtt_send_msg from maskcam.mqtt_common import ( MQTT_BROKER_IP, MQTT_BROKER_PORT, MQTT_DEVICE_DESCRIPTION, MQTT_DEVICE_NAME, ) from maskcam.mqtt_common import ( MQTT_TOPIC_ALERTS, MQTT_TOPIC_FILES, MQTT_TOPIC_HELLO, MQTT_TOPIC_STATS, MQTT_TOPIC_UPDATE, MQTT_TOPIC_COMMANDS, ) from maskcam.maskcam_inference import main as inference_main from maskcam.maskcam_filesave import main as filesave_main from maskcam.maskcam_fileserver import main as fileserver_main from maskcam.maskcam_streaming import main as streaming_main udp_ports_pool = set() console = Console() # Use threading.Event instead of mp.Event() for sigint_handler, see: # https://bugs.python.org/issue41606 e_interrupt = threading.Event() q_commands = mp.Queue(maxsize=4) active_filesave_processes = [] P_INFERENCE = "inference" P_STREAMING = "streaming" P_FILESERVER = "file-server" P_FILESAVE_PREFIX = "file-save-" processes_info = {} def sigint_handler(sig, frame): print("[red]Ctrl+C pressed. Interrupting all processes...[/red]") e_interrupt.set() def start_process(name, target_function, config, **kwargs): e_interrupt_process = mp.Event() process = mp.Process( name=name, target=target_function, kwargs=dict( e_external_interrupt=e_interrupt_process, config=config, **kwargs, ), ) processes_info[name] = {"started": datetime.now(), "running": True} process.start() print(f"Process [yellow]{name}[/yellow] started with PID: {process.pid}") return process, e_interrupt_process def terminate_process(name, process, e_interrupt_process, delete_info=False): print(f"Sending interrupt to {name} process") e_interrupt_process.set() print(f"Waiting for process [yellow]{name}[/yellow] to terminate...") process.join(timeout=10) if process.is_alive(): print( f"[red]Forcing termination of process:[/red] [bold]{name}[/bold]", warning=True, ) process.terminate() if name in processes_info: if delete_info: del processes_info[name] # Sequential processes, avoid filling memory else: processes_info[name].update({"ended": datetime.now(), "running": False}) print(f"Process terminated: [yellow]{name}[/yellow]\n") def new_command(command): if q_commands.full(): print(f"Command {command} IGNORED. Queue is full.", error=True) return print(f"Received command: [yellow]{command}[/yellow]") q_commands.put_nowait(command) def mqtt_init(config): if MQTT_BROKER_IP is None or MQTT_DEVICE_NAME is None: print( "[red]MQTT is DISABLED[/red]" " since MQTT_BROKER_IP or MQTT_DEVICE_NAME env vars are not defined\n", warning=True, ) mqtt_client = None else: print(f"Connecting to MQTT server {MQTT_BROKER_IP}:{MQTT_BROKER_PORT}") print(f"Device name: [green]{MQTT_DEVICE_NAME}[/green]\n\n") mqtt_client = mqtt_connect_broker( client_id=MQTT_DEVICE_NAME, broker_ip=MQTT_BROKER_IP, broker_port=MQTT_BROKER_PORT, subscribe_to=[(MQTT_TOPIC_COMMANDS, 2)], # handles re-subscription cb_success=mqtt_on_connect, ) mqtt_client.on_message = mqtt_process_message return mqtt_client def mqtt_on_connect(mqtt_client): mqtt_say_hello(mqtt_client) mqtt_send_file_list(mqtt_client) def mqtt_process_message(mqtt_client, userdata, message): topic = message.topic if topic == MQTT_TOPIC_COMMANDS: payload = json.loads(message.payload.decode()) if payload["device_id"] != MQTT_DEVICE_NAME: return command = payload["command"] new_command(command) def mqtt_say_hello(mqtt_client): return mqtt_send_msg( mqtt_client, MQTT_TOPIC_HELLO, {"device_id": MQTT_DEVICE_NAME, "description": MQTT_DEVICE_DESCRIPTION}, enqueue=False, # Will be resent on_connect ) def mqtt_send_device_status(mqtt_client): t_now = datetime.now() device_address = get_ip_address() is_valid_address = device_address != ADDRESS_UNKNOWN_LABEL if P_INFERENCE in processes_info and processes_info[P_INFERENCE]["running"]: inference_runtime = t_now - processes_info[P_INFERENCE]["started"] else: inference_runtime = None if P_FILESERVER in processes_info and processes_info[P_FILESERVER]["running"]: fileserver_runtime = t_now - processes_info[P_FILESERVER]["started"] else: fileserver_runtime = None if P_STREAMING in processes_info and processes_info[P_STREAMING]["running"]: streaming_address = get_streaming_address( device_address, config["maskcam"]["streaming-port"], config["maskcam"]["streaming-path"], ) else: streaming_address = "N/A" total_fsave = len(active_filesave_processes) keep_n = len([p for p in active_filesave_processes if p["flag_keep_file"]]) return mqtt_send_msg( mqtt_client, MQTT_TOPIC_UPDATE, { "device_id": MQTT_DEVICE_NAME, "inference_runtime": format_tdelta(inference_runtime), "fileserver_runtime": format_tdelta(fileserver_runtime), "streaming_address": streaming_address, "device_address": device_address if is_valid_address else None, "save_current_files": f"{keep_n}/{total_fsave}", "time": f"{t_now:%H:%M:%S}", }, enqueue=False, # Only latest status is interesting ) def mqtt_send_file_list(mqtt_client): server_address = get_ip_address() server_port = int(config["maskcam"]["fileserver-port"]) try: file_list = sorted(os.listdir(config["maskcam"]["fileserver-hdd-dir"])) except FileNotFoundError: # directory not created file_list = [] return mqtt_send_msg( mqtt_client, MQTT_TOPIC_FILES, { "device_id": MQTT_DEVICE_NAME, "file_server": f"http://{server_address}:{server_port}", "file_list": file_list, }, enqueue=False, # Will be resent on_connect or when something changes ) def is_alert_condition(statistics, config): # Thresholds config max_total_people = int(config["maskcam"]["alert-max-total-people"]) min_visible_people = int(config["maskcam"]["alert-min-visible-people"]) max_no_mask = float(config["maskcam"]["alert-no-mask-fraction"]) # Calculate visible people without_mask = int(statistics["people_without_mask"]) with_mask = int(statistics["people_with_mask"]) visible_people = with_mask + without_mask is_alert = False if statistics["people_total"] > max_total_people: is_alert = True elif visible_people >= min_visible_people: no_mask_fraction = float(statistics["people_without_mask"]) / visible_people is_alert = no_mask_fraction > max_no_mask print(f"[yellow]ALERT condition: {is_alert}[/yellow]") return is_alert def handle_statistics(mqtt_client, stats_queue, config, is_live_input): while not stats_queue.empty(): statistics = stats_queue.get_nowait() if is_live_input: # Alert conditions detection raise_alert = is_alert_condition(statistics, config) if raise_alert: flag_keep_current_files() if mqtt_client is not None: topic = MQTT_TOPIC_ALERTS if raise_alert else MQTT_TOPIC_STATS message = {"device_id": MQTT_DEVICE_NAME, **statistics} mqtt_send_msg(mqtt_client, topic, message, enqueue=True) def allocate_free_udp_port(): new_port = udp_ports_pool.pop() print(f"Allocating UDP port: {new_port}") return new_port def release_udp_port(port_number): print(f"Releasing UDP port: {port_number}") udp_ports_pool.add(port_number) def handle_file_saving( video_period, video_duration, ram_dir, hdd_dir, force_save, mqtt_client=None ): period = timedelta(seconds=video_period) duration = timedelta(seconds=video_duration) latest_start = None latest_number = 0 # Handle termination of previous file-saving processes and move files RAM->HDD terminated_idxs = [] for idx, active_process in enumerate(active_filesave_processes): if datetime.now() - active_process["started"] >= duration: finish_filesave_process(active_process, hdd_dir, force_save, mqtt_client=mqtt_client) terminated_idxs.append(idx) if latest_start is None or active_process["started"] > latest_start: latest_start = active_process["started"] latest_number = active_process["number"] # Remove terminated processes from list in a separated loop for idx in sorted(terminated_idxs, reverse=True): del active_filesave_processes[idx] # Start new file-saving process if time has elapsed if latest_start is None or (datetime.now() - latest_start >= period): print( "[green]Time to start a new video file [/green]" f" (latest started at: {format_tdelta(latest_start)})" ) new_process_number = latest_number + 1 new_process_name = f"{P_FILESAVE_PREFIX}{new_process_number}" new_filename = f"{datetime.today().strftime('%Y%m%d_%H%M%S')}_{new_process_number}.mp4" new_filepath = f"{ram_dir}/{new_filename}" new_udp_port = allocate_free_udp_port() process_handler, e_interrupt_process = start_process( new_process_name, filesave_main, config, output_filename=new_filepath, udp_port=new_udp_port, ) active_filesave_processes.append( dict( number=new_process_number, name=new_process_name, filepath=new_filepath, filename=new_filename, started=datetime.now(), process_handler=process_handler, e_interrupt=e_interrupt_process, flag_keep_file=False, udp_port=new_udp_port, ) ) def finish_filesave_process(active_process, hdd_dir, force_filesave, mqtt_client=None): terminate_process( active_process["name"], active_process["process_handler"], active_process["e_interrupt"], delete_info=True, ) release_udp_port(active_process["udp_port"]) # Move file to its definitive place if flagged, otherwise remove it if active_process["flag_keep_file"] or force_filesave: definitive_filepath = f"{hdd_dir}/{active_process['filename']}" print(f"Force file saving: {bool(force_filesave)}") print(f"Permanent video file created: [green]{definitive_filepath}[/green]") # Must use shutil here to move RAM->HDD shutil.move(active_process["filepath"], definitive_filepath) # Send updated file list via MQTT (prints ignore if mqtt_client is None) mqtt_send_file_list(mqtt_client) else: print(f"Removing RAM video file: {active_process['filepath']}") os.remove(active_process["filepath"]) def flag_keep_current_files(): print("Request to [green]save current video files[/green]") for process in active_filesave_processes: print(f"Set flag to keep: [green]{process['filename']}[/green]") process["flag_keep_file"] = True if __name__ == "__main__": if len(sys.argv) > 2: print( """Usage: python3 maskcam_run.py [ URI ] Examples: \t$ python3 maskcam_run.py \t$ python3 maskcam_run.py file:///absolute/path/to/file.mp4 \t$ python3 maskcam_run.py v4l2:///dev/video1 \t$ python3 maskcam_run.py argus://0 Notes: \t - If no URI is provided, will use default-input defined in config_maskcam.txt \t - If a file:///path/file.mp4 is provided, the output will be ./output_file.mp4 \t - If the input is a live camera, the output will be consecutive \t video files under /dev/shm/date_time.mp4 \t according to the time interval defined in output-chunks-duration in config_maskcam.txt. """ ) sys.exit(0) try: # Print any ENV var config override to avoid confusions print_config_overrides() # Input source if len(sys.argv) > 1: input_filename = sys.argv[1] print(f"Provided input source: {input_filename}") else: input_filename = config["maskcam"]["default-input"] print(f"Using input from config file: {input_filename}") # Input type: file or live camera is_usbcamera = USBCAM_PROTOCOL in input_filename is_raspicamera = RASPICAM_PROTOCOL in input_filename is_live_input = is_usbcamera or is_raspicamera # Streaming enabled by default? streaming_autostart = int(config["maskcam"]["streaming-start-default"]) # Fileserver: sequentially save videos (only for camera input) fileserver_enabled = is_live_input and int(config["maskcam"]["fileserver-enabled"]) fileserver_period = int(config["maskcam"]["fileserver-video-period"]) fileserver_duration = int(config["maskcam"]["fileserver-video-duration"]) fileserver_force_save = int(config["maskcam"]["fileserver-force-save"]) fileserver_ram_dir = config["maskcam"]["fileserver-ram-dir"] fileserver_hdd_dir = config["maskcam"]["fileserver-hdd-dir"] # Inference restart timeout tout_inference_restart = int(config["maskcam"]["timeout-inference-restart"]) if is_live_input and tout_inference_restart: tout_inference_restart = timedelta(seconds=tout_inference_restart) else: tout_inference_restart = 0 # Filesave processes: load available ports load_udp_ports_filesaving(config, udp_ports_pool) # Should only have 1 element at a time unless this thread gets blocked stats_queue = mp.Queue(maxsize=5) # Init MQTT or set these to None if is_live_input: mqtt_client = mqtt_init(config) else: mqtt_client = None # SIGINT handler (Ctrl+C) signal.signal(signal.SIGINT, sigint_handler) print("[green bold]Press Ctrl+C to stop all processes[/green bold]") process_inference = None process_streaming = None process_fileserver = None e_inference_ready = mp.Event() if fileserver_enabled: process_fileserver, e_interrupt_fileserver = start_process( P_FILESERVER, fileserver_main, config, directory=fileserver_hdd_dir ) if streaming_autostart: print("[yellow]Starting streaming (streaming-start-default is set)[/yellow]") new_command(CMD_STREAMING_START) # Inference process: If input is a file, also saves file output_filename = None if is_live_input else f"output_{input_filename.split('/')[-1]}" process_inference, e_interrupt_inference = start_process( P_INFERENCE, inference_main, config, input_filename=input_filename, output_filename=output_filename, stats_queue=stats_queue, e_ready=e_inference_ready, ) while not e_interrupt.is_set(): # Send MQTT statistics, detect alarm events and request file-saving handle_statistics(mqtt_client, stats_queue, config, is_live_input) # Handle sequential file saving processes, only after inference process is ready if e_inference_ready.is_set(): if fileserver_enabled and is_live_input: # server can be enabled via MQTT handle_file_saving( fileserver_period, fileserver_duration, fileserver_ram_dir, fileserver_hdd_dir, fileserver_force_save, mqtt_client=mqtt_client, ) if not q_commands.empty(): command = q_commands.get_nowait() reply_updated_status = False print(f"Processing command: [yellow]{command}[yellow]") if command == CMD_STREAMING_START: if process_streaming is None or not process_streaming.is_alive(): process_streaming, e_interrupt_streaming = start_process( P_STREAMING, streaming_main, config ) reply_updated_status = True elif command == CMD_STREAMING_STOP: if process_streaming is not None and process_streaming.is_alive(): terminate_process(P_STREAMING, process_streaming, e_interrupt_streaming) reply_updated_status = True elif command == CMD_INFERENCE_RESTART: if process_inference.is_alive(): terminate_process(P_INFERENCE, process_inference, e_interrupt_inference) process_inference, e_interrupt_inference = start_process( P_INFERENCE, inference_main, config, input_filename=input_filename, output_filename=output_filename, stats_queue=stats_queue, ) reply_updated_status = True elif command == CMD_FILESERVER_RESTART: if process_fileserver is not None and process_fileserver.is_alive(): terminate_process(P_FILESERVER, process_fileserver, e_interrupt_fileserver) process_fileserver, e_interrupt_fileserver = start_process( P_FILESERVER, fileserver_main, config, directory=fileserver_hdd_dir, ) fileserver_enabled = True reply_updated_status = True elif command == CMD_FILE_SAVE: flag_keep_current_files() reply_updated_status = True elif command == CMD_STATUS_REQUEST: reply_updated_status = True else: print("[red]Command not recognized[/red]", error=True) if reply_updated_status: mqtt_send_device_status(mqtt_client) else: e_interrupt.wait(timeout=0.1) # Routine check: finish loop if the inference process is dead if not process_inference.is_alive(): e_interrupt.set() # Routine check: restart inference at given interval (only live_input) if tout_inference_restart: inference_runtime = datetime.now() - processes_info[P_INFERENCE]["started"] if inference_runtime > tout_inference_restart: print( "[yellow]Restarting inference due to timeout-inference-restart" f"(inference runtime: {format_tdelta(inference_runtime)})[/yellow]" ) new_command(CMD_INFERENCE_RESTART) except: # noqa console.print_exception() # Terminate all running processes, avoid breaking on any exception for active_file_process in active_filesave_processes: try: finish_filesave_process( active_file_process, fileserver_hdd_dir, fileserver_force_save, mqtt_client=mqtt_client, ) except: # noqa console.print_exception() try: if process_inference is not None and process_inference.is_alive(): terminate_process(P_INFERENCE, process_inference, e_interrupt_inference) except: # noqa console.print_exception() try: if process_fileserver is not None and process_fileserver.is_alive(): terminate_process(P_FILESERVER, process_fileserver, e_interrupt_fileserver) except: # noqa console.print_exception() try: if process_streaming is not None and process_streaming.is_alive(): terminate_process(P_STREAMING, process_streaming, e_interrupt_streaming) except: # noqa console.print_exception() ================================================ FILE: requirements-dev.in ================================================ # dev deps pip-tools black flake8 jupyter ipython ipdb ================================================ FILE: requirements.in ================================================ # General deps numpy scipy PyYAML ipdb # MQTT paho-mqtt # Tracker norfair # Colored prints and traceback rich ================================================ FILE: requirements.txt ================================================ # Versions frozen after `pip install -r requirements.in -c constraints.docker` # under balenalib/jetson-nano-ubuntu:bionic asn1crypto==0.24.0 attrs==17.4.0 backcall==0.2.0 colorama==0.4.4 commonmark==0.9.1 cryptography==2.1.4 cycler==0.10.0 Cython==0.26.1 dataclasses==0.7 decorator==4.1.2 filterpy==1.4.5 idna==2.6 ipdb==0.13.4 ipython==7.16.1 ipython-genutils==0.2.0 jedi==0.18.0 Jetson.GPIO==2.0.8 joblib==0.11 keyring==10.6.0 keyrings.alt==3.0 matplotlib==2.1.1 norfair==0.1.8 nose==1.3.7 numpy==1.13.3 olefile==0.45.1 opencv-python==3.2.0 paho-mqtt==1.5.1 parso==0.8.1 pexpect==4.8.0 pickleshare==0.7.5 Pillow==5.1.0 pluggy==0.6.0 prompt-toolkit==3.0.15 ptyprocess==0.7.0 py==1.5.2 pycrypto==2.6.1 pyds==1.0.1 Pygments==2.7.4 pygobject==3.26.1 pyparsing==2.2.0 pytest==3.3.2 python-dateutil==2.6.1 pytz==2018.3 pyxdg==0.25 PyYAML==5.4.1 rich==6.2.0 scikit==0.19.1 scikit-learn==0.19.1 scipy==0.19.1 SecretStorage==2.3.1 simplejson==3.13.2 six==1.11.0 tensorrt==7.1.3.0 traitlets==4.3.3 typing-extensions==3.7.4.3 wcwidth==0.2.5 ================================================ FILE: server/backend/Dockerfile ================================================ FROM tiangolo/uvicorn-gunicorn-fastapi:python3.7 COPY ./app /app COPY requirements.txt /app/requirements.txt ENV PYTHONPATH=/app WORKDIR /app RUN python -m pip install --upgrade pip && pip install -r requirements.txt ================================================ FILE: server/backend/alembic.ini ================================================ # A generic, single database configuration. [alembic] # path to migration scripts script_location = app/db/migrations # template used to generate migration files # file_template = %%(rev)s_%%(slug)s # timezone to use when rendering the date # within the migration file as well as the filename. # string value is passed to dateutil.tz.gettz() # leave blank for localtime # timezone = # max length of characters to apply to the # "slug" field # truncate_slug_length = 40 # set to 'true' to run the environment during # the 'revision' command, regardless of autogenerate # revision_environment = false # set to 'true' to allow .pyc and .pyo files without # a source .py file to be detected as revisions in the # versions/ directory # sourceless = false # version location specification; this defaults # to alembic/versions. When using multiple version # directories, initial revisions must be specified with --version-path # version_locations = %(here)s/bar %(here)s/bat alembic/versions # the output encoding used when revision files # are written from script.py.mako # output_encoding = utf-8 [post_write_hooks] # post_write_hooks defines scripts or Python functions that are run # on newly generated revision scripts. See the documentation for further # detail and examples # format using "black" - use the console_scripts runner, against the "black" entrypoint # hooks=black # black.type=console_scripts # black.entrypoint=black # black.options=-l 79 # Logging configuration [loggers] keys = root,sqlalchemy,alembic [handlers] keys = console [formatters] keys = generic [logger_root] level = WARN handlers = console qualname = [logger_sqlalchemy] level = WARN handlers = qualname = sqlalchemy.engine [logger_alembic] level = INFO handlers = qualname = alembic [handler_console] class = StreamHandler args = (sys.stderr,) level = NOTSET formatter = generic [formatter_generic] format = %(levelname)-5.5s [%(name)s] %(message)s datefmt = %H:%M:%S ================================================ FILE: server/backend/app/api/__init__.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from .exceptions import NoItemFoundException, GenericException, ItemAlreadyExist from .routes.device_routes import device_router from .routes.statistic_routes import statistic_router ================================================ FILE: server/backend/app/api/exceptions.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from fastapi import HTTPException class NoItemFoundException(HTTPException): def __init__(self): super().__init__( status_code=404, detail="No item was found for the provided ID", ) class ItemAlreadyExist(HTTPException): def __init__(self): super().__init__( status_code=500, detail="An instance with the same id already exist", ) class GenericException(HTTPException): def __init__(self, message: str): super().__init__( status_code=500, detail=f"An error occurred: \n{message}", ) ================================================ FILE: server/backend/app/api/routes/device_routes.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from typing import Dict, List, Optional from app.api import GenericException, ItemAlreadyExist, NoItemFoundException from app.db.cruds import ( create_device, delete_device, get_device, get_devices, update_device, get_files_by_device ) from app.db.schema import DeviceSchema, VideoFileSchema, get_db_generator from fastapi import APIRouter, Depends from fastapi.encoders import jsonable_encoder from sqlalchemy.exc import DataError, IntegrityError from sqlalchemy.orm import Session from sqlalchemy.orm.exc import NoResultFound device_router = APIRouter() @device_router.post("/devices", response_model=DeviceSchema) def create_device_item( device_information: DeviceSchema, db: Session = Depends(get_db_generator), ): """ Create device. Arguments: device_information {DeviceSchema} -- New device information. db {Session} -- Database session. Returns: Union[DeviceSchema, ItemAlreadyExist] -- Device instance that was added to the database or an error in case the device already exists. """ try: device_information = jsonable_encoder(device_information) return create_device( db_session=db, device_information=device_information ) except IntegrityError: raise ItemAlreadyExist() @device_router.get("/devices/{device_id}", response_model=DeviceSchema) def get_device_item( device_id: str, db: Session = Depends(get_db_generator), ): """ Get existing device. Arguments: device_id {str} -- Device id. db {Session} -- Database session. Returns: Union[DeviceSchema, NoItemFoundException] -- Device instance which id is device_id or an exception in case there's no matching device. """ try: return get_device(db_session=db, device_id=device_id) except NoResultFound: raise NoItemFoundException() @device_router.get( "/devices", response_model=List[DeviceSchema], response_model_include={"id", "description", "file_server_address"}, ) def get_devices_items(db: Session = Depends(get_db_generator)): """ Get all existing devices. Arguments: db {Session} -- Database session. Returns: List[DeviceSchema] -- All device instances present in the database. """ return get_devices(db_session=db) @device_router.put("/devices/{device_id}", response_model=DeviceSchema) def update_device_item( device_id: str, new_device_information: Dict = {}, db: Session = Depends(get_db_generator), ): """ Modify a device. Arguments: device_id {str} -- Device id. new_device_information {Dict} -- New device information. db {Session} -- Database session. Returns: Union[DeviceSchema, NoItemFoundException, GenericException] -- Device instance which id is device_id or an exception in case there's no matching device. """ try: return update_device( db_session=db, device_id=device_id, new_device_information=new_device_information, ) except NoResultFound: raise NoItemFoundException() except DataError as e: raise GenericException(e) @device_router.delete("/devices/{device_id}", response_model=DeviceSchema) def delete_device_item( device_id: str, db: Session = Depends(get_db_generator), ): """ Delete a device. Arguments: device_id {str} -- Device id. db {Session} -- Database session. Returns: Union[DeviceSchema, NoItemFoundException, GenericException] -- Device instance that was deleted or an exception in case there's no matching device. """ try: return delete_device(db_session=db, device_id=device_id) except NoResultFound: raise NoItemFoundException() except DataError as e: raise GenericException(e) @device_router.get("/files/{device_id}", response_model=List[VideoFileSchema]) def get_device_files( device_id: str, db: Session = Depends(get_db_generator), ): """ Get existing video files in device. Arguments: device_id {str} -- Device id. db {Session} -- Database session. Returns: List[VideoFileSchema] -- VideoFile instances which device_id matches """ return get_files_by_device(db_session=db, device_id=device_id) ================================================ FILE: server/backend/app/api/routes/statistic_routes.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from typing import Dict, List, Optional from app.api import GenericException, ItemAlreadyExist, NoItemFoundException from app.db.cruds import ( create_statistic, delete_statistic, get_statistic, get_statistics, get_statistics_from_to, update_statistic, ) from app.db.schema import StatisticSchema, get_db_generator from app.db.utils import convert_timestamp_to_datetime, get_enum_type from fastapi import APIRouter, Depends, Query from sqlalchemy.exc import DataError, IntegrityError from sqlalchemy.orm import Session from sqlalchemy.orm.exc import NoResultFound statistic_router = APIRouter() @statistic_router.post( "/devices/{device_id}/statistics", response_model=StatisticSchema ) def create_statistic_item( device_id: str, statistic_information: Dict = {}, db: Session = Depends(get_db_generator), ): """ Create new statistic entry. Arguments: device_id {str} -- Device id which sent the statistic. statistic_information {Dict} -- New statistic information. db {Session} -- Database session. Returns: Union[StatisticSchema, ItemAlreadyExist] -- Statistic instance that was added to the database or an exception in case a statistic already exists. """ try: # Format input data statistic_information["device_id"] = device_id statistic_information["datetime"] = convert_timestamp_to_datetime( statistic_information["datetime"] ) statistic_information["statistic_type"] = get_enum_type( statistic_information["statistic_type"] ) return create_statistic( db_session=db, statistic_information=statistic_information ) except IntegrityError: raise ItemAlreadyExist() @statistic_router.get( "/devices/{device_id}/statistics/{timestamp}", response_model=StatisticSchema, ) def get_statistic_item( device_id: str, timestamp: float, db: Session = Depends(get_db_generator), ): """ Get a specific statistic. Arguments: device_id {str} -- Device id. timestamp {float} -- Timestamp when the device registered the information. db {Session} -- Database session. Returns: Union[StatisticSchema, NoItemFoundException] -- Statistic instance defined by device_id and timestamp or an exception in case there's no matching statistic. """ try: return get_statistic( db_session=db, device_id=device_id, datetime=convert_timestamp_to_datetime(timestamp), ) except NoResultFound: raise NoItemFoundException() @statistic_router.get( "/devices/{device_id}/statistics", response_model=List[StatisticSchema], ) def get_all_device_statistics_items( device_id: str, datefrom: Optional[str] = Query(None), dateto: Optional[str] = Query(None), timestampfrom: Optional[float] = Query(None), timestampto: Optional[float] = Query(None), db: Session = Depends(get_db_generator), ): """ Get all statistics of a specific device. Arguments: device_id {str} -- Device id. datefrom {Optional[str]} -- Datetime to show information from. dateto {Optional[str]} -- Datetime to show information to. timestampfrom {Optional[float]} -- Timestamp to show information from. timestampto {Optional[float]} -- Timestamp to show information from. db {Session} -- Database session. Returns: List[StatisticSchema] -- Statistic instances defined by device_id and datetime range. """ from_datetime = datefrom to_datetime = dateto if not from_datetime and timestampfrom: from_datetime = convert_timestamp_to_datetime(timestampfrom) if not to_datetime and timestampto: to_datetime = convert_timestamp_to_datetime(timestampto) return get_statistics_from_to( db_session=db, device_id=device_id, from_date=from_datetime, to_date=to_datetime, ) @statistic_router.get( "/statistics", response_model=List[StatisticSchema], ) def get_all_statistics_items(db: Session = Depends(get_db_generator)): """ Get all statistics from all devices. Arguments: db {Session} -- Database session. Returns: List[StatisticSchema] -- All statistic instances present in the database. """ return get_statistics(db_session=db) @statistic_router.put( "/devices/{device_id}/statistics/{timestamp}", response_model=StatisticSchema, ) def update_statistic_item( device_id: str, timestamp: float, new_statistic_information: Dict = {}, db: Session = Depends(get_db_generator), ): """ Modify a specific statistic. Arguments: device_id {str} -- Device id. timestamp {float} -- Timestamp when the device registered the information. new_statistic_information {Dict} -- New statistic information. db {Session} -- Database session. Returns: Union[StatisticSchema, NoItemFoundException, GenericException] -- Updated statistic instance defined by device_id and datetime or an exception in case there's no matching statistic. """ try: return update_statistic( db_session=db, device_id=device_id, datetime=convert_timestamp_to_datetime(timestamp), new_statistic_information=new_statistic_information, ) except NoResultFound: raise NoItemFoundException() except DataError as e: raise GenericException(e) @statistic_router.delete( "/devices/{device_id}/statistics/{timestamp}", response_model=StatisticSchema, ) def delete_statistic_item( device_id: str, timestamp: float, db: Session = Depends(get_db_generator), ): """ Delete a specific statistic. Arguments: device_id {str} -- Device id. timestamp {float} -- Timestamp when the device registered the information. db {Session} -- Database session. Returns: Union[StatisticSchema, NoItemFoundException, GenericException] -- Statistic instance that was deleted or an exception in case there's no matching statistic. """ try: return delete_statistic( db_session=db, device_id=device_id, datetime=convert_timestamp_to_datetime(timestamp), ) except NoResultFound: raise NoItemFoundException() except DataError as e: raise GenericException(e) ================================================ FILE: server/backend/app/core/config.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ import os # Database configuration DB_USER = os.environ["POSTGRES_USER"] DB_PASSWORD = os.environ["POSTGRES_PASSWORD"] DB_NAME = os.environ["POSTGRES_DB"] DB_PORT = os.environ["POSTGRES_PORT"] DB_URI = f"postgresql://{DB_USER}:{DB_PASSWORD}@db:{DB_PORT}/{DB_NAME}" # MQTT broker configuration MQTT_BROKER = os.environ["MQTT_BROKER"] MQTT_BROKER_PORT = int(os.environ["MQTT_BROKER_PORT"]) # MQTT subscriber configuration SUBSCRIBER_CLIENT_ID = os.environ["SUBSCRIBER_CLIENT_ID"] # Topic configuration MQTT_HELLO_TOPIC = "hello" MQTT_ALERT_TOPIC = "alerts" MQTT_REPORT_TOPIC = "receive-from-jetson" MQTT_SEND_TOPIC = "send-to-jetson" MQTT_FILES_TOPIC = "video-files" ================================================ FILE: server/backend/app/db/cruds/__init__.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from .crud_device import ( create_device, delete_device, get_device, get_devices, update_device, ) from .crud_statistic import ( create_statistic, delete_statistic, get_statistic, get_statistics, get_statistics_from_to, update_statistic, ) from .crud_video_file import ( update_files, get_files_by_device, ) ================================================ FILE: server/backend/app/db/cruds/crud_device.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from typing import List, Union, Dict from app.db.schema import DeviceModel from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from sqlalchemy.orm.exc import NoResultFound def create_device( db_session: Session, device_information: Dict = {} ) -> Union[DeviceModel, IntegrityError]: """ Register new Jetson device. Arguments: db_session {Session} -- Database session. device_information {Dict} -- New device information. Returns: Union[DeviceModel, IntegrityError] -- Device instance that was added to the database or an exception in case the device already exists. """ try: # Replace empty spaces in device id device_information["id"] = device_information["id"].replace(" ", "_") # Create device device = DeviceModel(**device_information) db_session.add(device) db_session.commit() db_session.refresh(device) return device except IntegrityError: db_session.rollback() raise def get_device( db_session: Session, device_id: str ) -> Union[DeviceModel, NoResultFound]: """ Get a specific device. Arguments: db_session {Session} -- Database session. device_id {str} -- Jetson id. Returns: Union[DeviceModel, NoResultFound] -- Device instance which id is device_id or an exception in case there's no matching device. """ try: return get_device_by_id(db_session, device_id) except NoResultFound: raise def get_devices(db_session: Session) -> List[DeviceModel]: """ Get all devices. Arguments: db_session {Session} -- Database session. Returns: List[DeviceModel] -- All device instances present in the database. """ return db_session.query(DeviceModel).all() def update_device( db_session: Session, device_id: str, new_device_information: Dict = {} ) -> Union[DeviceModel, NoResultFound]: """ Modify a specific Jetson device. Arguments: db_session {Session} -- Database session. device_id {str} -- Jetson id. new_device_information {Dict} -- New device information. Returns: Union[DeviceModel, NoResultFound] -- Device instance which id is device_id or an exception in case there's no matching device. """ try: try: # Remove device id as it can't be modified del new_device_information["id"] except KeyError: pass device = get_device_by_id(db_session, device_id) for key, value in new_device_information.items(): if hasattr(device, key): setattr(device, key, value) db_session.commit() return device except NoResultFound: raise def delete_device( db_session: Session, device_id: str ) -> Union[DeviceModel, NoResultFound]: """ Delete a device. Arguments: db_session {Session} -- Database session. device_id {str} -- Jetson id. Returns: Union[DeviceModel, NoResultFound] -- Device instance that was deleted or an exception in case there's no matching device. """ try: device = get_device_by_id(db_session, device_id) db_session.delete(device) db_session.commit() return device except NoResultFound: raise def get_device_by_id( db_session: Session, device_id: str ) -> Union[DeviceModel, NoResultFound]: """ Get a device using the table's primary key. Arguments: db_session {Session} -- Database session. device_id {str} -- Jetson id. Returns: Union[DeviceModel, NoResultFound] -- Device instance which id is device_id or an exception in case there's no matching device. """ device = db_session.query(DeviceModel).get(device_id) if not device: raise NoResultFound() return device ================================================ FILE: server/backend/app/db/cruds/crud_statistic.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from datetime import datetime, timezone from typing import Dict, List, Optional, Union from app.db.schema import StatisticsModel from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from sqlalchemy.orm.exc import NoResultFound def create_statistic( db_session: Session, statistic_information: Dict = {} ) -> Union[StatisticsModel, IntegrityError]: """ Register new statistic entry. Arguments: db_session {Session} -- Database session. statistic_information {Dict} -- New statistic information. Returns: Union[StatisticsModel, IntegrityError] -- Statistic instance that was added to the database or an exception in case a statistic already exists. """ try: statistic = StatisticsModel(**statistic_information) db_session.add(statistic) db_session.commit() db_session.refresh(statistic) return statistic except IntegrityError: db_session.rollback() raise def get_statistic( db_session: Session, device_id: str, datetime: datetime ) -> Union[StatisticsModel, NoResultFound]: """ Get a specific statistic. Arguments: db_session {Session} -- Database session. device_id {str} -- Jetson id which sent the information. datetime {datetime} -- Datetime when the device registered the information. Returns: Union[StatisticsModel, NoResultFound] -- Statistic instance defined by device_id and datetime or an exception in case there's no matching statistic. """ try: return get_statistic_by_id_and_datetime(db_session, device_id, datetime) except NoResultFound: raise def get_statistics( db_session: Session, device_id: Optional[str] = None ) -> List[StatisticsModel]: """ Get all statistics. Arguments: db_session {Session} -- Database session. device_id {Optional[str]} -- Device id. Returns: List[StatisticsModel] -- All statistic instances present in the database or all statistics from a specific device. """ if device_id: # Get all statistics from a specific device query = db_session.query(StatisticsModel) return query.filter(StatisticsModel.device_id == device_id).all() else: # Get all statistics from form the database return db_session.query(StatisticsModel).all() def get_statistics_from_to( db_session: Session, device_id: str, from_date: Optional[str] = None, to_date: Optional[str] = None, ) -> List[StatisticsModel]: """ Get all statistics within a datetime range. Arguments: db_session {Session} -- Database session. device_id {str} -- Device id. from_date {Optional[str]} -- Beginning of datetime range. to_date {Optional[str]} -- End of datetime range. Returns: List[StatisticsModel] -- All statistic instances present in the database within a given datetime range. """ query = db_session.query(StatisticsModel) query = query.filter(StatisticsModel.device_id == device_id) if to_date is None: # By default, show information until the current day to_date = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S") if from_date: return query.filter( StatisticsModel.datetime.between(from_date, to_date) ).all() return query.filter(StatisticsModel.datetime <= to_date).all() def update_statistic( db_session: Session, device_id: str, datetime: datetime, new_statistic_information: Dict = {}, ) -> Union[StatisticsModel, NoResultFound]: """ Modify a specific statistic. Arguments: db_session {Session} -- Database session. device_id {str} -- Jetson id which sent the information. datetime {datetime} -- Datetime when the device registered the information. new_statistic_information {Dict} -- New statistic information. Returns: Union[StatisticsModel, NoResultFound] -- Updated statistic instance defined by device_id and datetime or an exception in case there's no matching statistic. """ try: try: # Remove device id as it can't be modified del new_statistic_information["device_id"] except KeyError: pass try: # Remove datetime as it can't be modified del new_statistic_information["datetime"] except KeyError: pass statistic = get_statistic_by_id_and_datetime( db_session, device_id, datetime ) for key, value in new_statistic_information.items(): if hasattr(statistic, key): setattr(statistic, key, value) db_session.commit() return statistic except NoResultFound: raise def delete_statistic( db_session: Session, device_id: str, datetime: datetime ) -> Union[StatisticsModel, NoResultFound]: """ Delete a specific statistic. Arguments: db_session {Session} -- Database session. device_id {str} -- Jetson id which sent the information. datetime {datetime} -- Datetime when the device registered the information. Returns: Union[StatisticsModel, NoResultFound] -- Statistic instance that was deleted or an exception in case there's no matching statistic. """ try: statistic = get_statistic_by_id_and_datetime( db_session, device_id, datetime ) db_session.delete(statistic) db_session.commit() return statistic except NoResultFound: raise def get_statistic_by_id_and_datetime( db_session: Session, device_id: str, datetime: datetime ) -> Union[StatisticsModel, NoResultFound]: """ Get a statistic using the table's primary keys. Arguments: db_session {Session} -- Database session. device_id {str} -- Jetson id which sent the information. datetime {datetime} -- Datetime when the device registered the information. Returns: Union[StatisticsModel, NoResultFound] -- Statistic instance defined by device_id and datetime or an exception in case there's no matching statistic. """ statistic = db_session.query(StatisticsModel).get((device_id, datetime)) if not statistic: raise NoResultFound() return statistic ================================================ FILE: server/backend/app/db/cruds/crud_video_file.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from datetime import datetime, timezone from typing import Dict, List, Optional, Union from app.db.schema import VideoFilesModel from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from sqlalchemy.orm.exc import NoResultFound def update_files( db_session: Session, device_id: str, file_list: List, ) -> List[VideoFilesModel]: """ Update the whole list of available files for this device Arguments: db_session {Session} -- Database session. device_id {str} -- Jetson id which sent the information. file_list {List} -- List of all available files in the device Returns: VideoFilesModel -- Updated video_files instance defined by device_id """ # Remove all previous files for device_id query = db_session.query(VideoFilesModel) query = query.filter(VideoFilesModel.device_id == device_id) query.delete(synchronize_session=False) db_session.commit() # Add new list result = [] for new_file in file_list: file_add = VideoFilesModel(device_id=device_id, video_name=new_file) db_session.add(file_add) result.append(file_add) db_session.commit() return result def get_files_by_device( db_session: Session, device_id: str ) -> List[VideoFilesModel]: """ Get a file using the table's primary keys. Arguments: db_session {Session} -- Database session. device_id {str} -- Jetson id to query files Returns: List[VideoFilesModel] -- All video files for the device """ query = db_session.query(VideoFilesModel) return query.filter(VideoFilesModel.device_id == device_id).all() ================================================ FILE: server/backend/app/db/migrations/env.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from logging.config import fileConfig from alembic import context from app.core import config as app_config from app.db.schema import Base from sqlalchemy import engine_from_config, pool # this is the Alembic Config object, which provides # access to the values within the .ini file in use. config = context.config config.set_main_option("sqlalchemy.url", app_config.DB_URI) # Interpret the config file for Python logging. # This line sets up loggers basically. fileConfig(config.config_file_name) target_metadata = Base.metadata def run_migrations_offline(): """ Run migrations in 'offline' mode. This configures the context with just a URL and not an Engine, though an Engine is acceptable here as well. By skipping the Engine creation we don't even need a DBAPI to be available. Calls to context.execute() here emit the given string to the script output. """ url = config.get_main_option("sqlalchemy.url") context.configure( url=url, target_metadata=target_metadata, literal_binds=True, dialect_opts={"paramstyle": "named"}, ) with context.begin_transaction(): context.run_migrations() def run_migrations_online(): """ Run migrations in 'online' mode. In this scenario we need to create an Engine and associate a connection with the context. """ connectable = engine_from_config( config.get_section(config.config_ini_section), prefix="sqlalchemy.", poolclass=pool.NullPool, ) with connectable.connect() as connection: context.configure( connection=connection, target_metadata=target_metadata ) with context.begin_transaction(): context.run_migrations() if context.is_offline_mode(): run_migrations_offline() else: run_migrations_online() ================================================ FILE: server/backend/app/db/migrations/script.py.mako ================================================ """${message} Revision ID: ${up_revision} Revises: ${down_revision | comma,n} Create Date: ${create_date} """ from alembic import op import sqlalchemy as sa ${imports if imports else ""} # revision identifiers, used by Alembic. revision = ${repr(up_revision)} down_revision = ${repr(down_revision)} branch_labels = ${repr(branch_labels)} depends_on = ${repr(depends_on)} def upgrade(): ${upgrades if upgrades else "pass"} def downgrade(): ${downgrades if downgrades else "pass"} ================================================ FILE: server/backend/app/db/migrations/versions/6a4d853aabce_added_database.py ================================================ """Added database Revision ID: 6a4d853aabce Revises: 8f58cd776eda Create Date: 2020-12-23 13:49:00.640607 """ from alembic import op import sqlalchemy as sa # revision identifiers, used by Alembic. revision = '6a4d853aabce' down_revision = '8f58cd776eda' branch_labels = None depends_on = None def upgrade(): # ### commands auto generated by Alembic - please adjust! ### pass # ### end Alembic commands ### def downgrade(): # ### commands auto generated by Alembic - please adjust! ### pass # ### end Alembic commands ### ================================================ FILE: server/backend/app/db/migrations/versions/6d5c250f098c_added_device_file_server_address.py ================================================ """Added Device.file_server_address Revision ID: 6d5c250f098c Revises: fb245977373f Create Date: 2021-01-26 14:55:11.504148 """ from alembic import op import sqlalchemy as sa # revision identifiers, used by Alembic. revision = '6d5c250f098c' down_revision = 'fb245977373f' branch_labels = None depends_on = None def upgrade(): # ### commands auto generated by Alembic - please adjust! ### op.add_column('device', sa.Column('file_server_address', sa.String(), nullable=True)) # ### end Alembic commands ### def downgrade(): # ### commands auto generated by Alembic - please adjust! ### op.drop_column('device', 'file_server_address') # ### end Alembic commands ### ================================================ FILE: server/backend/app/db/migrations/versions/8f58cd776eda_added_database.py ================================================ """Added database Revision ID: 8f58cd776eda Revises: Create Date: 2020-12-23 13:38:47.314749 """ from alembic import op import sqlalchemy as sa # revision identifiers, used by Alembic. revision = '8f58cd776eda' down_revision = None branch_labels = None depends_on = None def upgrade(): # ### commands auto generated by Alembic - please adjust! ### op.create_table('device', sa.Column('id', sa.String(), nullable=False), sa.Column('description', sa.String(), nullable=True), sa.PrimaryKeyConstraint('id') ) op.create_table('statistic', sa.Column('device_id', sa.String(), nullable=False), sa.Column('datetime', sa.DateTime(), nullable=False), sa.Column('statistic_type', sa.Enum('REPORT', 'ALERT', name='statistictypeenum'), nullable=True), sa.Column('people_with_mask', sa.Integer(), nullable=True), sa.Column('people_without_mask', sa.Integer(), nullable=True), sa.Column('people_total', sa.Integer(), nullable=True), sa.ForeignKeyConstraint(['device_id'], ['device.id'], ), sa.PrimaryKeyConstraint('device_id', 'datetime') ) # ### end Alembic commands ### def downgrade(): # ### commands auto generated by Alembic - please adjust! ### op.drop_table('statistic') op.drop_table('device') # ### end Alembic commands ### ================================================ FILE: server/backend/app/db/migrations/versions/fb245977373f_added_video_file_table.py ================================================ """Added video_file table Revision ID: fb245977373f Revises: 6a4d853aabce Create Date: 2021-01-21 22:47:20.335928 """ from alembic import op import sqlalchemy as sa # revision identifiers, used by Alembic. revision = 'fb245977373f' down_revision = '6a4d853aabce' branch_labels = None depends_on = None def upgrade(): # ### commands auto generated by Alembic - please adjust! ### op.create_table('video_file', sa.Column('device_id', sa.String(), nullable=False), sa.Column('video_name', sa.String(), nullable=False), sa.ForeignKeyConstraint(['device_id'], ['device.id'], ), sa.PrimaryKeyConstraint('device_id', 'video_name') ) op.alter_column('statistic', 'people_total', existing_type=sa.INTEGER(), nullable=False) op.alter_column('statistic', 'people_with_mask', existing_type=sa.INTEGER(), nullable=False) op.alter_column('statistic', 'people_without_mask', existing_type=sa.INTEGER(), nullable=False) # ### end Alembic commands ### def downgrade(): # ### commands auto generated by Alembic - please adjust! ### op.alter_column('statistic', 'people_without_mask', existing_type=sa.INTEGER(), nullable=True) op.alter_column('statistic', 'people_with_mask', existing_type=sa.INTEGER(), nullable=True) op.alter_column('statistic', 'people_total', existing_type=sa.INTEGER(), nullable=True) op.drop_table('video_file') # ### end Alembic commands ### ================================================ FILE: server/backend/app/db/schema/__init__.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from .base import Base, engine, get_db_session, get_db_generator from .models import DeviceModel, StatisticsModel, VideoFilesModel from .schemas import DeviceSchema, StatisticSchema, VideoFileSchema ================================================ FILE: server/backend/app/db/schema/base.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from typing import Generator, Union from app.core.config import DB_URI from sqlalchemy import create_engine from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import Session, sessionmaker def get_db_session() -> Session: """ Get a new create a new database session. Returns: Session -- New database session. """ db = SessionLocal() try: return db finally: db.close() def get_db_generator() -> Generator: """ Get a new create a new database generator. Returns: Generator -- New database generator. """ db = SessionLocal() try: yield db finally: db.close() # Create ORM engine and session engine = create_engine(DB_URI) SessionLocal = sessionmaker(bind=engine) # Construct a base class for declarative class definitions Base = declarative_base() ================================================ FILE: server/backend/app/db/schema/models.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from app.db.schema import Base from app.db.utils import StatisticTypeEnum from sqlalchemy import Column, DateTime, Enum, ForeignKey, Integer, String from sqlalchemy.orm import relationship class StatisticsModel(Base): __tablename__ = "statistic" device_id = Column( String, ForeignKey("device.id"), primary_key=True, ) datetime = Column(DateTime(timezone=True), primary_key=True) statistic_type = Column(Enum(StatisticTypeEnum, nullable=False)) people_with_mask = Column(Integer, nullable=False) people_without_mask = Column(Integer, nullable=False) people_total = Column(Integer, nullable=False) class VideoFilesModel(Base): __tablename__ = "video_file" device_id = Column( String, ForeignKey("device.id"), primary_key=True, ) video_name = Column(String, primary_key=True) class DeviceModel(Base): __tablename__ = "device" id = Column(String, primary_key=True) description = Column(String) file_server_address = Column(String) statistics = relationship("StatisticsModel", cascade="all, delete") video_files = relationship("VideoFilesModel", cascade="all, delete") ================================================ FILE: server/backend/app/db/schema/schemas.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from datetime import datetime from typing import Optional, List from app.db.utils import StatisticTypeEnum from pydantic import BaseModel class StatisticSchema(BaseModel): device_id: str datetime: datetime statistic_type: StatisticTypeEnum people_with_mask: int people_without_mask: int people_total: int class Config: orm_mode = True class DeviceSchema(BaseModel): id: str description: Optional[str] = None file_server_address: Optional[str] = None statistics: List[StatisticSchema] = [] class Config: orm_mode = True class VideoFileSchema(BaseModel): device_id: str video_name: str class Config: orm_mode = True ================================================ FILE: server/backend/app/db/utils/__init__.py ================================================ from .enums import StatisticTypeEnum from .utils import convert_timestamp_to_datetime, get_enum_type ================================================ FILE: server/backend/app/db/utils/enums.py ================================================ from enum import Enum class StatisticTypeEnum(str, Enum): REPORT = "REPORT" ALERT = "ALERT" ================================================ FILE: server/backend/app/db/utils/utils.py ================================================ from datetime import datetime, timezone from .enums import StatisticTypeEnum def convert_timestamp_to_datetime(timestamp: float) -> datetime: """ Convert timestamp date format to datetime. Arguments: timestamp {float} -- Input timestamp. Returns: datetime -- Datetime formatted object which represents the same information as timestamp. """ return datetime.fromtimestamp(timestamp, timezone.utc) def get_enum_type(statistic_type: str) -> StatisticTypeEnum: """ Convert string object to enum. Arguments: statistic_type {str} -- Input string. Returns: StatisticTypeEnum -- Enum corresponding to statistic_type. """ return ( StatisticTypeEnum.ALERT if statistic_type.lower() == "alerts" else StatisticTypeEnum.REPORT ) ================================================ FILE: server/backend/app/main.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from fastapi import FastAPI from app.api import device_router, statistic_router app = FastAPI() app.include_router(device_router) app.include_router(statistic_router) @app.get("/") def health_check(): """ API health check used by the load balancer. """ return {"statusCode": 200} ================================================ FILE: server/backend/app/mqtt/broker.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from app.core.config import MQTT_BROKER, MQTT_BROKER_PORT from paho.mqtt import client as mqtt_client from typing import Callable def connect_mqtt_broker(client_id: str, cb_connect: Callable=None) -> mqtt_client: """ Connect to MQTT broker. Arguments: client_id {str} -- Client process id. cb_connect {Callable} -- Callback for on_connect Returns: mqtt_client -- MQTT client. """ def on_connect(client, userdata, flags, code): if code == 0: print("Connected to MQTT Broker") if cb_connect is not None: cb_connect(client) else: print(f"Failed to connect, return code {code}\n") def on_disconnect(client, userdata, code): print("MQTT Broker disconnected") client = mqtt_client.Client(client_id) client.on_connect = on_connect client.on_disconnect = on_disconnect client.connect(MQTT_BROKER, MQTT_BROKER_PORT) return client ================================================ FILE: server/backend/app/mqtt/publisher.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ import json import random import time from datetime import datetime, timezone, timedelta from paho.mqtt import client as mqtt_client def publish(client): for msg_count in range(0, 6): time.sleep(1) device_id = f"Device_{msg_count}" # Test hello topic = "hello" hello_msg = { "id": device_id, "description": f"Description {msg_count}", } hello_result = client.publish(topic, json.dumps(hello_msg)) if hello_result[0] == 0: print(f"Send `{hello_msg}` to topic `{topic}`") else: print(f"Failed to send message to topic {topic}") # Test alert time.sleep(1) topics = ["alerts", "receive-from-jetson"] now = datetime.now(timezone.utc) for _ in range(0, 600): topic = random.choice(topics) if topic == "alerts": people_with_mask = random.randint(2000, 3000) people_without_mask = random.randint(0, 1000) else: people_with_mask = random.randint(0, 1000) people_without_mask = random.randint(500, 1000) alert_msg = { "device_id": device_id, "timestamp": datetime.timestamp(now), "people_with_mask": people_with_mask, "people_without_mask": people_without_mask, "people_total": people_with_mask + people_without_mask, } result = client.publish(topic, json.dumps(alert_msg)) if result[0] == 0: print(f"Send `{alert_msg}` to topic `{topic}`") else: print(f"Failed to send message to topic {topic}") now += timedelta(minutes=1) def connect_mqtt_broker(client_id: str): def on_connect(client, userdata, flags, code): if code == 0: print("Connected to MQTT Broker") else: print(f"Failed to connect, return code {code}\n") client = mqtt_client.Client(client_id) client.on_connect = on_connect # client.connect("3.17.17.197", 1883) client.connect("0.0.0.0", 1883) return client def run(): client = connect_mqtt_broker(client_id=f"publisher-{random.randint(0, 10)}") client.loop_start() publish(client) if __name__ == "__main__": run() ================================================ FILE: server/backend/app/mqtt/subscriber.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ import json from app.core.config import SUBSCRIBER_CLIENT_ID, MQTT_HELLO_TOPIC,\ MQTT_ALERT_TOPIC, MQTT_SEND_TOPIC,\ MQTT_REPORT_TOPIC, MQTT_FILES_TOPIC from app.db.cruds import create_device, create_statistic, update_files, update_device from app.db.schema import get_db_session from app.db.utils import convert_timestamp_to_datetime, get_enum_type from broker import connect_mqtt_broker from paho.mqtt import client as mqtt_client from sqlalchemy.exc import IntegrityError MQTT_CLIENT_TOPICS = [ # topic, QoS (MQTT_HELLO_TOPIC, 2), (MQTT_FILES_TOPIC, 2), (MQTT_ALERT_TOPIC, 2), (MQTT_REPORT_TOPIC, 2), (MQTT_SEND_TOPIC, 2), ] def subscribe(client: mqtt_client): """ Subscribe client to topic. Arguments: client {mqtt_client} -- Client process id. """ def on_message(client, userdata, msg): database_session = get_db_session() try: process_message(database_session, msg) finally: database_session.close() client.subscribe(MQTT_CLIENT_TOPICS) client.on_message = on_message def process_message(database_session, msg): """ Process message sent to topic. Arguments: database_session {Session} -- Database session. msg {str} -- Received message. """ message = json.loads(msg.payload.decode()) topic = msg.topic if topic == "hello": # Register new Jetson device device_id = message["device_id"] try: device_information = { "id": device_id, "description": message["description"], } device = create_device( db_session=database_session, device_information=device_information, ) print("Added device") except IntegrityError: print(f"A device with id={device_id} already exists") elif topic in ["alerts", "receive-from-jetson"]: try: # Receive alert or report and save it to the database statistic_information = { "device_id": message["device_id"], "datetime": convert_timestamp_to_datetime(message["timestamp"]), "statistic_type": get_enum_type(topic), "people_with_mask": message["people_with_mask"], "people_without_mask": message["people_without_mask"], "people_total": message["people_total"], } statistic = create_statistic( db_session=database_session, statistic_information=statistic_information, ) print(f"Added statistic") except IntegrityError: print(f"Error, the statistic already exist") elif topic == "video-files": try: print(f"Adding files for device_id: {message['device_id']}") new_information = {"file_server_address": message["file_server"]} update_device(db_session=database_session, device_id=message["device_id"], new_device_information=new_information) update_files(db_session=database_session, device_id=message["device_id"], file_list=message["file_list"]) except Exception as e: print(f"Exception trying to update files: {e}") elif topic == "send-to-jetson": # Just monitoring this channel, useful for debugging print(f"Detected info sent to device_id: {message['device_id']}") def main(): client = connect_mqtt_broker(client_id=SUBSCRIBER_CLIENT_ID, cb_connect=subscribe) client.loop_forever() if __name__ == "__main__": main() ================================================ FILE: server/backend/prestart.sh ================================================ #!/bin/bash # Wait database initialization and apply migrations if needed sleep 3 alembic upgrade head # Init subscriber process python app/mqtt/subscriber.py & ================================================ FILE: server/backend/requirements.txt ================================================ paho-mqtt==1.5.1 SQLAlchemy==1.3.21 python-dotenv==0.15.0 psycopg2-binary==2.8.6 alembic==1.4.3 pytest==6.2.1 pydantic==1.7.3 ================================================ FILE: server/backend/test_crud.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ import random from datetime import datetime, timezone import pytest from sqlalchemy.exc import IntegrityError from sqlalchemy.orm.exc import NoResultFound from app.db.cruds import ( create_device, create_statistic, delete_device, delete_statistic, get_device, get_devices, get_statistic, get_statistics, update_device, update_statistic, ) from app.db.schema import get_db_session from app.db.utils import StatisticTypeEnum, convert_timestamp_to_datetime DEVICE_ID = "test" database_session = get_db_session() # Device def test_create_device(): info = { "id": DEVICE_ID, "description": "test description", } device = create_device(db_session=database_session, device_information=info) assert device.id == DEVICE_ID assert device.description == "test description" def test_create_device_more_fields(): with pytest.raises(TypeError): info = { "id": DEVICE_ID, "description": "test description", "test_field_1": 1, "test_field_2": 2, } create_device(db_session=database_session, device_information=info) def test_create_same_device(): with pytest.raises(IntegrityError): info = { "id": DEVICE_ID, "description": "test description", } device = create_device( db_session=database_session, device_information=info ) def test_get_device(): device = get_device(db_session=database_session, device_id=DEVICE_ID) assert device.id == DEVICE_ID assert device.description == "test description" def test_update_device(): device = update_device( db_session=database_session, device_id=DEVICE_ID, new_device_information={"description": "new description"}, ) assert device.id == DEVICE_ID assert device.description == "new description" # Statistics def test_create_statistic(): people_with_mask = 4 people_without_mask = 7 now = convert_timestamp_to_datetime(1609780971.514455) # now = datetime(2021, 1, 4, 17, 22, 51, 514455, tzinfo=timezone.utc) stat_info = { "device_id": DEVICE_ID, "datetime": now, "statistic_type": StatisticTypeEnum.ALERT, "people_with_mask": people_with_mask, "people_without_mask": people_without_mask, "people_total": people_with_mask + people_without_mask, } statistic = create_statistic( db_session=database_session, statistic_information=stat_info ) assert statistic.device_id == DEVICE_ID assert statistic.datetime == now.replace(tzinfo=None) assert statistic.statistic_type == StatisticTypeEnum.ALERT assert statistic.people_with_mask == people_with_mask assert statistic.people_without_mask == people_without_mask assert statistic.people_total == people_with_mask + people_without_mask def test_create_same_statistic(): people_with_mask = 4 people_without_mask = 7 # now = datetime(2021, 1, 4, 17, 22, 51, 514455, tzinfo=timezone.utc) now = convert_timestamp_to_datetime(1609780971.514455) with pytest.raises(IntegrityError): stat_info = { "device_id": DEVICE_ID, "datetime": now, "statistic_type": StatisticTypeEnum.ALERT, "people_with_mask": people_with_mask, "people_without_mask": people_without_mask, "people_total": people_with_mask + people_without_mask, } create_statistic( db_session=database_session, statistic_information=stat_info ) def test_create_another_statistic(): people_with_mask = 5 people_without_mask = 8 # now = datetime(2021, 1, 5, 17, 22, 51, 514455, tzinfo=timezone.utc) now = convert_timestamp_to_datetime(1609867371.514455) stat_info = { "device_id": DEVICE_ID, "datetime": now, "statistic_type": StatisticTypeEnum.ALERT, "people_with_mask": people_with_mask, "people_without_mask": people_without_mask, "people_total": people_with_mask + people_without_mask, } statistic = create_statistic( db_session=database_session, statistic_information=stat_info ) assert statistic.device_id == DEVICE_ID assert statistic.datetime == now.replace(tzinfo=None) assert statistic.statistic_type == StatisticTypeEnum.ALERT assert statistic.people_with_mask == people_with_mask assert statistic.people_without_mask == people_without_mask assert statistic.people_total == people_with_mask + people_without_mask def test_get_statistic(): people_with_mask = 4 people_without_mask = 7 # now = datetime(2021, 1, 4, 17, 22, 51, 514455, tzinfo=timezone.utc) now = convert_timestamp_to_datetime(1609780971.514455) statistic = get_statistic( db_session=database_session, device_id=DEVICE_ID, datetime=now ) assert statistic.device_id == DEVICE_ID assert statistic.datetime == now.replace(tzinfo=None) assert statistic.statistic_type == StatisticTypeEnum.ALERT assert statistic.people_with_mask == people_with_mask assert statistic.people_without_mask == people_without_mask assert statistic.people_total == people_with_mask + people_without_mask def test_update_statistic(): people_without_mask = 7 # now = datetime(2021, 1, 4, 17, 22, 51, 514455, tzinfo=timezone.utc) now = convert_timestamp_to_datetime(1609780971.514455) statistic = update_statistic( db_session=database_session, device_id=DEVICE_ID, datetime=now, new_statistic_information={"people_with_mask": 20, "people_total": 27}, ) assert statistic.device_id == DEVICE_ID assert statistic.datetime == now.replace(tzinfo=None) assert statistic.statistic_type == StatisticTypeEnum.ALERT assert statistic.people_with_mask == 20 assert statistic.people_without_mask == people_without_mask assert statistic.people_total == 27 def test_delete_statistic(): people_with_mask = 20 people_without_mask = 7 # now = datetime(2021, 1, 4, 17, 22, 51, 514455, tzinfo=timezone.utc) now = convert_timestamp_to_datetime(1609780971.514455) statistic = delete_statistic( db_session=database_session, device_id=DEVICE_ID, datetime=now, ) assert statistic.device_id == DEVICE_ID assert statistic.datetime == now.replace(tzinfo=None) assert statistic.statistic_type == StatisticTypeEnum.ALERT assert statistic.people_with_mask == people_with_mask assert statistic.people_without_mask == people_without_mask assert statistic.people_total == people_with_mask + people_without_mask def test_delete_device(): device = delete_device(db_session=database_session, device_id=DEVICE_ID) assert device.id == DEVICE_ID assert device.description == "new description" def test_get_deleted_device(): with pytest.raises(NoResultFound): get_device(db_session=database_session, device_id=DEVICE_ID) def test_update_deleted_device(): with pytest.raises(NoResultFound): update_device( db_session=database_session, device_id=DEVICE_ID, new_device_information={"description": "new test description"}, ) def test_get_devices(): devices = get_devices(db_session=database_session) assert devices == [] def test_get_devices(): stats = get_statistics(db_session=database_session) assert stats == [] database_session.close() ================================================ FILE: server/backend.env.template ================================================ MQTT_BROKER=mosquitto MQTT_BROKER_PORT=1883 SUBSCRIBER_CLIENT_ID=server_backend MQTT_HELLO_TOPIC=hello MQTT_ALERT_TOPIC=alerts MQTT_REPORT_TOPIC=receive-from-jetson MQTT_SEND_TOPIC=send-to-jetson ================================================ FILE: server/build_docker.sh ================================================ #!/bin/bash sudo docker-compose build sudo docker-compose up -d ================================================ FILE: server/database.env.template ================================================ POSTGRES_USER= POSTGRES_PASSWORD= POSTGRES_PORT=5432 POSTGRES_DB= ================================================ FILE: server/docker-compose.yml ================================================ version: '3.1' services: db: image: postgres:13.2 env_file: - database.env volumes: - pgdata:/var/lib/postgresql/data ports: - 5432:5432 mosquitto: image: eclipse-mosquitto:1.6.13 ports: - 1883:1883 - 8883:8883 volumes: - mosquitto-data:/mosquitto/data - mosquitto-logs:/mosquitto/logs - mosquitto-conf:/mosquitto/config - ./mosquitto.conf:/mosquitto/config/mosquitto.conf restart: unless-stopped backend: image: backend volumes: - ./backend:/app build: context: backend depends_on: - db ports: - 80:80 command: bash -c "sleep 5 && /start-reload.sh" env_file: - backend.env - database.env streamlit: build: frontend command: "streamlit run main.py" ports: - "8501:8501" volumes: - "./frontend:/usr/src/app" env_file: - frontend.env depends_on: - db - backend links: - backend volumes: pgdata: mosquitto-data: mosquitto-logs: mosquitto-conf: ================================================ FILE: server/frontend/Dockerfile ================================================ FROM python:3.7 EXPOSE 8501 WORKDIR /usr/src/app COPY requirements.txt ./ RUN pip install -r requirements.txt COPY . . ================================================ FILE: server/frontend/main.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ import os import json import streamlit as st from datetime import datetime, time, timezone from session_manager import get_state from utils.api_utils import ( get_device, get_devices, get_statistics_from_to, get_device_files, ) from utils.format_utils import create_chart, format_data from paho.mqtt import client as mqtt_client MQTT_BROKER = os.environ["MQTT_BROKER"] MQTT_BROKER_PORT = int(os.environ["MQTT_BROKER_PORT"]) MQTT_CLIENT_ID = os.environ["MQTT_CLIENT_ID"] MQTT_TOPIC_COMMANDS = "commands" MQTT_TOPIC_STATUS = "device-status" CMD_FILE_SAVE = "save_file" CMD_STREAMING_START = "streaming_start" CMD_STREAMING_STOP = "streaming_stop" CMD_INFERENCE_RESTART = "inference_restart" CMD_FILESERVER_RESTART = "fileserver_restart" CMD_REQUEST_STATUS = "status_request" state = get_state() def display_sidebar(all_devices, state): """ Display sidebar information. """ st.sidebar.subheader("Device selection") state.selected_device = st.sidebar.selectbox( "Selected device", all_devices, index=all_devices.index(state.selected_device if state.selected_device else None), ) st.sidebar.subheader("Filters") state.date_filter = st.sidebar.date_input( "From date - To date", ( state.date_filter[0] if state.date_filter and len(state.date_filter) == 2 else datetime.now(timezone.utc), state.date_filter[1] if state.date_filter and len(state.date_filter) == 2 else datetime.now(timezone.utc), ), ) first_column, second_column = st.sidebar.beta_columns(2) state.from_time = first_column.time_input( "From time", state.from_time if state.from_time else time(0, 0) ) state.to_time = second_column.time_input( "To time", state.to_time if state.to_time else time(23, 45) ) state.group_data_by = st.sidebar.selectbox( "Group data by", ["Second", "Minute", "Hour", "Day", "Week", "Month"], index=2, ) state.show_only_one_chart = st.sidebar.checkbox("Show only one chart", value=True) def display_device(state): """ Display specific device information. """ selected_device = state.selected_device device = get_device(selected_device) if device is None: st.write("Seems that something went wrong while getting the device information.") else: st.header(f"Device: {device['id']}") if state.mqtt_last_status: status = state.mqtt_last_status # shortcut if not status["device_address"]: st.write( ":warning: **Set MASKCAM_DEVICE_ADDRESS on the device to enable " "streaming and file download links**" ) device_status = st.beta_container() col1, col2 = device_status.beta_columns(2) col1.write("🟢 Device connected " f"*(Last update: {status['time']})*") if not status["streaming_address"] or status["streaming_address"] == "N/A": col2.write(":red_circle: Streaming is stopped") else: if status["device_address"]: col2.write( f"🟢 " "Streaming enabled", unsafe_allow_html=True, ) else: col2.write( "🟢 Streaming enabled (unknown device address)", ) device_status.write( f"**Save videos: {status['save_current_files']}**" f" | *Inference runtime: {status['inference_runtime']}*" f" | *Fileserver runtime: {status['fileserver_runtime']}*" ) mqtt_status = st.empty() # Might be changed in real time during connection if not state.mqtt_last_status: if not state.mqtt_status: # Loading page, first connection attempt to device send_mqtt_command(device["id"], CMD_REQUEST_STATUS, mqtt_status) else: mqtt_set_status(mqtt_status, state.mqtt_status) cols = st.beta_columns(6) # Buttons from right to left with cols.pop(): if st.button("Restart Deepstream"): send_mqtt_command(device["id"], CMD_INFERENCE_RESTART, mqtt_status) with cols.pop(): if st.button("Restart file server"): send_mqtt_command(device["id"], CMD_FILESERVER_RESTART, mqtt_status) with cols.pop(): if st.button("Stop streaming"): send_mqtt_command(device["id"], CMD_STREAMING_STOP, mqtt_status) with cols.pop(): if st.button("Start streaming"): send_mqtt_command(device["id"], CMD_STREAMING_START, mqtt_status) with cols.pop(): if st.button("Save a video"): send_mqtt_command(device["id"], CMD_FILE_SAVE, mqtt_status) with cols.pop(): if st.button("Refresh status"): send_mqtt_command(device["id"], CMD_REQUEST_STATUS, mqtt_status) st.header("Reported statistics") device_statistics = None date_filter = state.date_filter if len(date_filter) == 2: datetime_from = f"{date_filter[0]}T{state.from_time}" datetime_to = f"{date_filter[1]}T{state.to_time}" device_statistics = get_statistics_from_to(selected_device, datetime_from, datetime_to) if not device_statistics: st.write("The selected device has no statistics to show for the given filters.") else: reports, alerts = format_data(device_statistics, state.group_data_by) if state.show_only_one_chart: complete_chart = create_chart(reports=reports, alerts=alerts) st.plotly_chart(complete_chart, use_container_width=True) else: st.subheader("Reports") if reports: report_chart = create_chart(reports=reports) st.plotly_chart(report_chart, use_container_width=True) else: st.write("The selected device has no reports to show for the given filters.") st.subheader("Alerts") if alerts: alerts_chart = create_chart(alerts=alerts) st.plotly_chart(alerts_chart, use_container_width=True) else: st.write("The selected device has no alerts to show for the given filters.") device_files = get_device_files(device_id=selected_device) st.subheader("Saved video files on device") if not device_files: st.write("The selected device has no saved files yet") else: server_address = None if not state.mqtt_last_status: st.write(":warning: **Downloads will fail since device is NOT connected**") elif state.mqtt_last_status["device_address"] is None: st.write( ":warning: **Set MASKCAM_DEVICE_ADDRESS on device to enable download links**" ) else: # file_server_address is valid when device_address=MASKCAM_DEVICE_ADDRESS is set server_address = f"{device['file_server_address']}" for file_instance in device_files: if server_address: url = f"{server_address}/{file_instance['video_name']}" st.markdown(f"[{file_instance['video_name']}]({url})") else: st.markdown(f"{file_instance['video_name']}") def mqtt_set_status(mqtt_status, text): state.mqtt_status = text mqtt_status.empty() mqtt_status.markdown(f"**MQTT status:** {text}") def _on_connect(client, userdata, flags, rc): if rc == 0: state.mqtt_connected = True def _on_message(client, userdata, msg): # This is the only topic the frontend subscribes to assert msg.topic == MQTT_TOPIC_STATUS if not state.selected_device: return message = json.loads(msg.payload.decode()) if message["device_id"] != state.selected_device: return state.mqtt_last_status = message @st.cache(allow_output_mutation=True) def restore_client(): client = mqtt_client.Client(MQTT_CLIENT_ID) client.connect(MQTT_BROKER, MQTT_BROKER_PORT) return client def get_mqtt_client(): client = restore_client() client.on_connect = _on_connect client.on_message = _on_message return client def mqtt_wait_connection(client, timeout): while not state.mqtt_connected and timeout: client.loop(timeout=1) timeout -= 1 def mqtt_wait_response(client, timeout): while not state.mqtt_last_status and timeout: client.loop(timeout=1) timeout -= 1 def send_mqtt_message_wait_response(topic, message, mqtt_status): # This function connects, sends a message and waits for the reply. Reconnects as needed try: client = get_mqtt_client() if not state.mqtt_connected: mqtt_set_status(mqtt_status, "Connecting...") mqtt_wait_connection(client, 5) state.mqtt_last_status = None # Reset status to await updated response # Since we're not running the client.loop() permanently, # the way to ensure that the MQTT client is connected is to try # to send a message and if it fails, try reconnecting. retry_publish = 2 # mosquitto disconnects after a while so at least use 2 here while retry_publish: retry_publish -= 1 client.subscribe(MQTT_TOPIC_STATUS) # Must be done after reconnection msg_info = client.publish(topic, json.dumps(message)) mqtt_set_status(mqtt_status, "Sending message...") timeout = 5 while not msg_info.rc and not msg_info.is_published() and timeout: client.loop(1) timeout -= 1 if msg_info.is_published(): mqtt_set_status(mqtt_status, ":clock3: Waiting device response...") retry_publish = 0 # Success: exit retry loop elif msg_info.rc: state.mqtt_connected = False mqtt_set_status(mqtt_status, ":clock3: Reconnecting...") client.reconnect() mqtt_wait_connection(client, 5) if not msg_info.is_published(): mqtt_set_status(mqtt_status, ":o: Message failed") return mqtt_wait_response(client, 5) if not state.mqtt_last_status: mqtt_set_status(mqtt_status, ":red_circle: Device not responding") except Exception as e: mqtt_set_status(mqtt_status, f":red_circle: Could not connect to MQTT broker: {e}") def send_mqtt_command(device_id, command, mqtt_status): send_mqtt_message_wait_response( MQTT_TOPIC_COMMANDS, {"device_id": device_id, "command": command}, mqtt_status ) def main(): st.set_page_config(page_title="Maskcam") st.title("MaskCam dashboard") all_devices = get_devices() display_sidebar(all_devices, state) if state.selected_device is None: st.write("Please select a device.") else: display_device(state) state.sync() if __name__ == "__main__": main() ================================================ FILE: server/frontend/requirements.txt ================================================ streamlit==0.74.1 requests==2.25.1 plotly==4.14.3 paho-mqtt==1.5.1 ================================================ FILE: server/frontend/session_manager.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ import streamlit as st from streamlit.hashing import _CodeHasher from streamlit.report_thread import get_report_ctx from streamlit.server.server import Server class _SessionState: def __init__(self, session, hash_funcs): """ Initialize SessionState instance. """ self.__dict__["_state"] = { "data": {}, "hash": None, "hasher": _CodeHasher(hash_funcs), "is_rerun": False, "session": session, } def __call__(self, **kwargs): """ Initialize state data once. """ for item, value in kwargs.items(): if item not in self._state["data"]: self._state["data"][item] = value def __getitem__(self, item): """ Return a saved state value, None if item is undefined. """ return self._state["data"].get(item, None) def __getattr__(self, item): """ Return a saved state value, None if item is undefined. """ return self._state["data"].get(item, None) def __setitem__(self, item, value): """Set state value.""" self._state["data"][item] = value def __setattr__(self, item, value): """ Set state value. """ self._state["data"][item] = value def clear(self): """ Clear session state and request a rerun. """ self._state["data"].clear() self._state["session"].request_rerun() def sync(self): """ Rerun the app with all state values up to date from the beginning to fix rollbacks. """ # Ensure to rerun only once to avoid infinite loops # caused by a constantly changing state value at each run. # # Example: state.value += 1 if self._state["is_rerun"]: self._state["is_rerun"] = False elif self._state["hash"] is not None: if self._state["hash"] != self._state["hasher"].to_bytes( self._state["data"], None ): self._state["is_rerun"] = True self._state["session"].request_rerun() self._state["hash"] = self._state["hasher"].to_bytes( self._state["data"], None ) def _get_session(): session_id = get_report_ctx().session_id session_info = Server.get_current()._get_session_info(session_id) if session_info is None: raise RuntimeError("Couldn't get your Streamlit Session object.") return session_info.session def get_state(hash_funcs=None): session = _get_session() if not hasattr(session, "_custom_session_state"): session._custom_session_state = _SessionState(session, hash_funcs) return session._custom_session_state ================================================ FILE: server/frontend/utils/api_utils.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ import json import os from typing import Dict, List import requests SERVER_URL = os.environ["SERVER_URL"] def get_devices(): """ Get all devices. """ response = requests.get(f"http://{SERVER_URL}/devices") devices_json = json.loads(response.content) if response.ok else [] devices = [None] devices.extend([device["id"] for device in devices_json]) return devices def get_device(device_id: str): """ Get specific device. Arguments: device_id {str} -- Device id. """ response = requests.get(f"http://{SERVER_URL}/devices/{device_id}") return json.loads(response.content) if response.ok else None def get_statistics_from_to(device_id, datetime_from, datetime_to): """ Get statistics from a specific device within a datetime range. Arguments: device_id {str} -- Device id. datetime_from {str} -- Datetime from. datetime_to {str} -- Datetime to. """ response = requests.get( f"http://{SERVER_URL}/devices/{device_id}/statistics?datefrom={datetime_from}&dateto={datetime_to}" ) return json.loads(response.content) if response.ok else None def get_device_files(device_id): """ Get files from a specific device Arguments: device_id {str} -- Device id. """ response = requests.get( f"http://{SERVER_URL}/files/{device_id}" ) return json.loads(response.content) if response.ok else None ================================================ FILE: server/frontend/utils/format_utils.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ from typing import Dict, List import pandas as pd import numpy as np import plotly.graph_objects as go from plotly.subplots import make_subplots def format_data(statistics: List = [], group_data_by: str = None): """ Format data to be displayed in the carts. Arguments: statistics {List} -- All device statistics. group_data_by {str} -- User selected aggregation option. """ reports, alerts = {}, {} for statistic in statistics: # Separate reports from alerts if statistic["statistic_type"] == "REPORT": if not reports: reports = create_statistics_dict() reports = add_information(reports, statistic) else: if not alerts: alerts = create_statistics_dict() alerts = add_information(alerts, statistic) # Aggregate data if reports: reports = group_data(reports, group_data_by) if alerts: alerts = group_data(alerts, group_data_by) return reports, alerts def group_data(data: Dict, group_data_by: str): """ Aggregate data using different criteria. Arguments: data {Dict} -- Data to aggregate. group_data_by {str} -- User selected aggregation option. """ data_df = pd.DataFrame.from_dict(data) data_df["dates"] = pd.to_datetime(data_df["dates"]) criterion = "H" if group_data_by == "Second": criterion = "S" elif group_data_by == "Minute": criterion = "T" elif group_data_by == "Day": criterion = "D" elif group_data_by == "Week": criterion = "W" elif group_data_by == "Month": criterion = "M" group = ( data_df.resample(criterion, on="dates") .agg( { "people_total": "sum", "people_with_mask": "sum", "people_without_mask": "sum", } ) .reset_index() ) # Calculate new mask percentage group["mask_percentage"] = ( group["people_with_mask"] * 100 / group["people_total"] ) group["mask_percentage"] = group["mask_percentage"].replace( [np.inf, -np.inf], 0 ) group["visible_people"] = ( group["people_with_mask"] + group["people_without_mask"] ) # Drop empty lines group.dropna(subset=["mask_percentage"], inplace=True) group = group.to_dict() grouped_data = { "dates": [t.to_pydatetime() for t in group["dates"].values()], "people_with_mask": list(group["people_with_mask"].values()), "people_total": list(group["people_total"].values()), "mask_percentage": list(group["mask_percentage"].values()), "visible_people": list(group["visible_people"].values()), } return grouped_data def create_statistics_dict(): """ Chart data structure. """ return { "dates": [], "mask_percentage": [], "people_total": [], "people_with_mask": [], "people_without_mask": [], "visible_people": [], } def add_information(statistic_dict: Dict, statistic_information: Dict): """ Add information to existing dict. Arguments: statistic_dict {Dict} -- Existing dict. statistic_information {Dict} -- Data to add. """ statistic_dict["dates"].append(statistic_information["datetime"]) total = statistic_information["people_total"] people_with_mask = statistic_information["people_with_mask"] people_without_mask = statistic_information["people_without_mask"] statistic_dict["people_total"].append(total) statistic_dict["people_with_mask"].append(people_with_mask) statistic_dict["people_without_mask"].append(people_without_mask) mask_percentage = ( statistic_information["people_with_mask"] * 100 / total if total != 0 else 0 ) statistic_dict["mask_percentage"].append(mask_percentage) statistic_dict["visible_people"].append( people_with_mask + people_without_mask ) return statistic_dict def create_chart(reports: Dict = {}, alerts: Dict = {}): """ Create Plotly chart. Arguments: reports {Dict} -- Reports data. alerts {Dict} -- Alerts data. """ # Create figure with secondary y-axis figure = make_subplots(specs=[[{"secondary_y": True}]]) if reports: report_colors = { "people_total": "darkslategray", "people_with_mask": "cadetblue", "mask_percentage": "limegreen", "visible_people": "royalblue", } figure = add_trace(reports, figure, report_colors, trace_type="report") if alerts: alert_colors = { "people_total": "darkred", "people_with_mask": "salmon", "mask_percentage": "orange", "visible_people": "indianred", } figure = add_trace(alerts, figure, alert_colors, trace_type="alert") # Set x-axis title figure.update_xaxes(title_text="Datetime") # Set y-axes titles figure.update_yaxes(title_text="Number of people", secondary_y=False) figure.update_yaxes( title_text="Mask Percentage", secondary_y=True, rangemode="tozero" ) figure.update_layout( xaxis_tickformat="%H:%M:%S
%Y/%d/%m", barmode="group", bargap=0.4, bargroupgap=0.1, legend={ "orientation": "h", "yanchor": "bottom", "y": 1.02, "xanchor": "right", "x": 1, }, margin=dict(t=50, b=30, r=30), font=dict( size=10, ), ) return figure def add_trace(trace_information: Dict, figure, colors: Dict, trace_type=""): """ Add new trace to Plotly chart. Arguments: trace_information {Dict} -- Trace to add. figure -- Plotly chart. colors {Dict} -- Colors to use in the new trace. trace_type -- Alert or Report. """ figure.add_trace( go.Bar( x=trace_information["dates"], y=trace_information["people_total"], name="People" if not trace_type else f"People {trace_type}", marker_color=colors["people_total"], ), secondary_y=False, ) figure.add_trace( go.Bar( x=trace_information["dates"], y=trace_information["people_with_mask"], name="Masks" if not trace_type else f"Masks {trace_type}", marker_color=colors["people_with_mask"], ), secondary_y=False, ) figure.add_trace( go.Bar( x=trace_information["dates"], y=trace_information["visible_people"], name="Visible" if not trace_type else f"Visible {trace_type}", marker_color=colors["visible_people"], ), secondary_y=False, ) figure.add_trace( go.Scatter( x=trace_information["dates"], y=trace_information["mask_percentage"], name="Mask %" if not trace_type else f"Mask % {trace_type}", marker_color=colors["mask_percentage"], ), secondary_y=True, ) return figure ================================================ FILE: server/frontend.env.template ================================================ SERVER_URL=backend MQTT_BROKER=mosquitto MQTT_BROKER_PORT=1883 MQTT_CLIENT_ID=server_frontend ================================================ FILE: server/mosquitto.conf ================================================ bind_address 0.0.0.0 ================================================ FILE: server/server_setup.sh ================================================ #!/bin/bash sudo apt-get update sudo apt install docker.io -y sudo systemctl start docker sudo curl -L "https://github.com/docker/compose/releases/download/1.27.4/docker-compose-$(uname -s)-$(uname -m)" -o /usr/local/bin/docker-compose sudo chmod +x /usr/local/bin/docker-compose sudo bash ./build_docker.sh ================================================ FILE: server/stop_docker.sh ================================================ #!/bin/bash sudo docker-compose down ================================================ FILE: utils/combine_coco.py ================================================ # %% import json import sys # %% def merge_2_into_1(json1, json2): # ID offsets to put annotations and images from json2 into json1 annotations1 = json1["annotations"] new_id_ann = 1 + max(*[ann["id"] for ann in annotations1]) images1 = json1["images"] new_id_im = 1 + max(*[im["id"] for im in images1]) categories1 = json1["categories"] categories2 = json2["categories"] annotations2 = json2["annotations"] images2 = json2["images"] # Append images2 to images1 and build a map from old to new IDs img_ids_1 = {im1["file_name"]: im1["id"] for im1 in images1} map_im2_to_im1 = {} for im2 in images2: if im2["file_name"] in img_ids_1: # im2 was already in JSON 1 map_im2_to_im1[im2["id"]] = img_ids_1[im2["file_name"]] else: map_im2_to_im1[im2["id"]] = new_id_im im2["id"] = new_id_im images1.append(im2) # Add to JSON 1 new_id_im += 1 cat_ids_1 = {cat1["name"]: cat1["id"] for cat1 in categories1} map_cat2_to_cat1 = {cat2["id"]: cat_ids_1[cat2["name"]] for cat2 in categories2} print(f"Cats 1: {categories1}") print(f"Cats 2: {categories2}") print(f"Map: {map_cat2_to_cat1}") for annotation2 in annotations2: annotation2["image_id"] = map_im2_to_im1[annotation2["image_id"]] annotation2["category_id"] = map_cat2_to_cat1[annotation2["category_id"]] annotation2["id"] = new_id_ann annotations1.append(annotation2) new_id_ann += 1 def print_coco(js): print( f"Images: {len(js['images'])} " f"| Annotations: {len(js['annotations'])} " f"| Categories: {len(js['categories'])}" ) def open_coco(file_name): with open(file_name, "r") as f: js = json.load(f) print(f"Loaded JSON file: {file_name}") print_coco(js) return js # %% file1 = sys.argv[1] other_files = sys.argv[2:] # file1 = "../new_dataset/annotations/instances_1_30s.json" # file2 = "../new_dataset/annotations/instances_2_30s.json" json1 = open_coco(file1) for file2 in other_files: merge_2_into_1(json1, open_coco(file2)) output_filename = "merge_result.json" with open(output_filename, "w") as output_file: json.dump(json1, output_file) print(f"Saved {output_filename}") print_coco(json1) ================================================ FILE: utils/gst_capabilities.sh ================================================ gst-launch-1.0 --gst-debug=v4l2src:5 v4l2src device=/dev/video0 ! fakesink 2>&1 | sed -une '/caps of src/ s/[:;] /\n/gp' ================================================ FILE: utils/mqtt-test/broker.py ================================================ import os from paho.mqtt import client as mqtt_client MQTT_BROKER = os.environ["MQTT_BROKER"] MQTT_BROKER_PORT = 1883 def connect_mqtt_broker(client_id: str) -> mqtt_client: def on_connect(client, userdata, flags, code): if code == 0: print("Connected to MQTT Broker") else: print(f"Failed to connect, return code {code}\n") client = mqtt_client.Client(client_id) client.on_connect = on_connect client.connect(MQTT_BROKER, MQTT_BROKER_PORT) return client ================================================ FILE: utils/mqtt-test/publisher.py ================================================ import json import random import time from datetime import datetime, timezone from broker import connect_mqtt_broker def publish(client): while True: time.sleep(1) # Test topic test_topic = "test" test_result = client.publish( test_topic, json.dumps( { "msg": f"testing", "timestamp": datetime.timestamp(datetime.now(timezone.utc)), } ), ) if test_result[0] == 0: print(f"Send test msg to test topic") else: print(f"Failed to send message to test topic") def run(): client = connect_mqtt_broker(client_id=f"publisher-test") client.loop_start() publish(client) if __name__ == "__main__": run() ================================================ FILE: utils/mqtt-test/suscriber.py ================================================ import json from broker import connect_mqtt_broker from paho.mqtt import client as mqtt_client def subscribe(client: mqtt_client): def on_message(client, userdata, msg): message = json.loads(msg.payload.decode()) topic = msg.topic print(f"Message received in topic: {topic}") print(message) test_topic = "test" client.subscribe(test_topic) client.on_message = on_message def main(): client = connect_mqtt_broker(client_id=f"subscriber-test") subscribe(client) client.loop_forever() if __name__ == "__main__": main() ================================================ FILE: utils/onnx_fix_mobilenet.py ================================================ import sys import onnx_graphsurgeon as gs import onnx import numpy as np """ This code is intended to change the uint8 input type (not supported by TensorRT) The input ONNX file can be produced by converting with: https://github.com/onnx/tensorflow-onnx After conversion, the output ONNX should be able to be converted as: /usr/src/tensorrt/bin/trtexec --fp16 --onnx=input_file.onnx --explicitBatch --saveEngine=output_file.trt But that step still fails due to NonMaxSuppression plugin not found (operation not supported by TensorRT) """ input_onnx = sys.argv[1] output_onnx = sys.argv[2] graph = gs.import_onnx(onnx.load(input_onnx)) for inp in graph.inputs: inp.dtype = np.float32 onnx.save(gs.export_onnx(graph), output_onnx) ================================================ FILE: utils/remove_images_coco.py ================================================ # %% import json import sys # %% def merge_2_into_1(json1, json2): # TODO: Remap category IDs according to their name # ID offsets to put annotations and images from json2 into json1 annotations1 = json1["annotations"] new_id_ann = 1 + max(*[ann["id"] for ann in annotations1]) images1 = json1["images"] new_id_im = 1 + max(*[im["id"] for im in images1]) annotations2 = json2["annotations"] images2 = json2["images"] # Append images2 to images1 and build a map from old to new IDs img_ids_1 = {im1["file_name"]: im1["id"] for im1 in images1} map_im2_to_im1 = {} for im2 in images2: if im2["file_name"] in img_ids_1: # im2 was already in JSON 1 map_im2_to_im1[im2["id"]] = img_ids_1[im2["file_name"]] else: map_im2_to_im1[im2["id"]] = new_id_im im2["id"] = new_id_im images1.append(im2) # Add to JSON 1 new_id_im += 1 for annotation2 in annotations2: annotation2["image_id"] = map_im2_to_im1[annotation2["image_id"]] annotation2["id"] = new_id_ann annotations1.append(annotation2) new_id_ann += 1 def print_coco(js): print( f"Images: {len(js['images'])} " f"| Annotations: {len(js['annotations'])} " f"| Categories: {len(js['categories'])}" ) def open_coco(file_name): with open(file_name, "r") as f: js = json.load(f) print(f"Loaded JSON file: {file_name}") print_coco(js) return js # %% file1 = sys.argv[1] images_to_remove = sys.argv[2:] json1 = open_coco(file1) images = json1["images"] annotations = json1["annotations"] new_images = [] new_annotations = [] img_ids_to_keep = set() # Append images2 to images1 and build a map from old to new IDs for im in images: if im["file_name"] not in images_to_remove: img_ids_to_keep.add(im["id"]) new_images.append(im) for ann in annotations: if ann["image_id"] in img_ids_to_keep: new_annotations.append(ann) json1["images"] = new_images json1["annotations"] = new_annotations output_filename = "clean_result.json" with open(output_filename, "w") as output_file: json.dump(json1, output_file) print(f"Saved {output_filename}") print_coco(json1) ================================================ FILE: utils/tf1_trt_inference.py ================================================ import tensorflow as tf from tensorflow.python.compiler.tensorrt import trt_convert import tensorflow.contrib.tensorrt as trt """ Code based on: - https://docs.nvidia.com/deeplearning/frameworks/tf-trt-user-guide/index.html#using-savedmodel - https://github.com/NVIDIA-AI-IOT/tf_trt_models/blob/master/examples/detection/detection.ipynb """ # converter = trt_convert.TrtGraphConverter( # input_saved_model_dir=input_saved_model_dir, # precision_mode=”FP16”, # maximum_cached_engines=100) # converter.convert() # converter.save(output_saved_model_dir) # with tf.Session() as sess: # # First load the SavedModel into the session # tf.saved_model.loader.load( # sess, [tf.saved_model.tag_constants.SERVING], # output_saved_model_dir) # output = sess.run([output_tensor], feed_dict={input_tensor: input_data}) input_name = ["image_tensor"] output_names = [ "detection_boxes", "detection_classes", "detection_scores", "num_detections", ] tf_config.gpu_options.allow_growth = True tf_sess = tf.Session(config=tf_config) tf.import_graph_def(trt_graph, name="") tf_input = tf_sess.graph.get_tensor_by_name(input_names[0] + ":0") tf_scores = tf_sess.graph.get_tensor_by_name("detection_scores:0") tf_boxes = tf_sess.graph.get_tensor_by_name("detection_boxes:0") tf_classes = tf_sess.graph.get_tensor_by_name("detection_classes:0") tf_num_detections = tf_sess.graph.get_tensor_by_name("num_detections:0") trt_graph = trt.create_inference_graph( input_graph_def=frozen_graph, outputs=output_names, max_batch_size=1, max_workspace_size_bytes=1 << 25, precision_mode="FP16", minimum_segment_size=50, ) ================================================ FILE: utils/tf2_trt_convert.py ================================================ import sys import tensorflow as tf from tensorflow.python.compiler.tensorrt import trt_convert as trt input_saved_model_dir = sys.argv[1] output_saved_model_dir = sys.argv[2] """ This code is based on: - https://docs.nvidia.com/deeplearning/frameworks/tf-trt-user-guide/index.html#worflow-with-savedmodel - https://sayak.dev/tf.keras/tensorrt/tensorflow/2020/07/01/accelerated-inference-trt.html Currently fails after a while (jetson, tensorflow==2.3.1) with: 2020-11-27 16:32:23.116411: W tensorflow/core/framework/op_kernel.cc:1767] OP_REQUIRES failed at trt_engine_resource_ops.cc:196 : Not found: Container TF-TRT does not exist. (Could not find resource: TF-TRT/TRTEngineOp_0_0) """ conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS # conversion_params = conversion_params._replace(max_workspace_size_bytes=(1 << 32)) conversion_params = conversion_params._replace(precision_mode="FP16") conversion_params = conversion_params._replace(maximum_cached_engines=100) converter = trt.TrtGraphConverterV2( input_saved_model_dir=input_saved_model_dir, conversion_params=conversion_params ) converter.convert() converter.save(output_saved_model_dir) # saved_model_loaded = tf.saved_model.load( # output_saved_model_dir, tags=[tag_constants.SERVING]) # graph_func = saved_model_loaded.signatures[ # signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] # frozen_func = convert_to_constants.convert_variables_to_constants_v2( # graph_func) # output = frozen_func(input_data)[0].numpy() ================================================ FILE: utils/tf_trt_convert.py ================================================ # %% import sys import os import time import urllib import matplotlib import numpy as np import tensorflow as tf import tensorflow.contrib.tensorrt as trt matplotlib.use("Agg") from PIL import Image # Install this from: https://github.com/NVIDIA-AI-IOT/tf_trt_models # Tested on Nano with python 3.6.9 and TF-1.15.4 from tf_trt_models.detection import download_detection_model, build_detection_graph import matplotlib.pyplot as plt import matplotlib.patches as patches # %% # inference_graph_path = "../inference_graph_1024x608" # inference_graph_path = "../inference_graph_300x300" inference_graph_path = "../inference_graph_400x708" config_path = f"{inference_graph_path}/pipeline.config" checkpoint_path = f"{inference_graph_path}/model.ckpt" batch_size = 4 score_threshold = 0.5 # TODO: Try this, with 0.5 in 300x300 and check results w/slack frozen_graph, input_names, output_names = build_detection_graph( config=config_path, checkpoint=checkpoint_path, batch_size=batch_size, score_threshold=score_threshold, ) # %% trt_graph = trt.create_inference_graph( input_graph_def=frozen_graph, outputs=output_names, max_batch_size=batch_size, max_workspace_size_bytes=1 << 25, precision_mode="FP16", minimum_segment_size=50, ) # %% converted_trt_graph_file = f"converted_trt_708_400_bs{batch_size}.pb" # %% with open(converted_trt_graph_file, "wb") as f: f.write(trt_graph.SerializeToString()) # %% trt_graph = tf.GraphDef() with open(converted_trt_graph_file, "rb") as f: trt_graph.ParseFromString(f.read()) # %% input_names = ["image_tensor"] output_names = [ "detection_boxes", "detection_classes", "detection_scores", "num_detections", ] # %% tf_config = tf.ConfigProto() tf_config.gpu_options.allow_growth = True tf_sess = tf.Session(config=tf_config) tf.import_graph_def(trt_graph, name="") tf_input = tf_sess.graph.get_tensor_by_name(input_names[0] + ":0") tf_scores = tf_sess.graph.get_tensor_by_name("detection_scores:0") tf_boxes = tf_sess.graph.get_tensor_by_name("detection_boxes:0") tf_classes = tf_sess.graph.get_tensor_by_name("detection_classes:0") tf_num_detections = tf_sess.graph.get_tensor_by_name("num_detections:0") # %% paths = ["../yolo/data/obj_train_data/images/hdstock_2_90.jpg"] paths += ["../yolo/data/obj_train_data/images/1_30s_0.jpg"] image = Image.open(paths[0]) plt.imshow(image) # image_resized = np.array(image.resize((1024, 608))) # image_resized = np.array(image.resize((300, 300))) image_resized = np.array(image.resize((708, 400))) image = np.array(image) # %% scores, boxes, classes, num_detections = tf_sess.run( [tf_scores, tf_boxes, tf_classes, tf_num_detections], feed_dict={tf_input: np.stack([image_resized] * batch_size)}, ) boxes = boxes[0] # index by 0 to remove batch dimension scores = scores[0] classes = classes[0] num_detections = num_detections[0] # %% fig = plt.figure() ax = fig.add_subplot(1, 1, 1) ax.imshow(image) # plot boxes exceeding score threshold for i in range(int(num_detections)): # scale box to image coordinates box = boxes[i] * np.array( [image.shape[0], image.shape[1], image.shape[0], image.shape[1]] ) # display rectangle patch = patches.Rectangle( (box[1], box[0]), box[3] - box[1], box[2] - box[0], color="g", alpha=0.3 ) ax.add_patch(patch) # display class index and score plt.text( x=box[1] + 10, y=box[2] - 10, s="%d (%0.2f) " % (classes[i], scores[i]), color="w", ) plt.savefig("detections.png") # %% num_samples = 50 input_batch = np.stack([image_resized] * batch_size) t0 = time.time() for i in range(num_samples): scores, boxes, classes, num_detections = tf_sess.run( [tf_scores, tf_boxes, tf_classes, tf_num_detections], feed_dict={tf_input: input_batch}, ) t1 = time.time() print("Average runtime: %f seconds" % (float(t1 - t0) / num_samples)) # %% tf_sess.close() ================================================ FILE: yolo/config.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ import yaml class Config: def __init__(self, config_file_path): # Load config file with open(config_file_path, "r") as stream: self._config = yaml.load(stream, Loader=yaml.FullLoader) # Define colors to be used internally through the app, and also externally if wanted self.colors = { "green": (0, 128, 0), "white": (255, 255, 255), "olive": (0, 128, 128), "black": (0, 0, 0), "navy": (128, 0, 0), "red": (0, 0, 255), "pink": (128, 128, 255), "maroon": (0, 0, 128), "grey": (128, 128, 128), "purple": (128, 0, 128), "yellow": (0, 255, 255), "lime": (0, 255, 0), "fuchsia": (255, 0, 255), "aqua": (255, 255, 0), "blue": (255, 0, 0), "teal": (128, 128, 0), "silver": (192, 192, 192), } def __getitem__(self, name): return self._config[name] ================================================ FILE: yolo/config_images.yml ================================================ yolo_generic: detection_threshold: 0.1 nms_threshold: 0.1 # Lower values filter more distance_threshold: 1 yolo_trt_tiny: names_file: data/obj.names # engine_file: ../yolo/yolov4-facemask-tiny-fp16.trt # engine_file: ../../tensorrt_batch8_fp16.trt # engine_file: facemask_y4tiny_1024_608_fp16.trt # engine_file: maskcam_y4t_1120_640_fp16.trt engine_file: maskcam_y4t_1024_608_fp16.trt use_cuda: true input_width: 1024 input_height: 608 batch_size: 1 min_detection_size: 8 # Also see: face_min_size in face_mask_detector debug: output_detector_resolution: true draw_detections: true profiler: true ================================================ FILE: yolo/data/obj.data ================================================ classes = 4 train = data/train.txt valid = data/valid.txt names = data/obj.names backup = backup/ ================================================ FILE: yolo/data/obj.names ================================================ mask no_mask not_visible misplaced ================================================ FILE: yolo/facemask-yolov4-tiny.cfg ================================================ [net] # Testing #batch=1 #subdivisions=1 # Training batch=64 subdivisions=16 # Sizes: multiples of 32. Ratio (hd/4k videos): 1.78 # width=1120 # height=640 width=1024 height=608 channels=3 momentum=0.9 decay=0.0005 angle=10 saturation = 1.5 exposure = 1.5 hue=.2 learning_rate=0.00261 burn_in=1000 policy=steps # max_batches = 8000 # steps=6400,7200 max_batches = 12000 steps=9600,10800 scales=.1,.1 [convolutional] batch_normalize=1 filters=32 size=3 stride=2 pad=1 activation=leaky [convolutional] batch_normalize=1 filters=64 size=3 stride=2 pad=1 activation=leaky [convolutional] batch_normalize=1 filters=64 size=3 stride=1 pad=1 activation=leaky [route] layers=-1 groups=2 group_id=1 [convolutional] batch_normalize=1 filters=32 size=3 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 filters=32 size=3 stride=1 pad=1 activation=leaky [route] layers = -1,-2 [convolutional] batch_normalize=1 filters=64 size=1 stride=1 pad=1 activation=leaky [route] layers = -6,-1 [maxpool] size=2 stride=2 [convolutional] batch_normalize=1 filters=128 size=3 stride=1 pad=1 activation=leaky [route] layers=-1 groups=2 group_id=1 [convolutional] batch_normalize=1 filters=64 size=3 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 filters=64 size=3 stride=1 pad=1 activation=leaky [route] layers = -1,-2 [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=leaky [route] layers = -6,-1 [maxpool] size=2 stride=2 [convolutional] batch_normalize=1 filters=256 size=3 stride=1 pad=1 activation=leaky [route] layers=-1 groups=2 group_id=1 [convolutional] batch_normalize=1 filters=128 size=3 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 filters=128 size=3 stride=1 pad=1 activation=leaky [route] layers = -1,-2 [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=leaky [route] layers = -6,-1 [maxpool] size=2 stride=2 [convolutional] batch_normalize=1 filters=512 size=3 stride=1 pad=1 activation=leaky ################################## [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 filters=512 size=3 stride=1 pad=1 activation=leaky [convolutional] size=1 stride=1 pad=1 filters=27 activation=linear [yolo] mask = 3,4,5 anchors = 10,14, 23,27, 37,58, 81,82, 135,169, 344,319 classes=4 num=6 jitter=.3 scale_x_y = 1.05 cls_normalizer=1.0 iou_normalizer=0.07 iou_loss=ciou ignore_thresh = .7 truth_thresh = 1 random=0 resize=1.5 nms_kind=greedynms beta_nms=0.6 [route] layers = -4 [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=leaky [upsample] stride=2 [route] layers = -1, 23 [convolutional] batch_normalize=1 filters=256 size=3 stride=1 pad=1 activation=leaky [convolutional] size=1 stride=1 pad=1 filters=27 activation=linear [yolo] mask = 0,1,2 anchors = 10,14, 23,27, 37,58, 81,82, 135,169, 344,319 classes=4 num=6 jitter=.3 scale_x_y = 1.05 cls_normalizer=1.0 iou_normalizer=0.07 iou_loss=ciou ignore_thresh = .7 truth_thresh = 1 random=0 resize=1.5 nms_kind=greedynms beta_nms=0.6 ================================================ FILE: yolo/facemask-yolov4.cfg ================================================ [net] # Testing #batch=1 #subdivisions=1 # Training batch=64 subdivisions=32 width=608 height=608 channels=3 momentum=0.949 decay=0.0005 angle=0 saturation = 1.5 exposure = 1.5 hue=.1 learning_rate=0.001 burn_in=1000 max_batches = 8000 policy=steps steps=400000,450000 scales=.1,.1 #cutmix=1 mosaic=1 #:104x104 54:52x52 85:26x26 104:13x13 for 416 [convolutional] batch_normalize=1 filters=32 size=3 stride=1 pad=1 activation=mish # Downsample [convolutional] batch_normalize=1 filters=64 size=3 stride=2 pad=1 activation=mish [convolutional] batch_normalize=1 filters=64 size=1 stride=1 pad=1 activation=mish [route] layers = -2 [convolutional] batch_normalize=1 filters=64 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=32 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=64 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=64 size=1 stride=1 pad=1 activation=mish [route] layers = -1,-7 [convolutional] batch_normalize=1 filters=64 size=1 stride=1 pad=1 activation=mish # Downsample [convolutional] batch_normalize=1 filters=128 size=3 stride=2 pad=1 activation=mish [convolutional] batch_normalize=1 filters=64 size=1 stride=1 pad=1 activation=mish [route] layers = -2 [convolutional] batch_normalize=1 filters=64 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=64 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=64 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=64 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=64 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=64 size=1 stride=1 pad=1 activation=mish [route] layers = -1,-10 [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=mish # Downsample [convolutional] batch_normalize=1 filters=256 size=3 stride=2 pad=1 activation=mish [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=mish [route] layers = -2 [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=128 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=128 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=128 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=128 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=128 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=128 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=128 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=128 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=mish [route] layers = -1,-28 [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=mish # Downsample [convolutional] batch_normalize=1 filters=512 size=3 stride=2 pad=1 activation=mish [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=mish [route] layers = -2 [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=256 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=256 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=256 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=256 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=256 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=256 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=256 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=256 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=mish [route] layers = -1,-28 [convolutional] batch_normalize=1 filters=512 size=1 stride=1 pad=1 activation=mish # Downsample [convolutional] batch_normalize=1 filters=1024 size=3 stride=2 pad=1 activation=mish [convolutional] batch_normalize=1 filters=512 size=1 stride=1 pad=1 activation=mish [route] layers = -2 [convolutional] batch_normalize=1 filters=512 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=512 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=512 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=512 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=512 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=512 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=512 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=512 size=1 stride=1 pad=1 activation=mish [convolutional] batch_normalize=1 filters=512 size=3 stride=1 pad=1 activation=mish [shortcut] from=-3 activation=linear [convolutional] batch_normalize=1 filters=512 size=1 stride=1 pad=1 activation=mish [route] layers = -1,-16 [convolutional] batch_normalize=1 filters=1024 size=1 stride=1 pad=1 activation=mish stopbackward=800 ########################## [convolutional] batch_normalize=1 filters=512 size=1 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 size=3 stride=1 pad=1 filters=1024 activation=leaky [convolutional] batch_normalize=1 filters=512 size=1 stride=1 pad=1 activation=leaky ### SPP ### [maxpool] stride=1 size=5 [route] layers=-2 [maxpool] stride=1 size=9 [route] layers=-4 [maxpool] stride=1 size=13 [route] layers=-1,-3,-5,-6 ### End SPP ### [convolutional] batch_normalize=1 filters=512 size=1 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 size=3 stride=1 pad=1 filters=1024 activation=leaky [convolutional] batch_normalize=1 filters=512 size=1 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=leaky [upsample] stride=2 [route] layers = 85 [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=leaky [route] layers = -1, -3 [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 size=3 stride=1 pad=1 filters=512 activation=leaky [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 size=3 stride=1 pad=1 filters=512 activation=leaky [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=leaky [upsample] stride=2 [route] layers = 54 [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=leaky [route] layers = -1, -3 [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 size=3 stride=1 pad=1 filters=256 activation=leaky [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 size=3 stride=1 pad=1 filters=256 activation=leaky [convolutional] batch_normalize=1 filters=128 size=1 stride=1 pad=1 activation=leaky ########################## [convolutional] batch_normalize=1 size=3 stride=1 pad=1 filters=256 activation=leaky [convolutional] size=1 stride=1 pad=1 filters=27 activation=linear [yolo] mask = 0,1,2 anchors = 12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401 classes=4 num=9 jitter=.3 ignore_thresh = .7 truth_thresh = 1 scale_x_y = 1.2 iou_thresh=0.213 cls_normalizer=1.0 iou_normalizer=0.07 iou_loss=ciou nms_kind=greedynms beta_nms=0.6 max_delta=5 [route] layers = -4 [convolutional] batch_normalize=1 size=3 stride=2 pad=1 filters=256 activation=leaky [route] layers = -1, -16 [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 size=3 stride=1 pad=1 filters=512 activation=leaky [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 size=3 stride=1 pad=1 filters=512 activation=leaky [convolutional] batch_normalize=1 filters=256 size=1 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 size=3 stride=1 pad=1 filters=512 activation=leaky [convolutional] size=1 stride=1 pad=1 filters=27 activation=linear [yolo] mask = 3,4,5 anchors = 12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401 classes=4 num=9 jitter=.3 ignore_thresh = .7 truth_thresh = 1 scale_x_y = 1.1 iou_thresh=0.213 cls_normalizer=1.0 iou_normalizer=0.07 iou_loss=ciou nms_kind=greedynms beta_nms=0.6 max_delta=5 [route] layers = -4 [convolutional] batch_normalize=1 size=3 stride=2 pad=1 filters=512 activation=leaky [route] layers = -1, -37 [convolutional] batch_normalize=1 filters=512 size=1 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 size=3 stride=1 pad=1 filters=1024 activation=leaky [convolutional] batch_normalize=1 filters=512 size=1 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 size=3 stride=1 pad=1 filters=1024 activation=leaky [convolutional] batch_normalize=1 filters=512 size=1 stride=1 pad=1 activation=leaky [convolutional] batch_normalize=1 size=3 stride=1 pad=1 filters=1024 activation=leaky [convolutional] size=1 stride=1 pad=1 filters=27 activation=linear [yolo] mask = 6,7,8 anchors = 12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401 classes=4 num=9 jitter=.3 ignore_thresh = .7 truth_thresh = 1 random=1 scale_x_y = 1.05 iou_thresh=0.213 cls_normalizer=1.0 iou_normalizer=0.07 iou_loss=ciou nms_kind=greedynms beta_nms=0.6 max_delta=5 ================================================ FILE: yolo/integrations/yolo/detector_trt.py ================================================ import sys import os import time import argparse import numpy as np import cv2 import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit from norfair.tracker import Detection from integrations.yolo.utils_pytorch import load_class_names, post_processing """ Code based on these implementations: - (main) https://github.com/NVIDIA/object-detection-tensorrt-example/blob/master/SSD_Model/utils/inference.py - (originally found here) https://github.com/Tianxiaomo/pytorch-YOLOv4/blob/master/demo_trt.py """ # Simple helper data class that's a little nicer to use than a 2-tuple. class HostDeviceMem(object): def __init__(self, host_mem, device_mem): self.host = host_mem self.device = device_mem def __str__(self): return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device) def __repr__(self): return self.__str__() class DetectorYoloTRT: """ Adaptor for the original Yolo implementation (AlexeyAB/darknet) """ def __init__(self, config): self.batch_size = config["batch_size"] self.input_h = config["input_height"] self.input_w = config["input_width"] self.detection_threshold = config["detection_threshold"] self.nms_threshold = config["nms_threshold"] self.engine_path = config["engine_file"] self.class_names = load_class_names(config["names_file"]) if "min_detection_size" in config: self.min_size = config["min_detection_size"] else: self.min_size = 0 self.logger = trt.Logger() self.runtime = trt.Runtime(self.logger) print("Reading engine from file {}".format(self.engine_path)) with open(self.engine_path, "rb") as f: self.engine = self.runtime.deserialize_cuda_engine(f.read()) self.context = self.engine.create_execution_context() self.buffers = self._allocate_buffers(self.engine) self.context.set_binding_shape( 0, (self.batch_size, 3, self.input_h, self.input_w) ) self.img_batch = np.zeros((self.batch_size, 3 * self.input_h * self.input_w)) self.timer_preprocess = 0.0 self.timer_inference = 0.0 self.timer_execute = 0.0 self.timer_postprocess = 0.0 self.n_frames = 0 self.n_inferences = 0 def print_profiler(self): print( f"Batch size: {self.batch_size}" f" | Frames processed: {self.n_frames}" f" | # inferences executed: {self.n_inferences}" ) print( f"Avg preprocess time/frame:\t{self.timer_preprocess/self.n_frames:.4f}s" f"\t| FPS: {self.n_frames / self.timer_preprocess:.1f}" ) print( f"Avg inference time/frame:\t{self.timer_inference/self.n_frames:.4f}s" f"\t| FPS: {self.n_frames / self.timer_inference:.1f}" ) print( f"Avg postprocess time/frame:\t{self.timer_postprocess/self.n_frames:.4f}s" f"\t| FPS: {self.n_frames / self.timer_postprocess:.1f}" ) print( f"Avg execute time/frame:\t{self.timer_execute/self.n_frames:.4f}s" f"\t| FPS: {self.n_frames / self.timer_execute:.1f}" ) print( f"Avg execute time/inference:\t{self.timer_execute/self.n_inferences:.4f}s" ) def detect(self, frames, rescale_detections=True): inputs, outputs, bindings, stream = self.buffers frames_resized = [] self.n_frames += len(frames) tick = time.time() for idx, frame in enumerate(frames): orig_height, orig_width = frame.shape[:2] # Input frame_resized = cv2.resize( frame, (self.input_w, self.input_h), interpolation=cv2.INTER_LINEAR ) frames_resized.append(frame_resized) img_in = cv2.cvtColor(frame_resized, cv2.COLOR_BGR2RGB) img_in = np.transpose(img_in, (2, 0, 1)).astype(np.float32) # img_in = np.expand_dims(img_in, axis=0) img_in /= 255.0 # img_in = np.ascontiguousarray(img_in) self.img_batch[idx] = img_in.ravel() np.copyto(inputs[0].host, self.img_batch.ravel()) self.timer_preprocess += time.time() - tick tick = time.time() trt_outputs = self._do_inference( self.context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream, ) self.timer_inference += time.time() - tick trt_outputs[0] = trt_outputs[0].reshape(self.batch_size, -1, 1, 4) trt_outputs[1] = trt_outputs[1].reshape( self.batch_size, -1, len(self.class_names) ) # detection threshold + NMS filtering tick = time.time() detections = post_processing( img_in, self.detection_threshold, self.nms_threshold, trt_outputs ) self.timer_postprocess += time.time() - tick dets_batches = [] for batch_idx in range(len(detections)): width = orig_width if rescale_detections else self.input_w height = orig_height if rescale_detections else self.input_h dets = [] for k, d in enumerate(detections[batch_idx]): d[0] *= width d[1] *= height d[2] *= width d[3] *= height if self.min_size: detection_width = d[2] - d[0] detection_height = d[3] - d[1] if min(detection_height, detection_width) < self.min_size: break p = d[4] label = self.class_names[d[6]] dets.append( Detection( np.array((d[0:2], d[2:4])), data={"label": label, "p": p}, ) ) dets_batches.append(dets) return dets_batches, frames_resized # This function is generalized for multiple inputs/outputs. # inputs and outputs are expected to be lists of HostDeviceMem objects. def _do_inference(self, context, bindings, inputs, outputs, stream): self.n_inferences += 1 tick = time.time() # Transfer input data to the GPU. [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs] # Run inference. context.execute_async_v2(bindings=bindings, stream_handle=stream.handle) # Transfer predictions back from the GPU. [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs] # Synchronize the stream stream.synchronize() # In this case, we measure the memcpy + execute ops together (small difference) self.timer_execute += time.time() - tick # Return only the host outputs. return [out.host for out in outputs] def _do_inference_sync(self, context, bindings, inputs, outputs, stream): self.n_inferences += 1 # Transfer input data to the GPU. [cuda.memcpy_htod(inp.device, inp.host) for inp in inputs] # Run inference. tick = time.time() context.execute_v2(bindings=bindings) self.timer_execute += time.time() - tick # Transfer predictions back from the GPU. [cuda.memcpy_dtoh(out.host, out.device) for out in outputs] # Return only the host outputs. return [out.host for out in outputs] # Allocates all buffers required for an engine, i.e. host/device inputs/outputs. def _allocate_buffers(self, engine): inputs = [] outputs = [] bindings = [] stream = cuda.Stream() for binding in engine: # engine.max_batch_size is 1 for static batch size = ( trt.volume(engine.get_binding_shape(binding)) * self.engine.max_batch_size ) dims = engine.get_binding_shape(binding) # in case batch dimension is -1 (dynamic) if dims[0] < 0: size *= -1 dtype = trt.nptype(engine.get_binding_dtype(binding)) # Allocate host and device buffers host_mem = cuda.pagelocked_empty(size, dtype) device_mem = cuda.mem_alloc(host_mem.nbytes) # Append the device buffer to device bindings. bindings.append(int(device_mem)) # Append to the appropriate list. if engine.binding_is_input(binding): inputs.append(HostDeviceMem(host_mem, device_mem)) else: outputs.append(HostDeviceMem(host_mem, device_mem)) return inputs, outputs, bindings, stream # Allocates all buffers required for an engine, i.e. host/device inputs/outputs. def _allocate_buffers(self, engine): inputs = [] outputs = [] bindings = [] stream = cuda.Stream() for binding in engine: # engine.max_batch_size is 1 for static batch size = ( trt.volume(engine.get_binding_shape(binding)) * self.engine.max_batch_size ) dims = engine.get_binding_shape(binding) # in case batch dimension is -1 (dynamic) if dims[0] < 0: size *= -1 dtype = trt.nptype(engine.get_binding_dtype(binding)) # Allocate host and device buffers host_mem = cuda.pagelocked_empty(size, dtype) device_mem = cuda.mem_alloc(host_mem.nbytes) # Append the device buffer to device bindings. bindings.append(int(device_mem)) # Append to the appropriate list. if engine.binding_is_input(binding): inputs.append(HostDeviceMem(host_mem, device_mem)) else: outputs.append(HostDeviceMem(host_mem, device_mem)) return inputs, outputs, bindings, stream ================================================ FILE: yolo/integrations/yolo/utils_pytorch.py ================================================ import numpy as np """ Functions copied from pytorch-YOLOv4/tool/utils.py with all calls to print() removed """ def nms_cpu(boxes, confs, nms_thresh=0.5, min_mode=False): # print(boxes.shape) x1 = boxes[:, 0] y1 = boxes[:, 1] x2 = boxes[:, 2] y2 = boxes[:, 3] areas = (x2 - x1) * (y2 - y1) order = confs.argsort()[::-1] keep = [] while order.size > 0: idx_self = order[0] idx_other = order[1:] keep.append(idx_self) xx1 = np.maximum(x1[idx_self], x1[idx_other]) yy1 = np.maximum(y1[idx_self], y1[idx_other]) xx2 = np.minimum(x2[idx_self], x2[idx_other]) yy2 = np.minimum(y2[idx_self], y2[idx_other]) w = np.maximum(0.0, xx2 - xx1) h = np.maximum(0.0, yy2 - yy1) inter = w * h if min_mode: over = inter / np.minimum(areas[order[0]], areas[order[1:]]) else: over = inter / (areas[order[0]] + areas[order[1:]] - inter) inds = np.where(over <= nms_thresh)[0] order = order[inds + 1] return np.array(keep) def load_class_names(namesfile): class_names = [] with open(namesfile, "r") as fp: lines = fp.readlines() for line in lines: line = line.rstrip() class_names.append(line) return class_names def post_processing(img, conf_thresh, nms_thresh, output): # anchors = [12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401] # num_anchors = 9 # anchor_masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]] # strides = [8, 16, 32] # anchor_step = len(anchors) // num_anchors # [batch, num, 1, 4] box_array = output[0] # [batch, num, num_classes] confs = output[1] if type(box_array).__name__ != "ndarray": box_array = box_array.cpu().detach().numpy() confs = confs.cpu().detach().numpy() num_classes = confs.shape[2] # [batch, num, 4] box_array = box_array[:, :, 0] # [batch, num, num_classes] --> [batch, num] max_conf = np.max(confs, axis=2) max_id = np.argmax(confs, axis=2) bboxes_batch = [] for i in range(box_array.shape[0]): argwhere = max_conf[i] > conf_thresh l_box_array = box_array[i, argwhere, :] l_max_conf = max_conf[i, argwhere] l_max_id = max_id[i, argwhere] bboxes = [] # nms for each class for j in range(num_classes): cls_argwhere = l_max_id == j ll_box_array = l_box_array[cls_argwhere, :] ll_max_conf = l_max_conf[cls_argwhere] ll_max_id = l_max_id[cls_argwhere] keep = nms_cpu(ll_box_array, ll_max_conf, nms_thresh) if keep.size > 0: ll_box_array = ll_box_array[keep, :] ll_max_conf = ll_max_conf[keep] ll_max_id = ll_max_id[keep] for k in range(ll_box_array.shape[0]): bboxes.append( [ ll_box_array[k, 0], ll_box_array[k, 1], ll_box_array[k, 2], ll_box_array[k, 3], ll_max_conf[k], ll_max_conf[k], ll_max_id[k], ] ) bboxes_batch.append(bboxes) return bboxes_batch ================================================ FILE: yolo/integrations/yolo/yolo_adaptor.py ================================================ import cv2 import numpy as np from norfair.drawing import Color class YoloAdaptor: def __init__(self, config): self.detection_threshold = config["detection_threshold"] self.distance_threshold = config["distance_threshold"] def classify_people(self, tracked_people): p_masks = [] for d in tracked_people: meta = d.last_detection.data if meta["label"] == "mask": p_mask = float(meta["p"]) elif meta["label"] == "no_mask" or meta["label"] == "misplaced": p_mask = 1 - float(meta["p"]) elif meta["label"] == "not_visible": p_mask = 0.5 else: raise # Unknown label p_masks.append(p_mask) return p_masks def keypoints_distance(self, detected_pose, tracked_pose): detected_points = detected_pose.points estimated_pose = tracked_pose.estimate min_box_size = min( max( detected_points[1][0] - detected_points[0][0], # x2 - x1 detected_points[1][1] - detected_points[0][1], # y2 - y1 1, ), max( estimated_pose[1][0] - estimated_pose[0][0], # x2 - x1 estimated_pose[1][1] - estimated_pose[0][1], # y2 - y1 1, ), ) mean_distance_normalized = ( np.mean(np.linalg.norm(detected_points - estimated_pose, axis=1)) / min_box_size ) return mean_distance_normalized def person_has_face(self, person): return person.last_detection.data["label"] != "not_visible" def get_person_head(self, person): if person.live_points.sum() < 2: return None p1, p2 = person.estimate.astype(int) return (tuple(p1), tuple(p2)) def draw_raw_detections(self, frame, detections): for d in detections: p1, p2 = d.points.astype(int) bbox = (tuple(p1), tuple(p2)) label = d.data["label"] p = float(d.data["p"]) color = ( Color.green if label == "mask" else ( Color.red if label == "no_mask" else (Color.yellow if label == "misplaced" else Color.white) ) ) cv2.rectangle(frame, bbox[0], bbox[1], color, 1) cv2.putText( frame, f"{label}: {p:.2f}", (bbox[0][0], bbox[0][1] - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA, ) # # Draw debugging info # cv2.putText( # frame, # f"width: {bbox[1][0] - bbox[0][0]}", # (bbox[1][0], bbox[1][1] + 10), # cv2.FONT_HERSHEY_SIMPLEX, # 0.5, # color, # 1, # cv2.LINE_AA, # ) ================================================ FILE: yolo/run_yolo_images.py ================================================ ################################################################################ # Copyright (c) 2020-2021, Berkeley Design Technology, Inc. All rights reserved. # # Permission 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: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE 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. ################################################################################ # %% import os import cv2 import sys import yaml import time import glob from integrations.yolo.yolo_adaptor import YoloAdaptor # Requires python tensorrt, usually compiled for python 3.6 at system level from integrations.yolo.detector_trt import DetectorYoloTRT # %% images_folder = sys.argv[1] output_folder = sys.argv[2] print(f"Scanning input directory: {images_folder}") images = [] for filetype in ["png", "jpg", "jpeg"]: # No uppercase for now images += glob.glob(f"{images_folder}/*.{filetype}") print(f"Found {len(images)} images") if input(f"Confirm output to [{output_folder}] [y/n]").strip() != "y": print("Not confirmed. Exiting") exit(0) os.system(f"mkdir -p {output_folder}") # %% with open("config_images.yml", "r") as stream: # Not using Loader=yaml.FullLoader since it doesn't work on jetson PyYAML version config = yaml.load(stream) yolo_config = {**config["yolo_trt_tiny"], **config["yolo_generic"]} detector = DetectorYoloTRT(yolo_config) # Converter functions from Yolo -> Tracker + FaceMaskDetector pose_adaptor = YoloAdaptor(config["yolo_generic"]) detector_output = config["debug"]["output_detector_resolution"] for k, image_filename in enumerate(images): frame = cv2.imread(image_filename) if ( detector_output ): # Only for debugging purposes: use resized frame in video output detections, frames_resized = detector.detect([frame], rescale_detections=False) frame = frames_resized[0] else: detections, _ = detector.detect([frame], rescale_detections=True) detections = detections[0] # Remove batch dimension # Drawing functions if config["debug"]["draw_detections"]: # Raw yolo detections pose_adaptor.draw_raw_detections(frame, detections) im_basename = image_filename.split("/")[-1] image_outfile = f"{output_folder}/{im_basename}" cv2.imwrite(image_outfile, frame) print(f"Writing [{k}/{len(images)}]: {image_outfile}") if config["debug"]["profiler"]: detector.print_profiler() ================================================ FILE: yolo/train_cu90.sh ================================================ mkdir -p backup LD_LIBRARY_PATH=/usr/local/cuda-9.0/lib64/ ./darknet detector train data/obj.data facemask-yolov4-tiny.cfg yolov4-tiny.conv.29 -dont_show -mjpeg_port 8090 -map