Tiny YOLOv4 TensorFlow Lite model on Jetson Xavier

This sample is an example of running an AI container on the Jetson platform. This container utilizes the GPU on the Jetson (with NVIDIA drivers, CUDA and cuDNN installed) using an NVIDIA L4T (linux for Tegra) base image with TensorFlow 2 installed. The Jetson must have been flashed with Jetpack 4.4.

The model used in this project is converted from the original Darknet tiny YOLO v4 weights, trained on the COCO dataset and able to recognize the 80 COCO classes. To train on your own dataset see https://github.com/michhar/yolov4-darknet-notes.

Xavier Setup

• Flashed with JetPack 4.4 (L4T R32.4.3) with all ML and CV tools (including nvidia-docker)
• Samsung NVMe 512 GB to store docker images
• 16 GB swap file on NVMe mount
• Azure CLI for pushing image to Azure Container Registry
• [Optional] Docker may be configured to run with non-root user as in Manage Docker as a non-root user allowing the omission of using sudo with docker

Build image

The following instructions will enable you to build a docker container with a YOLOv4 (tiny) TensorFlow Lite model using nginx, gunicorn, flask, and runit. The app code is based on the tensorflow-yolov4-tflite project. This project uses TensorFlow v2.

Note: References to third-party software in this repo are for informational and convenience purposes only. Microsoft does not endorse nor provide rights for the third-party software. For more information on third-party software please see the links provided above.

Prerequisites for building image

1. Ensure NVIDIA Docker on your Jetson
2. Install curl

Building the docker container

1. Create a new directory on your machine and copy all the files (including the sub-folders) from this GitHub folder to that directory.
2. Build the container image (will take several minutes) by running the following docker command from a terminal window in that directory.

sudo nvidia-docker build . -t tiny-yolov4-tflite:arm64v8-cuda-cudnn -f arm64v8-gpu-cudnn.dockerfile

Running and testing

The REST endpoint accepts an image with the size of 416 pixels by 416 pixels. This is requirement by the tiny YOLOv4 model. Since the LVA edge module is capable of sending specified size image in specified format, we are not preprocessing the incoming images to resize them. This is mainly because of the performance improvement.

Run the container using the following docker command.

sudo docker run --runtime=nvidia --name my_yolo_container -p 80:80 -d  -i tiny-yolov4-tflite:arm64v8-cuda-cudnn

Test the container using the following commands.

/score

To get a list of detected objects, use the following command.

curl -X POST http://127.0.0.1/score -H "Content-Type: image/jpeg" --data-binary @<full_path_to_image_file_in_jpeg>

If successful, you will see JSON printed on your screen that looks something like this

{
  "inferences": [
    {
      "type": "entity",
      "entity": {
        "tag": {
          "value": "zebra",
          "confidence": "0.8333446"
        },
        "box": {
          "l": "0.6046585",
          "t": "0.4014857",
          "w": "0.21853799",
          "h": "0.49041268"
        }
      }
    },
    {
      "type": "entity",
      "entity": {
        "tag": {
          "value": "giraffe",
          "confidence": "0.769461"
        },
        "box": {
          "l": "0.33088243",
          "t": "0.0015953871",
          "w": "0.5128964",
          "h": "0.83996487"
        }
      }
    }
  ]
}

Terminate the container using the following docker commands.

docker stop my_yolo_container
docker rm my_yolo_container

Upload docker image to Azure container registry

Follow instruction in Push and Pull Docker images - Azure Container Registry to save your image for later use on another machine.

IMPORTANT: Docker may need to be configured to run with non-root user as in Manage Docker as a non-root user.

Deploy as an Azure IoT Edge module

Follow instruction in Deploy module from Azure portal to deploy the container image as an IoT Edge module (use the IoT Edge module option).

Troubleshooting

Troubleshooting a running container

To troubleshoot a running container you may enter it with ssh by using the following command.

sudo docker exec -it my_yolo_container /bin/bash

Helpful links

• darknet implementation for YOLOv4
• TensorFlow YOLOv4 converters and implementations
• Live Video Analytics on Edge documentation
• A method to train a tiny YOLO v4 model on your own data

Credits

• TensorFlow Container for Jetson and JetPack
• Live Video Analytics Tiny YOLOv4 TensorFlow Lite model sample